AlgoTune/algotune_agent

"""AlgoTune Agent's edit-time guard against scorer and runtime tampering.

This is the AST policy from ``AlgoTuner/security/code_validator.py``. It lives inside the
scaffold so enabling AlgoTune does not change the parser or evaluator used by any other scaffold.
"""
from __future__ import annotations

import ast
from pathlib import PurePath


class TamperingDetector(ast.NodeVisitor):
    """Detect the runtime mutation and code-evasion patterns rejected upstream."""

    PROTECTED_MODULES = {
        "hmac", "hashlib", "os", "sys", "subprocess", "importlib", "builtins",
        "__builtins__", "types", "gc", "inspect", "base64", "gzip", "random", "time",
        "tempfile", "math", "ast", "collections", "enum", "itertools", "logging",
        "multiprocessing", "numbers", "re", "string", "queue", "traceback", "typing",
        "numpy", "scipy", "cvxpy", "networkx", "ortools", "sklearn", "sympy", "pysat",
        "mpmath", "numba", "cryptography", "faiss", "hdbscan", "ot", "orjson",
    }
    PROTECTED_ATTRIBUTES = {
        "hmac.compare_digest", "hashlib.sha256", "hashlib.sha512", "hashlib.md5",
        "os.system", "subprocess.run", "sys.modules",
    }
    DISALLOWED_CALLS = {"exec", "eval", "compile", "gc.get_objects"}

    def __init__(self):
        self.violations: list[dict] = []
        self.import_aliases: dict[str, str] = {}

    def _record_alias(self, bound_name: str, imported_name: str) -> None:
        if bound_name:
            self.import_aliases[bound_name] = imported_name

    def _name(self, node) -> str | None:
        if isinstance(node, ast.Name):
            return self.import_aliases.get(node.id, node.id)
        if isinstance(node, ast.Attribute):
            base = self._name(node.value)
            if base:
                return f"{base}.{node.attr}"
        return None

    @staticmethod
    def _string(node) -> str | None:
        return node.value if isinstance(node, ast.Constant) and isinstance(node.value, str) else None

    def _protected(self, name: str | None) -> bool:
        return bool(name) and (
            name in self.PROTECTED_MODULES or name.split(".", 1)[0] in self.PROTECTED_MODULES
        )

    def _violation(self, node, kind: str, module: str, attribute: str, code: str,
                   message: str | None = None) -> None:
        item = {
            "line": node.lineno,
            "type": kind,
            "module": module,
            "attribute": attribute,
            "code": code,
        }
        if message:
            item["message"] = message
        self.violations.append(item)

    def visit_Import(self, node: ast.Import) -> None:
        for alias in node.names:
            imported = alias.name or ""
            self._record_alias(alias.asname or imported.split(".", 1)[0], imported)
            if imported.split(".", 1)[0] == "ctypes":
                self._violation(
                    node, "import", "ctypes", "<module>", "import ctypes",
                    "Importing ctypes is not allowed in solver code.",
                )
        self.generic_visit(node)

    def visit_ImportFrom(self, node: ast.ImportFrom) -> None:
        for alias in node.names:
            if alias.name != "*":
                imported = f"{node.module}.{alias.name}" if node.module else alias.name
                self._record_alias(alias.asname or alias.name, imported)
        if (node.module or "").split(".", 1)[0] == "ctypes":
            self._violation(
                node, "import", "ctypes", "<module>", f"from {node.module} import ...",
                "Importing ctypes is not allowed in solver code.",
            )
        self.generic_visit(node)

    def visit_Subscript(self, node: ast.Subscript) -> None:
        if self._name(node.value) == "sys.modules":
            self._violation(
                node, "subscript", "sys", "modules", "sys.modules[...]",
                "Accessing sys.modules is not allowed in solver code.",
            )
        self.generic_visit(node)

    def _assignment(self, target, node) -> None:
        if isinstance(target, ast.Attribute):
            module = self._name(target.value)
            if target.attr == "is_solution":
                self._violation(
                    node, "assignment", module or "<dynamic>", "is_solution",
                    f"{module or '<dynamic>'}.is_solution = ...",
                    "Overriding is_solution at runtime is not allowed.",
                )
            elif self._protected(module):
                self._violation(
                    node, "assignment", module or "<unknown>", target.attr,
                    f"{module}.{target.attr} = ...",
                )
        elif isinstance(target, ast.Subscript) and isinstance(target.value, ast.Attribute):
            name = self._name(target.value)
            if name in self.PROTECTED_ATTRIBUTES:
                module, _, attribute = (name or "").partition(".")
                self._violation(
                    node, "subscript_assignment", module, attribute or "<unknown>",
                    f"{name}[...] = ...",
                )
            elif target.value.attr == "__dict__":
                module = self._name(target.value.value)
                if self._protected(module):
                    self._violation(
                        node, "dict_assignment", module or "<unknown>", "<unknown>",
                        f"{module}.__dict__[...] = ...",
                    )

    def visit_Assign(self, node: ast.Assign) -> None:
        for target in node.targets:
            self._assignment(target, node)
        self.generic_visit(node)

    def visit_AugAssign(self, node: ast.AugAssign) -> None:
        self._assignment(node.target, node)
        self.generic_visit(node)

    def visit_Call(self, node: ast.Call) -> None:
        name = self._name(node.func)
        if name in self.DISALLOWED_CALLS:
            self._violation(
                node, "call", name.split(".", 1)[0], name.split(".", 1)[-1], f"{name}(...)",
                f"Calling {name} is not allowed in solver code.",
            )
        if isinstance(node.func, ast.Name) and node.func.id == "__import__" and node.args:
            imported = self._string(node.args[0])
            if imported and imported.split(".", 1)[0] == "ctypes":
                self._violation(
                    node, "dynamic_import", "ctypes", "<module>", "__import__('ctypes')",
                    "Dynamically importing ctypes is not allowed in solver code.",
                )
        if name and name.startswith("sys.modules"):
            self._violation(
                node, "call", "sys", "modules", f"{name}()",
                "Accessing sys.modules is not allowed in solver code.",
            )
        if isinstance(node.func, ast.Name) and node.func.id == "setattr" and len(node.args) >= 2:
            module = self._name(node.args[0])
            attribute = self._string(node.args[1])
            if self._protected(module):
                self._violation(
                    node, "setattr", module or "<unknown>", attribute or "<dynamic>",
                    f"setattr({module}, ...)",
                )
            if attribute == "is_solution":
                self._violation(
                    node, "setattr", module or "<dynamic>", "is_solution",
                    "setattr(..., 'is_solution', ...)",
                    "Overriding is_solution at runtime is not allowed.",
                )
        elif isinstance(node.func, ast.Attribute) and node.func.attr == "__setattr__":
            module = self._name(node.func.value)
            if self._protected(module):
                self._violation(
                    node, "setattr", module or "<unknown>", "<unknown>",
                    f"{module}.__setattr__(...)",
                )
        self.generic_visit(node)


def tampering_error(tree: ast.AST, source: str) -> str | None:
    detector = TamperingDetector()
    detector.visit(tree)
    if not detector.violations:
        return None
    lines = source.splitlines()
    output = ["Error: Code contains security violations"]
    for violation in detector.violations:
        number = int(violation["line"])
        content = lines[number - 1].strip() if number <= len(lines) else ""
        message = violation.get("message") or (
            f"Tampering with {violation.get('module', '<unknown>')}."
            f"{violation.get('attribute', '')}"
        )
        output.extend([
            f"\nLine {number}: {content}",
            f"  {message}",
            f"  Detected: {violation['code']}",
        ])
    output.append(
        "\nOnly algorithmic solver logic is allowed. Runtime tampering, dynamic code injection, "
        "and harness introspection are prohibited."
    )
    return "\n".join(output)


def protected_filename_error(filename: str) -> str | None:
    """Reject files that could shadow the modules protected by the upstream validator."""
    path = PurePath(filename)
    name = path.name
    module = name[:-3].lower() if name.lower().endswith(".py") else ""
    protected = {item.split(".", 1)[0].lower() for item in TamperingDetector.PROTECTED_MODULES}
    protected.update({"sitecustomize", "usercustomize"})
    protected_parent = next(
        (part.lower() for part in path.parts[:-1] if part.lower() in protected), None
    )
    if name == "__init__.py" and protected_parent:
        module = protected_parent
    if module in protected:
        return (
            f"That is not allowed: editing or creating {name!r} would shadow "
            f"the protected module {module!r}."
        )
    return None