test_workspace_security.py 7.07 KB
"""U-036: real unittest coverage for workspace_security.py -- adapted from
graphloop's test_workspace_security.py (2026-08-02). Module logic is
unchanged from graphloop ("adapted verbatim", see workspace_security.py's
own docstring); only the import path changed.

Run: python3 -m unittest discover -s supervisor/isolation -p "test_*.py" -v
"""

import os
import shutil
import sys
import tempfile
import unittest

from supervisor.isolation import workspace_security as ws


class ValidateArtifactsTests(unittest.TestCase):
    def setUp(self):
        self.workspace = tempfile.mkdtemp(prefix="wssec-workspace-")
        self.addCleanup(shutil.rmtree, self.workspace, ignore_errors=True)

        self.outside_dir = tempfile.mkdtemp(prefix="wssec-outside-")
        self.addCleanup(shutil.rmtree, self.outside_dir, ignore_errors=True)

    def _write(self, relative_path, content="content"):
        full_path = os.path.join(self.workspace, relative_path)
        os.makedirs(os.path.dirname(full_path), exist_ok=True)
        with open(full_path, "w", encoding="utf-8") as handle:
            handle.write(content)
        return full_path

    def test_accepts_and_normalizes_plain_relative_artifacts(self):
        self._write("output.txt")
        self._write(os.path.join("nested", "dir", "file.txt"))

        result = ws.validate_artifacts(["output.txt", "nested/dir/file.txt"], self.workspace)

        self.assertEqual(result, ["output.txt", "nested/dir/file.txt"])

    def test_normalizes_redundant_leading_dot_slash(self):
        self._write("output.txt")

        result = ws.validate_artifacts(["./output.txt"], self.workspace)

        self.assertEqual(result, ["output.txt"])

    def test_rejects_non_string_item(self):
        self._write("output.txt")

        with self.assertRaises(ws.WorkspaceSecurityError):
            ws.validate_artifacts(["output.txt", 42], self.workspace)

    def test_rejects_non_list_artifacts_argument(self):
        with self.assertRaises(ws.WorkspaceSecurityError):
            ws.validate_artifacts("output.txt", self.workspace)

    def test_rejects_absolute_path_artifact(self):
        absolute = self._write("output.txt")

        with self.assertRaises(ws.WorkspaceSecurityError):
            ws.validate_artifacts([absolute], self.workspace)

    def test_rejects_dot_dot_traversal_even_when_target_exists(self):
        self._write("output.txt")
        os.makedirs(os.path.join(self.workspace, "sibling"), exist_ok=True)

        with self.assertRaises(ws.WorkspaceSecurityError):
            ws.validate_artifacts(["sibling/../output.txt"], self.workspace)

    def test_rejects_dot_dot_traversal_escaping_workspace(self):
        secret = os.path.join(self.outside_dir, "secret.txt")
        with open(secret, "w", encoding="utf-8") as handle:
            handle.write("secret")

        with self.assertRaises(ws.WorkspaceSecurityError):
            ws.validate_artifacts(
                [os.path.join("..", os.path.basename(self.outside_dir), "secret.txt")],
                self.workspace,
            )

    def test_rejects_missing_artifact(self):
        with self.assertRaises(ws.WorkspaceSecurityError):
            ws.validate_artifacts(["does_not_exist.txt"], self.workspace)

    def test_rejects_symlink_artifact(self):
        real_file = self._write("real.txt")
        symlink_path = os.path.join(self.workspace, "link.txt")
        os.symlink(real_file, symlink_path)
        self.addCleanup(lambda: os.path.exists(symlink_path) and os.remove(symlink_path))

        with self.assertRaises(ws.WorkspaceSecurityError):
            ws.validate_artifacts(["link.txt"], self.workspace)

    def test_rejects_artifact_resolving_outside_workspace_via_symlinked_directory(self):
        secret = os.path.join(self.outside_dir, "secret.txt")
        with open(secret, "w", encoding="utf-8") as handle:
            handle.write("secret")

        symlinked_dir = os.path.join(self.workspace, "escape")
        os.symlink(self.outside_dir, symlinked_dir, target_is_directory=True)
        self.addCleanup(lambda: os.path.exists(symlinked_dir) and os.remove(symlinked_dir))

        with self.assertRaises(ws.WorkspaceSecurityError):
            ws.validate_artifacts(["escape/secret.txt"], self.workspace)

    def test_rejects_empty_string_artifact(self):
        with self.assertRaises(ws.WorkspaceSecurityError):
            ws.validate_artifacts([""], self.workspace)


class ValidateSubprocessCommandTests(unittest.TestCase):
    def setUp(self):
        self.workspace = tempfile.mkdtemp(prefix="wssec-cmd-workspace-")
        self.addCleanup(shutil.rmtree, self.workspace, ignore_errors=True)

        self.adapter_script = os.path.join(self.workspace, "existing_adapter.py")
        with open(self.adapter_script, "w", encoding="utf-8") as handle:
            handle.write("print('ok')\n")

    def test_accepts_python_executable_with_existing_script(self):
        result = ws.validate_subprocess_command([sys.executable, self.adapter_script])

        self.assertEqual(result, [sys.executable, self.adapter_script])

    def test_rejects_empty_command(self):
        with self.assertRaises(ws.WorkspaceSecurityError):
            ws.validate_subprocess_command([])

    def test_rejects_non_list_command(self):
        with self.assertRaises(ws.WorkspaceSecurityError):
            ws.validate_subprocess_command("{} {}".format(sys.executable, self.adapter_script))

    def test_rejects_non_string_argument(self):
        with self.assertRaises(ws.WorkspaceSecurityError):
            ws.validate_subprocess_command([sys.executable, self.adapter_script, 7])

    def test_rejects_shell_metacharacters_in_arguments(self):
        dangerous_args = [
            "foo; rm -rf /", "foo && rm -rf /", "foo | rm -rf /",
            "foo > /etc/passwd", "foo `rm -rf /`", "foo $(rm -rf /)",
        ]
        for dangerous_arg in dangerous_args:
            with self.subTest(dangerous_arg=dangerous_arg):
                with self.assertRaises(ws.WorkspaceSecurityError):
                    ws.validate_subprocess_command([sys.executable, self.adapter_script, dangerous_arg])

    def test_rejects_bare_shell_invocation(self):
        with self.assertRaises(ws.WorkspaceSecurityError):
            ws.validate_subprocess_command(["/bin/sh", "-c", "echo hi"])

    def test_rejects_script_argument_with_disallowed_extension(self):
        non_script = os.path.join(self.workspace, "not_a_script.txt")
        with open(non_script, "w", encoding="utf-8") as handle:
            handle.write("not python")

        with self.assertRaises(ws.WorkspaceSecurityError):
            ws.validate_subprocess_command([sys.executable, non_script])

    def test_rejects_script_argument_that_does_not_exist(self):
        missing_script = os.path.join(self.workspace, "missing_adapter.py")

        with self.assertRaises(ws.WorkspaceSecurityError):
            ws.validate_subprocess_command([sys.executable, missing_script])

    def test_rejects_newline_in_argument(self):
        with self.assertRaises(ws.WorkspaceSecurityError):
            ws.validate_subprocess_command([sys.executable, self.adapter_script, "foo\nbar"])


if __name__ == "__main__":
    unittest.main()