🔒 Fix security vulnerabilities in display_switch.py

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}")

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