activation_scanner_core.py

#!/usr/bin/env python3
"""Runtime scanner that scores text with a cached probe artifact."""

from __future__ import annotations

import argparse
import json
import math
import time
from pathlib import Path
from typing import Any

import numpy as np

try:
    from .benchmarks.model_registry import extract_features, extract_sae_features
    from .scanner_static import static_scan, static_score
except ImportError:
    import sys

    sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
    from research.benchmarks.model_registry import extract_features, extract_sae_features  # type: ignore
    from research.scanner_static import static_scan, static_score  # type: ignore


CORE_VERSION = "activation-scanner-core-2026-06-03-static-v3"
DECISION_POLICY_NAME = "corroborated-block-v3"
STATIC_BLOCK_BUNDLE_MAX_SPAN = 900
STATIC_BLOCK_BUNDLES = (
    frozenset(("static-secret-file", "static-exfiltration", "static-hidden-action")),
)
STATIC_ACTIVATION_CORROBORATORS = {
    "static-exfiltration",
    "static-hidden-action",
    "static-host-modification",
}


def read_input_text(args: argparse.Namespace) -> str:
    if args.text and args.file:
        raise SystemExit("Use either --text or --file, not both.")
    if args.text:
        return args.text
    if args.file:
        return Path(args.file).read_text()
    raise SystemExit("Provide --text or --file.")


def read_batch_inputs(path: Path) -> list[tuple[str | None, str]]:
    rows = json.loads(path.read_text())
    if not isinstance(rows, list):
        raise SystemExit("--batch-file must contain a JSON array.")
    out: list[tuple[str | None, str]] = []
    for idx, row in enumerate(rows):
        if isinstance(row, str):
            out.append((None, row))
            continue
        if isinstance(row, dict) and isinstance(row.get("text"), str):
            raw_id = row.get("id")
            out.append((str(raw_id) if raw_id is not None else None, row["text"]))
            continue
        raise SystemExit(f"Batch row {idx} must be a string or object with a text field.")
    if not out:
        raise SystemExit("--batch-file must contain at least one item.")
    return out


def artifact_dir_from_path(path: Path) -> Path:
    if path.is_file():
        return path.parent
    return path


def load_artifact(path: Path) -> tuple[Path, dict[str, Any], dict[str, np.ndarray]]:
    artifact_dir = artifact_dir_from_path(path)
    metadata_path = artifact_dir / "metadata.json"
    weights_path = artifact_dir / "probe_weights.npz"
    if not metadata_path.exists():
        raise SystemExit(f"Missing artifact metadata: {metadata_path}")
    if not weights_path.exists():
        raise SystemExit(f"Missing artifact weights: {weights_path}")

    metadata = json.loads(metadata_path.read_text())
    with np.load(weights_path) as data:
        weights = {
            "coef": np.asarray(data["coef"], dtype=np.float64),
            "intercept": np.asarray(data["intercept"], dtype=np.float64),
            "classes": np.asarray(data["classes"], dtype=np.int64),
        }
    return artifact_dir, metadata, weights


def runtime_value(cli_value: Any, metadata_value: Any, default: Any) -> Any:
    if cli_value is not None:
        return cli_value
    if metadata_value is not None:
        return metadata_value
    return default


def extract_runtime_matrix_for_texts(
    *,
    texts: list[str],
    metadata: dict[str, Any],
    args: argparse.Namespace,
) -> tuple[np.ndarray, dict[str, Any]]:
    if not texts:
        raise SystemExit("Need at least one text to scan.")

    sensor_model = metadata["sensor_model"]
    feature_details = metadata.get("feature_details", {})
    batch_size = runtime_value(args.batch_size, feature_details.get("batch_size"), 16)
    max_length = runtime_value(args.max_length, feature_details.get("max_length"), 256)
    device = runtime_value(args.device, feature_details.get("device"), "cpu")
    dtype = runtime_value(args.dtype, feature_details.get("dtype"), "float32")
    extractor = metadata.get("extractor") or ("sae" if metadata.get("feature_kind") == "sae" else "raw")

    if extractor == "sae":
        sae = metadata.get("sae") or {}
        sae_name = sae.get("name")
        if not sae_name:
            raise SystemExit("SAE artifact metadata is missing sae.name")
        bundle = extract_sae_features(
            sensor_model,
            sae_name,
            texts,
            batch_size=int(batch_size),
            max_length=int(max_length),
            device=str(device),
            dtype=str(dtype),
            local_files_only=args.local_files_only,
        )
        layer = bundle.layers[0]
        matrix = bundle.features_by_layer[layer]
    else:
        selected_layers = tuple(int(layer) for layer in metadata["layer_policy"]["selected_layers"])
        bundle = extract_features(
            sensor_model,
            texts,
            layers=selected_layers,
            layer_sweep=False,
            batch_size=int(batch_size),
            max_length=int(max_length),
            device=str(device),
            dtype=str(dtype),
            local_files_only=args.local_files_only,
        )
        matrix = np.concatenate([bundle.features_by_layer[layer] for layer in selected_layers], axis=1)

    expected_dim = int(metadata["feature_dim"])
    if matrix.shape[1] != expected_dim:
        raise SystemExit(f"Feature dimension mismatch: artifact expects {expected_dim}, runtime got {matrix.shape[1]}")

    details = {
        "feature_kind": bundle.feature_kind,
        "extractor": extractor,
        "feature_elapsed_seconds": bundle.elapsed_seconds,
        "feature_details": bundle.details,
        "layer_policy": metadata.get("layer_policy"),
    }
    return matrix, details


def extract_runtime_matrix(
    *,
    text: str,
    metadata: dict[str, Any],
    args: argparse.Namespace,
) -> tuple[np.ndarray, dict[str, Any]]:
    return extract_runtime_matrix_for_texts(texts=[text], metadata=metadata, args=args)


def sigmoid(value: float) -> float:
    if value >= 0:
        z = math.exp(-value)
        return 1.0 / (1.0 + z)
    z = math.exp(value)
    return z / (1.0 + z)


def softmax(row: np.ndarray) -> np.ndarray:
    shifted = row - np.max(row)
    exp = np.exp(shifted)
    return exp / exp.sum()


def positive_probabilities(matrix: np.ndarray, weights: dict[str, np.ndarray]) -> np.ndarray:
    coef = weights["coef"]
    intercept = weights["intercept"]
    classes = [int(value) for value in weights["classes"]]
    logits = matrix @ coef.T + intercept

    if logits.shape[1] == 1:
        raw = logits[:, 0]
        positive_class = classes[-1] if classes else 1
        probs = np.asarray([sigmoid(float(value)) for value in raw], dtype=np.float64)
        return probs if positive_class == 1 else 1.0 - probs

    if 1 not in classes:
        return np.zeros(logits.shape[0], dtype=np.float64)
    class_idx = classes.index(1)
    return np.asarray([softmax(row)[class_idx] for row in logits], dtype=np.float64)


def positive_probability(matrix: np.ndarray, weights: dict[str, np.ndarray]) -> float:
    return float(positive_probabilities(matrix, weights)[0])


def top_sae_feature_contributions(matrix: np.ndarray, weights: dict[str, np.ndarray], limit: int) -> list[dict[str, Any]]:
    if limit <= 0 or weights["coef"].shape[0] != 1:
        return []
    values = matrix[0]
    coef = weights["coef"][0]
    contributions = values * coef
    positive = np.flatnonzero(contributions > 0)
    if len(positive) == 0:
        return []
    ranked = positive[np.argsort(contributions[positive])[::-1]][:limit]
    return [
        {
            "feature_index": int(idx),
            "activation_value": float(values[idx]),
            "probe_weight": float(coef[idx]),
            "contribution": float(contributions[idx]),
        }
        for idx in ranked
    ]


def build_reasons(
    *,
    activation_score: float,
    static_findings: list[dict[str, Any]],
    feature_kind: str,
    artifact_id: str,
) -> list[str]:
    reasons = [f"cached {feature_kind} probe {artifact_id} score={activation_score:.3f}"]
    for finding in static_findings:
        reasons.append(f"{finding['severity']} static finding {finding['id']}: {finding['reason']}")
    if not static_findings and activation_score < 0.5:
        reasons.append("No strong cached-probe or static keyword signal.")
    return reasons


def static_finding_ids(static_findings: list[dict[str, Any]]) -> set[str]:
    return {str(finding.get("id")) for finding in static_findings if finding.get("id")}


def static_block_bundle(static_findings: list[dict[str, Any]]) -> frozenset[str] | None:
    ids = static_finding_ids(static_findings)
    for bundle in STATIC_BLOCK_BUNDLES:
        if not bundle.issubset(ids):
            continue
        bundle_findings = [row for row in static_findings if str(row.get("id")) in bundle]
        starts = [int(row.get("start", 0)) for row in bundle_findings]
        ends = [int(row.get("end", 0)) for row in bundle_findings]
        if starts and ends and max(ends) - min(starts) <= STATIC_BLOCK_BUNDLE_MAX_SPAN:
            return bundle
    return None


def activation_static_corroborated(static_findings: list[dict[str, Any]]) -> bool:
    ids = static_finding_ids(static_findings)
    if ids & STATIC_ACTIVATION_CORROBORATORS:
        return True
    return False


def choose_runtime_decision(
    *,
    activation_score: float,
    static_risk_score: float,
    risk_score: float,
    static_findings: list[dict[str, Any]],
    warn_threshold: float,
    block_threshold: float,
) -> tuple[str, list[str]]:
    """Choose the product decision without treating one high score as a hard block.

    Calibration showed that activation-only and single-keyword high scores can
    be useful review signals while still being too noisy for automatic block.
    The block tier therefore needs corroboration.
    """

    reasons: list[str] = []
    bundle = static_block_bundle(static_findings)
    if bundle:
        reasons.append(f"static block bundle: {','.join(sorted(bundle))}")
        return "block", reasons

    if (
        activation_score >= block_threshold
        and static_risk_score >= warn_threshold
        and activation_static_corroborated(static_findings)
    ):
        reasons.append(
            "activation block score corroborated by static finding "
            f"(activation={activation_score:.3f}, static={static_risk_score:.3f})"
        )
        return "block", reasons

    if risk_score >= warn_threshold:
        if risk_score >= block_threshold:
            reasons.append("high risk score without block corroboration; downgraded to warn")
        else:
            reasons.append("risk score reached warn threshold")
        return "warn", reasons

    reasons.append("risk score below warn threshold")
    return "allow", reasons


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--artifact", type=Path, required=True, help="Artifact directory or metadata.json path.")
    parser.add_argument("--text", help="Tool, MCP, skill, plugin, or hook description to scan.")
    parser.add_argument("--file", type=Path, help="Read text to scan from a local file.")
    parser.add_argument("--batch-file", type=Path, help="Read a JSON array of strings or {id,text} objects.")
    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("--pretty", action="store_true")
    return parser


def runtime_args(
    *,
    batch_size: int | None = None,
    max_length: int | None = None,
    device: str | None = None,
    dtype: str | None = None,
    local_files_only: bool = False,
    warn_threshold: float | None = None,
    block_threshold: float | None = None,
    top_sae_features: int = 8,
) -> argparse.Namespace:
    return argparse.Namespace(
        batch_size=batch_size,
        max_length=max_length,
        device=device,
        dtype=dtype,
        local_files_only=local_files_only,
        warn_threshold=warn_threshold,
        block_threshold=block_threshold,
        top_sae_features=top_sae_features,
    )


def build_risk_output(
    *,
    text: str,
    artifact_dir: Path,
    metadata: dict[str, Any],
    weights: dict[str, np.ndarray],
    matrix_row: np.ndarray,
    activation_score: float,
    activation_details: dict[str, Any],
    warn_threshold: float,
    block_threshold: float,
    top_sae_features: int,
    elapsed_seconds: float,
) -> dict[str, Any]:
    findings = static_scan(text)
    static_risk_score = static_score(findings)
    risk_score = max(activation_score, static_risk_score)
    decision, decision_policy_reasons = choose_runtime_decision(
        activation_score=activation_score,
        static_risk_score=static_risk_score,
        risk_score=risk_score,
        static_findings=findings,
        warn_threshold=warn_threshold,
        block_threshold=block_threshold,
    )

    sae_features = []
    if metadata.get("feature_kind") == "sae":
        sae_features = top_sae_feature_contributions(matrix_row, weights, top_sae_features)

    artifact_id = metadata.get("artifact_id", artifact_dir.name)
    return {
        "decision": decision,
        "risk_score": risk_score,
        "risk_reasons": build_reasons(
            activation_score=activation_score,
            static_findings=findings,
            feature_kind=str(metadata.get("feature_kind")),
            artifact_id=str(artifact_id),
        ),
        "activation_score": activation_score,
        "text_baseline_score": None,
        "static_score": static_risk_score,
        "sae_features": sae_features,
        "evidence_spans": findings,
        "scanner_version": CORE_VERSION,
        "artifact_id": artifact_id,
        "artifact_path": str(artifact_dir),
        "model_id": metadata.get("model_id"),
        "sensor_model": metadata.get("sensor_model"),
        "sae": metadata.get("sae", {}).get("name") if metadata.get("sae") else None,
        "mode": "cached_probe_runtime",
        "training_data": metadata.get("dataset"),
        "activation_details": activation_details,
        "thresholds": {
            "warn": warn_threshold,
            "block": block_threshold,
        },
        "decision_policy": {
            "name": DECISION_POLICY_NAME,
            "block_requires": (
                "nearby high-confidence static bundle or activation block score "
                "corroborated by action-oriented exfiltration, hidden-action, "
                "or host-modification finding"
            ),
            "static_block_bundles": [sorted(bundle) for bundle in STATIC_BLOCK_BUNDLES],
            "static_block_bundle_max_span": STATIC_BLOCK_BUNDLE_MAX_SPAN,
            "static_activation_corroborators": sorted(STATIC_ACTIVATION_CORROBORATORS),
            "policy_reasons": decision_policy_reasons,
        },
        "elapsed_seconds": elapsed_seconds,
    }


def scan_texts(
    target_texts: list[str],
    artifact: Path,
    *,
    batch_size: int | None = None,
    max_length: int | None = None,
    device: str | None = None,
    dtype: str | None = None,
    local_files_only: bool = False,
    warn_threshold: float | None = None,
    block_threshold: float | None = None,
    top_sae_features: int = 8,
) -> list[dict[str, Any]]:
    started = time.perf_counter()
    args = runtime_args(
        batch_size=batch_size,
        max_length=max_length,
        device=device,
        dtype=dtype,
        local_files_only=local_files_only,
        warn_threshold=warn_threshold,
        block_threshold=block_threshold,
        top_sae_features=top_sae_features,
    )

    artifact_dir, metadata, weights = load_artifact(artifact)
    matrix, activation_details = extract_runtime_matrix_for_texts(texts=target_texts, metadata=metadata, args=args)
    activation_scores = positive_probabilities(matrix, weights)

    thresholds = metadata.get("thresholds", {})
    warn = float(warn_threshold if warn_threshold is not None else thresholds.get("warn", 0.30))
    block = float(block_threshold if block_threshold is not None else thresholds.get("block", 0.85))
    elapsed_seconds = time.perf_counter() - started

    return [
        build_risk_output(
            text=text,
            artifact_dir=artifact_dir,
            metadata=metadata,
            weights=weights,
            matrix_row=matrix[idx : idx + 1],
            activation_score=float(activation_scores[idx]),
            activation_details=activation_details,
            warn_threshold=warn,
            block_threshold=block,
            top_sae_features=top_sae_features,
            elapsed_seconds=elapsed_seconds,
        )
        for idx, text in enumerate(target_texts)
    ]


def scan_text(
    target_text: str,
    artifact: Path,
    *,
    batch_size: int | None = None,
    max_length: int | None = None,
    device: str | None = None,
    dtype: str | None = None,
    local_files_only: bool = False,
    warn_threshold: float | None = None,
    block_threshold: float | None = None,
    top_sae_features: int = 8,
) -> dict[str, Any]:
    return scan_texts(
        [target_text],
        artifact,
        batch_size=batch_size,
        max_length=max_length,
        device=device,
        dtype=dtype,
        local_files_only=local_files_only,
        warn_threshold=warn_threshold,
        block_threshold=block_threshold,
        top_sae_features=top_sae_features,
    )[0]


def main(argv: list[str] | None = None) -> int:
    parser = build_parser()
    args = parser.parse_args(argv)
    if args.batch_file:
        if args.text or args.file:
            raise SystemExit("Use --batch-file by itself, not with --text or --file.")
        rows = read_batch_inputs(args.batch_file)
        outputs = scan_texts(
            [text for _, text in rows],
            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,
        )
        for idx, (input_id, _) in enumerate(rows):
            if input_id is not None:
                outputs[idx]["input_id"] = input_id
        output = {
            "mode": "cached_probe_batch_runtime",
            "scanner_version": CORE_VERSION,
            "artifact": str(args.artifact),
            "count": len(outputs),
            "results": outputs,
        }
    else:
        target_text = read_input_text(args)
        output = scan_text(
            target_text,
            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,
        )
    print(json.dumps(output, indent=2 if args.pretty else None, ensure_ascii=False))
    return 0


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