diff --git a/stapler-scripts/display-switch/display_switch.py b/stapler-scripts/display-switch/display_switch.py index b914456..882ae17 100755 --- a/stapler-scripts/display-switch/display_switch.py +++ b/stapler-scripts/display-switch/display_switch.py @@ -5,9 +5,22 @@ import os import stat import datetime +import shlex +import getpass -STATE_FILE = "/tmp/monitor_state.sh" -LOG_FILE = "/tmp/display_switch.log" +def _get_secure_path(filename): + """Returns a secure, user-specific path for a file.""" + user = getpass.getuser() + # Try XDG_RUNTIME_DIR first (typically /run/user/UID) + base_dir = os.environ.get("XDG_RUNTIME_DIR") + if not base_dir or not os.path.isdir(base_dir) or not os.access(base_dir, os.W_OK): + base_dir = "/tmp" + + # Use a user-prefixed filename to avoid collisions and predictable paths in /tmp + return os.path.join(base_dir, f"display-switch-{user}-{filename}") + +STATE_FILE = _get_secure_path("monitor_state.sh") +LOG_FILE = _get_secure_path("display_switch.log") def log(message): timestamp = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S") @@ -166,17 +179,21 @@ def save_state(): command_parts.extend(["--rotate", m['rotation']]) - full_command = " ".join(command_parts) + full_command = shlex.join(command_parts) log(f"Generated restore command: {full_command}") try: - with open(STATE_FILE, "w") as f: + # Securely create the file with restrictive permissions (0o600) + # os.O_EXCL ensures we don't follow symlinks or overwrite an existing file + fd = os.open(STATE_FILE, os.O_CREAT | os.O_EXCL | os.O_WRONLY, 0o600) + with os.fdopen(fd, 'w') as f: f.write("#!/bin/bash\n") f.write(full_command + "\n") - os.chmod(STATE_FILE, os.stat(STATE_FILE).st_mode | stat.S_IEXEC) log(f"State saved successfully to {STATE_FILE}") - except IOError as e: + except FileExistsError: + log(f"State file {STATE_FILE} already exists. Skipping save.") + except (IOError, OSError) as e: log(f"Failed to write state file: {e}") sys.exit(1) @@ -186,9 +203,19 @@ def restore_state(): log(f"No state file found at {STATE_FILE}. Nothing to restore.") return + # Security check: Ensure the file is owned by the current user and has restrictive permissions + st = os.stat(STATE_FILE) + if st.st_uid != os.getuid(): + log(f"Security error: {STATE_FILE} is not owned by the current user.") + return + if st.st_mode & (stat.S_IRWXG | stat.S_IRWXO): + log(f"Security error: {STATE_FILE} has insecure permissions: {oct(st.st_mode)}") + return + try: log(f"Executing {STATE_FILE}...") - subprocess.run(STATE_FILE, check=True, shell=True) + # Execute securely: No shell=True, and pass the file path as an argument to bash + subprocess.run(["/bin/bash", STATE_FILE], check=True) log("State restored successfully.") os.remove(STATE_FILE) log(f"Deleted {STATE_FILE}") diff --git a/stapler-scripts/display-switch/test_display_switch_security.py b/stapler-scripts/display-switch/test_display_switch_security.py new file mode 100644 index 0000000..4828063 --- /dev/null +++ b/stapler-scripts/display-switch/test_display_switch_security.py @@ -0,0 +1,88 @@ +import os +import stat +import subprocess +import pytest +import getpass +from unittest.mock import patch, MagicMock +import display_switch + +@pytest.fixture +def mock_env(tmp_path): + # Setup a clean environment for testing + state_file = tmp_path / "monitor_state.sh" + log_file = tmp_path / "display_switch.log" + with patch('display_switch.STATE_FILE', str(state_file)), \ + patch('display_switch.LOG_FILE', str(log_file)): + yield state_file, log_file + +def test_secure_path_generation(): + user = getpass.getuser() + with patch.dict(os.environ, {"XDG_RUNTIME_DIR": "/run/user/1000"}): + with patch('os.path.isdir', return_value=True), patch('os.access', return_value=True): + path = display_switch._get_secure_path("test.sh") + assert path == f"/run/user/1000/display-switch-{user}-test.sh" + + with patch.dict(os.environ, {}, clear=True): + path = display_switch._get_secure_path("test.sh") + assert path == f"/tmp/display-switch-{user}-test.sh" + +@patch('display_switch.get_current_state') +@patch('subprocess.run') +def test_save_state_secure_creation(mock_run, mock_get_state, mock_env): + state_file, _ = mock_env + mock_get_state.return_value = [{ + "name": "HDMI-0", "primary": True, "active": True, + "pos": "0x0", "rotation": "normal", "active_mode": "1920x1080", "rate": "60.00" + }] + + display_switch.save_state() + + assert state_file.exists() + st = os.stat(state_file) + # Check permissions are 0o600 (rw-------) + assert stat.S_IMODE(st.st_mode) == 0o600 + + with open(state_file, 'r') as f: + content = f.read() + assert "#!/bin/bash" in content + assert "xrandr --output HDMI-0 --primary --mode 1920x1080 --rate 60.00 --pos 0x0 --rotate normal" in content + +@patch('subprocess.run') +def test_restore_state_secure_execution(mock_run, mock_env): + state_file, _ = mock_env + state_file.write_text("#!/bin/bash\nxrandr --output HDMI-0 --auto") + os.chmod(state_file, 0o600) + + display_switch.restore_state() + + # Ensure it was called with bash and NO shell=True + mock_run.assert_called_once_with(["/bin/bash", str(state_file)], check=True) + assert not state_file.exists() + +@patch('subprocess.run') +def test_restore_state_rejects_insecure_permissions(mock_run, mock_env): + state_file, _ = mock_env + state_file.write_text("#!/bin/bash\nxrandr --output HDMI-0 --auto") + # Set insecure permissions (e.g., world-readable) + os.chmod(state_file, 0o644) + + display_switch.restore_state() + + mock_run.assert_not_called() + assert state_file.exists() + +@patch('display_switch.get_current_state') +def test_save_state_command_injection_prevention(mock_get_state, mock_env): + state_file, _ = mock_env + # Malicious monitor name + mock_get_state.return_value = [{ + "name": "HDMI-0; rm -rf /", "primary": False, "active": True, + "pos": "0x0", "rotation": "normal", "active_mode": "1920x1080", "rate": "60.00" + }] + + display_switch.save_state() + + with open(state_file, 'r') as f: + content = f.read() + # The malicious name should be quoted as a single argument + assert "xrandr --output 'HDMI-0; rm -rf /' --mode 1920x1080" in content