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 sha256_file(path: str | Path) -> str:
    digest = hashlib.sha256()
    with Path(path).open("rb") as handle:
        for chunk in iter(lambda: handle.read(1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def sha256_text_file(path: str | Path) -> str:
    return sha256_text(Path(path).read_text(encoding="utf-8"))


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.4.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_stats(input_source: str, output_source: str | None, audit: dict[str, Any]) -> dict[str, object]:
    input_bytes = len(input_source.encode("utf-8"))
    output_bytes = len(output_source.encode("utf-8")) if output_source is not None else None
    return {
        "input_bytes": input_bytes,
        "output_bytes": output_bytes,
        "ratio": round(output_bytes / max(input_bytes, 1), 4) if output_bytes is not None else None,
        "input_lines": input_source.count("\n") + (1 if input_source else 0),
        "output_lines": output_source.count("\n") + (1 if output_source else 0) if output_source is not None else None,
        "input_sha256": sha256_text(input_source),
        "output_sha256": sha256_text(output_source) if output_source is not None else None,
        "review_indicator_count": audit.get("review", {}).get("indicator_count", 0),
    }


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": {
            "cwd": str(Path.cwd()),
            "seed": seed,
            "options": options,
        },
        "stats": build_stats(input_source, output_source, audit),
        "audit": audit,
    }


def verify_manifest(path: str | Path) -> dict[str, Any]:
    manifest_path = Path(path)
    try:
        manifest = json.loads(manifest_path.read_text(encoding="utf-8"))
    except FileNotFoundError as exc:
        raise ValueError(f"manifest file not found: {manifest_path}") from exc
    except json.JSONDecodeError as exc:
        raise ValueError(f"invalid manifest JSON: {exc}") from exc
    if not isinstance(manifest, dict):
        raise ValueError("manifest must contain a JSON object")

    errors: list[str] = []
    if manifest.get("schema") != "patchwork.manifest.v1":
        errors.append("unsupported manifest schema")

    build = manifest.get("build", {})
    cwd = Path(build.get("cwd") or Path.cwd()) if isinstance(build, dict) else Path.cwd()
    checks: dict[str, dict[str, object]] = {}
    for section in ("input", "output"):
        item = manifest.get(section, {})
        if not isinstance(item, dict):
            errors.append(f"{section} section missing or invalid")
            continue
        raw_path = item.get("path")
        expected = item.get("sha256")
        if not isinstance(raw_path, str) or not isinstance(expected, str):
            errors.append(f"{section} path or sha256 missing")
            continue
        candidate = Path(raw_path)
        if not candidate.is_absolute():
            candidate = cwd / candidate
        exists = candidate.exists()
        try:
            actual = sha256_text_file(candidate) if exists else None
        except UnicodeDecodeError:
            actual = None
        ok = actual == expected
        if not ok:
            errors.append(f"{section} hash mismatch or file missing: {candidate}")
        checks[section] = {
            "path": str(candidate),
            "exists": exists,
            "expected_sha256": expected,
            "actual_sha256": actual,
            "ok": ok,
        }

    return {
        "valid": not errors,
        "errors": errors,
        "checks": checks,
        "manifest": str(manifest_path),
    }


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