activation_scanner_hook.py

#!/usr/bin/env python3
"""Hook-facing wrapper for the cached activation scanner.

This module accepts real-ish install/runtime payloads such as MCP configs,
tool definitions, skills, hook commands, or plain text. It normalizes them into
scanner subjects, redacts secret values, then emits a gate-friendly JSON result.
"""

from __future__ import annotations

import argparse
import hashlib
import json
import sys
import time
from dataclasses import dataclass
from pathlib import Path
from typing import Any

try:
    from .activation_scanner_cli import DEFAULT_ARTIFACT, artifact_complete
    from .activation_scanner_core import CORE_VERSION, DECISION_POLICY_NAME, scan_texts
except ImportError:
    sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
    from research.activation_scanner_cli import DEFAULT_ARTIFACT, artifact_complete  # type: ignore
    from research.activation_scanner_core import CORE_VERSION, DECISION_POLICY_NAME, scan_texts  # type: ignore


DECISION_RANK = {
    "allow": 0,
    "warn": 1,
    "block": 2,
    "quarantine": 3,
}
SENSITIVE_KEY_PARTS = (
    "api_key",
    "apikey",
    "auth",
    "bearer",
    "client_secret",
    "credential",
    "key",
    "password",
    "private",
    "secret",
    "token",
)
TEXT_KEYS = (
    "text",
    "description",
    "prompt",
    "instructions",
    "content",
    "body",
    "schema",
    "command",
    "script",
)


@dataclass(frozen=True)
class ScanSubject:
    subject_id: str
    kind: str
    text: str
    name: str | None = None
    source: str | None = None
    path: str | None = None


def is_sensitive_key(key: str) -> bool:
    lower = key.lower()
    return any(part in lower for part in SENSITIVE_KEY_PARTS)


def redact_for_scan(value: Any, parent_key: str = "") -> Any:
    if isinstance(value, dict):
        redacted: dict[str, Any] = {}
        for key, child in sorted(value.items(), key=lambda item: str(item[0])):
            key_text = str(key)
            if parent_key.lower() == "env" or is_sensitive_key(key_text):
                redacted[key_text] = redact_secret_value(child)
            else:
                redacted[key_text] = redact_for_scan(child, key_text)
        return redacted
    if isinstance(value, list):
        return [redact_for_scan(item, parent_key) for item in value]
    return value


def redact_secret_value(value: Any) -> str:
    if value in (None, ""):
        return "[REDACTED_EMPTY_VALUE]"
    if isinstance(value, str):
        return f"[REDACTED_VALUE len={len(value)}]"
    return f"[REDACTED_{type(value).__name__.upper()}]"


def stable_json(value: Any) -> str:
    return json.dumps(value, sort_keys=True, ensure_ascii=False, separators=(",", ":"))


def pretty_subject_text(value: Any) -> str:
    redacted = redact_for_scan(value)
    if isinstance(redacted, str):
        return redacted
    return json.dumps(redacted, sort_keys=True, ensure_ascii=False, indent=2)


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


def first_string(mapping: dict[str, Any], keys: tuple[str, ...]) -> str | None:
    for key in keys:
        value = mapping.get(key)
        if isinstance(value, str) and value:
            return value
    return None


def object_subject(mapping: dict[str, Any], subject_id: str, inherited_kind: str | None = None) -> ScanSubject:
    kind = first_string(mapping, ("kind", "type", "category", "subType")) or inherited_kind or "object"
    name = first_string(mapping, ("name", "id", "tool_name", "server", "title"))
    source = first_string(mapping, ("source", "origin", "scope"))
    path = first_string(mapping, ("path", "file", "file_path"))

    text: str | None = None
    for key in TEXT_KEYS:
        raw = mapping.get(key)
        if isinstance(raw, str) and raw.strip():
            text = raw
            break
    if text is None:
        text = pretty_subject_text(mapping)
    elif len(mapping) > 1:
        text = pretty_subject_text(mapping)

    return ScanSubject(
        subject_id=subject_id,
        kind=kind,
        name=name,
        source=source,
        path=path,
        text=text,
    )


def normalize_payload(payload: Any, subject_id: str = "input") -> list[ScanSubject]:
    if isinstance(payload, str):
        return [ScanSubject(subject_id=subject_id, kind="text", text=payload)]

    if isinstance(payload, list):
        subjects: list[ScanSubject] = []
        for idx, item in enumerate(payload):
            subjects.extend(normalize_payload(item, f"{subject_id}-{idx + 1}"))
        return subjects

    if not isinstance(payload, dict):
        return [ScanSubject(subject_id=subject_id, kind=type(payload).__name__, text=pretty_subject_text(payload))]

    item_groups = payload.get("items") or payload.get("subjects")
    if isinstance(item_groups, list):
        subjects = []
        for idx, item in enumerate(item_groups):
            subjects.extend(normalize_payload(item, f"{subject_id}-{idx + 1}"))
        return subjects

    for servers_key in ("mcpServers", "mcp_servers"):
        servers = payload.get(servers_key)
        if isinstance(servers, dict):
            subjects = []
            for name, config in sorted(servers.items(), key=lambda item: str(item[0])):
                server_payload = {
                    "kind": "mcp_server",
                    "name": str(name),
                    "mcp_config": config,
                    "source": payload.get("source"),
                    "path": payload.get("path"),
                }
                subjects.append(object_subject(server_payload, f"{subject_id}-mcp-{name}", "mcp_server"))
            return subjects

    tools = payload.get("tools")
    if isinstance(tools, list):
        subjects = []
        for idx, tool in enumerate(tools):
            if isinstance(tool, dict):
                tool_id = first_string(tool, ("name", "id", "tool_name")) or str(idx + 1)
                subjects.append(object_subject({"kind": "tool", **tool}, f"{subject_id}-tool-{tool_id}", "tool"))
            else:
                subjects.extend(normalize_payload(tool, f"{subject_id}-tool-{idx + 1}"))
        return subjects

    tool = payload.get("tool")
    if isinstance(tool, dict):
        tool_id = first_string(tool, ("name", "id", "tool_name")) or "tool"
        return [object_subject({"kind": "tool", **tool}, f"{subject_id}-{tool_id}", "tool")]

    return [object_subject(payload, subject_id)]


def parse_input(raw: str, input_format: str) -> Any:
    if input_format == "text":
        return raw
    if input_format == "json":
        return json.loads(raw)
    stripped = raw.lstrip()
    if stripped.startswith("{") or stripped.startswith("[") or stripped.startswith('"'):
        try:
            return json.loads(raw)
        except json.JSONDecodeError:
            return raw
    return raw


def read_payload(args: argparse.Namespace) -> Any:
    selected = [args.text is not None, args.input_file is not None, args.stdin]
    if sum(selected) > 1:
        raise SystemExit("Use only one of --text, --input-file, or --stdin.")
    if args.text is not None:
        raw = args.text
    elif args.input_file is not None:
        raw = args.input_file.read_text()
    elif args.stdin or not sys.stdin.isatty():
        raw = sys.stdin.read()
    else:
        raise SystemExit("Provide --text, --input-file, or pipe payload on stdin.")
    return parse_input(raw, args.input_format)


def max_decision(results: list[dict[str, Any]]) -> str:
    if not results:
        return "allow"
    return max((str(row.get("decision", "allow")) for row in results), key=lambda decision: DECISION_RANK.get(decision, -1))


def exit_code_for_decision(decision: str, fail_on: str) -> int:
    if fail_on == "never":
        return 0
    return 2 if DECISION_RANK.get(decision, -1) >= DECISION_RANK[fail_on] else 0


def subject_public(subject: ScanSubject) -> dict[str, Any]:
    return {
        "id": subject.subject_id,
        "kind": subject.kind,
        "name": subject.name,
        "source": subject.source,
        "path": subject.path,
        "content_sha256": content_sha256(subject.text),
        "text_length": len(subject.text),
    }


def build_normalized_payload(subjects: list[ScanSubject]) -> dict[str, Any]:
    return {
        "mode": "activation_scanner_hook_normalize",
        "scanner_version": CORE_VERSION,
        "decision_policy": DECISION_POLICY_NAME,
        "count": len(subjects),
        "subjects": [
            {
                "subject": subject_public(subject),
                "text": subject.text,
            }
            for subject in subjects
        ],
    }


def scan_subjects(subjects: list[ScanSubject], args: argparse.Namespace) -> dict[str, Any]:
    if not subjects:
        raise SystemExit("No scanner subjects found in payload.")
    if not artifact_complete(args.artifact):
        raise SystemExit(f"Missing scanner artifact: {args.artifact}")

    started = time.perf_counter()
    risk_results = scan_texts(
        [subject.text for subject in subjects],
        args.artifact,
        batch_size=args.batch_size,
        max_length=args.max_length,
        device=args.device,
        dtype=args.dtype,
        local_files_only=args.local_files_only,
        warn_threshold=args.warn_threshold,
        block_threshold=args.block_threshold,
        top_sae_features=args.top_sae_features,
    )
    results = []
    for subject, risk in zip(subjects, risk_results, strict=True):
        results.append(
            {
                "subject": subject_public(subject),
                "decision": risk.get("decision"),
                "risk_score": risk.get("risk_score"),
                "risk": risk,
            }
        )

    decision = max_decision(results)
    return {
        "mode": "activation_scanner_hook",
        "scanner_version": CORE_VERSION,
        "decision_policy": DECISION_POLICY_NAME,
        "artifact": str(args.artifact),
        "count": len(results),
        "max_decision": decision,
        "gate": {
            "decision": decision,
            "fail_on": args.fail_on,
            "exit_code": exit_code_for_decision(decision, args.fail_on),
        },
        "results": results,
        "elapsed_seconds": time.perf_counter() - started,
    }


def print_json(payload: Any, pretty: bool) -> None:
    print(json.dumps(payload, indent=2 if pretty else None, ensure_ascii=False))


def add_payload_args(parser: argparse.ArgumentParser) -> None:
    parser.add_argument("--text", help="Plain text or JSON payload to scan.")
    parser.add_argument("--input-file", type=Path, help="Read text or JSON payload from a file.")
    parser.add_argument("--stdin", action="store_true", help="Read text or JSON payload from stdin.")
    parser.add_argument("--input-format", choices=("auto", "json", "text"), default="auto")


def add_runtime_args(parser: argparse.ArgumentParser) -> None:
    parser.add_argument("--artifact", type=Path, default=DEFAULT_ARTIFACT, help="Cached scanner artifact directory.")
    parser.add_argument("--batch-size", type=int, default=None)
    parser.add_argument("--max-length", type=int, default=None)
    parser.add_argument("--device", default=None)
    parser.add_argument("--dtype", choices=("auto", "float32", "bfloat16"), default=None)
    parser.add_argument("--local-files-only", action="store_true")
    parser.add_argument("--warn-threshold", type=float, default=None)
    parser.add_argument("--block-threshold", type=float, default=None)
    parser.add_argument("--top-sae-features", type=int, default=8)
    parser.add_argument(
        "--fail-on",
        choices=("never", "warn", "block", "quarantine"),
        default="block",
        help="Set gate.exit_code=2 when max_decision reaches this decision or higher.",
    )
    parser.add_argument("--pretty", action="store_true")


def command_normalize(args: argparse.Namespace) -> int:
    subjects = normalize_payload(read_payload(args))
    print_json(build_normalized_payload(subjects), args.pretty)
    return 0


def command_scan(args: argparse.Namespace) -> int:
    subjects = normalize_payload(read_payload(args))
    payload = scan_subjects(subjects, args)
    print_json(payload, args.pretty)
    return int(payload["gate"]["exit_code"])


def error_payload(message: str, fail_on: str) -> dict[str, Any]:
    decision = "quarantine"
    return {
        "mode": "activation_scanner_hook_error",
        "scanner_version": CORE_VERSION,
        "decision_policy": DECISION_POLICY_NAME,
        "max_decision": decision,
        "gate": {
            "decision": decision,
            "fail_on": fail_on,
            "exit_code": exit_code_for_decision(decision, fail_on),
        },
        "error": message,
    }


def command_serve_jsonl(args: argparse.Namespace) -> int:
    if not artifact_complete(args.artifact):
        raise SystemExit(f"Missing scanner artifact: {args.artifact}")
    if args.warmup:
        warm_subject = ScanSubject("warmup", "text", "Returns today's weather forecast for a requested city as JSON.")
        scan_subjects([warm_subject], args)
        print(
            stable_json(
                {
                    "mode": "activation_scanner_hook_ready",
                    "scanner_version": CORE_VERSION,
                    "decision_policy": DECISION_POLICY_NAME,
                    "artifact": str(args.artifact),
                }
            ),
            file=sys.stderr,
            flush=True,
        )

    for raw_line in sys.stdin:
        raw_line = raw_line.strip()
        if not raw_line:
            continue
        try:
            payload = parse_input(raw_line, args.input_format)
            output = scan_subjects(normalize_payload(payload), args)
        except Exception as exc:  # keep the protocol alive and fail closed per request
            output = error_payload(str(exc), args.fail_on)
        print(stable_json(output), flush=True)
    return 0


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description=__doc__)
    subparsers = parser.add_subparsers(dest="command")

    normalize = subparsers.add_parser("normalize", help="Normalize/redact a hook payload without loading the model.")
    add_payload_args(normalize)
    normalize.add_argument("--pretty", action="store_true")
    normalize.set_defaults(func=command_normalize)

    scan = subparsers.add_parser("scan", help="Scan one hook payload and exit with a gate code.")
    add_payload_args(scan)
    add_runtime_args(scan)
    scan.set_defaults(func=command_scan)

    serve = subparsers.add_parser("serve-jsonl", help="Keep a warm scanner process and scan one JSON/text payload per line.")
    add_runtime_args(serve)
    serve.add_argument("--input-format", choices=("auto", "json", "text"), default="auto")
    serve.add_argument("--warmup", dest="warmup", action="store_true", default=True)
    serve.add_argument("--no-warmup", dest="warmup", action="store_false")
    serve.set_defaults(func=command_serve_jsonl)

    return parser


def main(argv: list[str] | None = None) -> int:
    parser = build_parser()
    args = parser.parse_args(argv)
    if not hasattr(args, "func"):
        parser.print_help()
        return 2
    return int(args.func(args))


if __name__ == "__main__":
    raise SystemExit(main())