bikini-patchwork/patchwork/audit.py

from __future__ import annotations

import ast
from datetime import datetime, timezone
import hashlib
from importlib.metadata import PackageNotFoundError, version
import json
import platform
from pathlib import Path
from typing import Any


SENSITIVE_IMPORTS = {
    "ctypes",
    "marshal",
    "mmap",
    "pickle",
    "socket",
    "subprocess",
    "winreg",
}

SENSITIVE_CALLS = {
    "__import__",
    "compile",
    "eval",
    "exec",
    "marshal.loads",
    "os.execl",
    "os.execle",
    "os.execlp",
    "os.execlpe",
    "os.execv",
    "os.execve",
    "os.execvp",
    "os.execvpe",
    "os.popen",
    "os.spawnl",
    "os.spawnle",
    "os.spawnlp",
    "os.spawnlpe",
    "os.spawnv",
    "os.spawnve",
    "os.spawnvp",
    "os.spawnvpe",
    "os.system",
    "pickle.loads",
    "subprocess.call",
    "subprocess.check_call",
    "subprocess.check_output",
    "subprocess.Popen",
    "subprocess.run",
}


def sha256_text(value: str) -> str:
    return hashlib.sha256(value.encode("utf-8")).hexdigest()


def _utc_now() -> str:
    return datetime.now(timezone.utc).isoformat().replace("+00:00", "Z")


def _tool_version() -> str:
    try:
        return version("patchwork")
    except PackageNotFoundError:
        return "0.2.0"


def _call_name(node: ast.AST) -> str | None:
    if isinstance(node, ast.Name):
        return node.id
    if isinstance(node, ast.Attribute):
        parent = _call_name(node.value)
        if parent:
            return f"{parent}.{node.attr}"
        return node.attr
    return None


class AuditVisitor(ast.NodeVisitor):
    def __init__(self) -> None:
        self.imports: set[str] = set()
        self.functions: list[str] = []
        self.classes: list[str] = []
        self.review_indicators: list[dict[str, object]] = []
        self.node_count = 0

    def generic_visit(self, node: ast.AST) -> None:
        self.node_count += 1
        super().generic_visit(node)

    def visit_Import(self, node: ast.Import) -> None:
        for alias in node.names:
            top_level = alias.name.split(".", 1)[0]
            self.imports.add(alias.name)
            if top_level in SENSITIVE_IMPORTS:
                self.review_indicators.append(
                    {
                        "kind": "import",
                        "name": alias.name,
                        "line": node.lineno,
                        "reason": "sensitive standard-library capability",
                    }
                )
        self.generic_visit(node)

    def visit_ImportFrom(self, node: ast.ImportFrom) -> None:
        module = node.module or ""
        if module:
            self.imports.add(module)
        top_level = module.split(".", 1)[0]
        if top_level in SENSITIVE_IMPORTS:
            self.review_indicators.append(
                {
                    "kind": "import",
                    "name": module,
                    "line": node.lineno,
                    "reason": "sensitive standard-library capability",
                }
            )
        self.generic_visit(node)

    def visit_FunctionDef(self, node: ast.FunctionDef) -> None:
        self.functions.append(node.name)
        self.generic_visit(node)

    def visit_AsyncFunctionDef(self, node: ast.AsyncFunctionDef) -> None:
        self.functions.append(node.name)
        self.generic_visit(node)

    def visit_ClassDef(self, node: ast.ClassDef) -> None:
        self.classes.append(node.name)
        self.generic_visit(node)

    def visit_Call(self, node: ast.Call) -> None:
        name = _call_name(node.func)
        if name in SENSITIVE_CALLS:
            self.review_indicators.append(
                {
                    "kind": "call",
                    "name": name,
                    "line": node.lineno,
                    "reason": "dynamic execution or process capability",
                }
            )
        self.generic_visit(node)


def analyze_source(source: str, *, path: str | Path | None = None) -> dict[str, Any]:
    tree = ast.parse(source, filename=str(path or "<memory>"))
    visitor = AuditVisitor()
    visitor.visit(tree)
    review_indicators = sorted(visitor.review_indicators, key=lambda item: (int(item["line"]), str(item["name"])))
    return {
        "path": str(path) if path else None,
        "sha256": sha256_text(source),
        "lines": source.count("\n") + (1 if source else 0),
        "bytes_utf8": len(source.encode("utf-8")),
        "ast_nodes": visitor.node_count,
        "imports": sorted(visitor.imports),
        "functions": sorted(visitor.functions),
        "classes": sorted(visitor.classes),
        "review": {
            "indicator_count": len(review_indicators),
            "requires_human_review": bool(review_indicators),
            "indicators": review_indicators,
        },
    }


def build_manifest(
    *,
    input_path: str | Path,
    output_path: str | Path,
    input_source: str,
    output_source: str,
    seed: int,
    options: dict[str, object],
    audit: dict[str, Any],
) -> dict[str, Any]:
    return {
        "schema": "patchwork.manifest.v1",
        "generated_at": _utc_now(),
        "tool": {
            "name": "patchwork",
            "version": _tool_version(),
            "python": platform.python_version(),
            "platform": platform.platform(),
        },
        "input": {
            "path": str(input_path),
            "sha256": sha256_text(input_source),
            "bytes_utf8": len(input_source.encode("utf-8")),
        },
        "output": {
            "path": str(output_path),
            "sha256": sha256_text(output_source),
            "bytes_utf8": len(output_source.encode("utf-8")),
        },
        "build": {
            "seed": seed,
            "options": options,
        },
        "audit": audit,
    }


def write_json(path: str | Path, data: dict[str, Any]) -> Path:
    output_path = Path(path)
    output_path.parent.mkdir(parents=True, exist_ok=True)
    output_path.write_text(json.dumps(data, indent=2, sort_keys=True) + "\n", encoding="utf-8")
    return output_path