Open_Ultra_SAST/benchmarks/valid/prover_arm.py

"""Run deterministic access-control proofs on local frozen changes and score the arm."""

from __future__ import annotations

import argparse
import hashlib
import json
import statistics
import tempfile
import time
from collections import Counter
from dataclasses import asdict
from pathlib import Path

from benchmarks.unseen import score
from benchmarks.unseen.extract import SHA, InstrumentFailure
from benchmarks.unseen.replay import canonical, load_slice
from benchmarks.valid import metrics
from benchmarks.valid.artifact_contract import (
    read_benchmark_artifact,
    read_benchmark_sequence,
    write_benchmark_artifact,
    write_benchmark_sequence,
)
from benchmarks.valid.catalog import local_git
from openultrasast.hunt.producers import Context, _Absent, _safe_path
from openultrasast.hunt.prover import prove_change


def _write_json(path, value, *, kind=None, replace=False):
    kind = kind or ("benchmark-arm-inputs" if Path(path).name == "inputs.json" else "benchmark-arm-record")
    if isinstance(value, list):
        write_benchmark_sequence(path, kind, value, replace=replace)
    else:
        write_benchmark_artifact(path, kind, value, replace=replace)


def _record_path(out, identity):
    return out / (hashlib.sha256(identity.encode()).hexdigest() + ".json")


def _prove(change, repos, *, max_depth):
    started = time.monotonic()
    instrument = dict(
        hook_exit=0,
        hook_bytes_read=0,
        changed_bytes={},
        file_bytes={},
        head_verified=False,
        base_verified=False,
        suspect_fast=False,
        failure=None,
    )
    record = {
        **change,
        "done": True,
        "status": "prover_covered",
        "instrument": instrument,
        "findings": [],
        "traces": [],
        "coverage": {},
        "degradations": [],
        "push": {"admission": {"defects": [], "dispositions": []}, "quick_tier": []},
    }
    try:
        if not all(SHA.fullmatch(change[side]) for side in ("base", "head")):
            raise InstrumentFailure("invalid revision")
        git = local_git(Path(repos[change["repo"]]).expanduser().resolve())
        git.run("cat-file", "-e", change["base"] + "^{commit}")
        instrument["base_verified"] = True
        with tempfile.TemporaryDirectory(prefix="ousast-prover-") as temporary:
            root = Path(temporary) / "head"
            added = False
            try:
                git.run("-c", "core.hooksPath=/dev/null", "worktree", "add", "--detach", str(root), change["head"])
                added = True
                checked_git = local_git(root)
                # Use the local-only Git runner even when expanding an inherited sparse checkout.
                if (
                    checked_git.runner(
                        ["git", "-C", str(root), "config", "--get", "core.sparseCheckout"], capture_output=True, text=True
                    ).stdout.strip()
                    == "true"
                ):
                    checked_git.run("sparse-checkout", "disable")
                if checked_git.run("status", "--porcelain", "--untracked-files=no").strip():
                    raise InstrumentFailure("incomplete checkout")
                ctx = Context(root, change)
                ctx.verify_head()
                instrument["head_verified"] = True
                changed = {file["path"] for file in change.get("files", [])}
                # Verify every source read, including hop files, before trusting a quiet verdict.
                for relative in sorted(changed | set(ctx.product_files)):
                    path = _safe_path(root, relative)
                    if not path.is_file():
                        held = checked_git.runner(["git", "-C", str(root), "cat-file", "-e", f"HEAD:{relative}"], capture_output=True)
                        if held.returncode == 0:
                            raise InstrumentFailure("unread tracked file")
                        continue  # a file deleted by this change has no head input
                    try:
                        text = ctx.read(relative)
                    except _Absent as error:
                        # A file git holds empty is read, not unread: the first smoke run failed four changes on
                        # legitimately empty files (`__init__.py`, an empty asset). Anything else absent is unread.
                        held = checked_git.runner(
                            ["git", "-C", str(root), "cat-file", "-s", f"HEAD:{relative}"], capture_output=True, text=True
                        )
                        if error.reason == "empty_source" and held.returncode == 0 and int(held.stdout.strip() or 0) == 0:
                            instrument["empty_files"] = instrument.get("empty_files", 0) + 1
                            continue
                        raise InstrumentFailure("unread tracked file") from error
                    # The byte table holds only files that were read; a zero here is a real zero-byte read, which
                    # stays on the record so the summary refuses it.
                    size = len(text.encode("utf-8"))
                    instrument["file_bytes"][relative] = size
                    if relative in changed:
                        instrument["changed_bytes"][relative] = size
                    print(f"{change['id']} {relative}: {size} bytes read", flush=True)
                    if not size:
                        raise InstrumentFailure("zero-byte read")
                result = prove_change(root, change, ctx, max_depth=max_depth)
                instrument["hook_bytes_read"] = sum(instrument["file_bytes"].values())
                record.update(
                    findings=list(result.findings),
                    traces=list(result.traces),
                    coverage=asdict(result.coverage),
                    degradations=list(result.degradations),
                )
                record["push"]["admission"]["defects"] = [
                    {**finding, "defect_id": finding["id"], "source": "prover", "locations": [[finding["path"], finding["line"]]]}
                    for finding in result.findings
                ]
            finally:
                if added:
                    git.run("worktree", "remove", "--force", str(root))
    except Exception as error:
        instrument.update(hook_exit=2, failure=type(error).__name__)
        record.update(status="instrument_failure", findings=[])
        record["push"]["admission"]["defects"] = []
        print(f"{change['id']}: instrument_failure ({type(error).__name__})", flush=True)
    record["seconds"] = time.monotonic() - started
    return record


def run(manifest, repos_path, out, *, slice_id="1", only="all", max_depth=3):
    if max_depth < 0 or only not in {"all", "ordinary", "vulnerability"}:
        raise ValueError("invalid prover arguments")
    manifest, repos_path, out = Path(manifest), Path(repos_path), Path(out)
    changes = load_slice(manifest, slice_id)
    if len({change["id"] for change in changes}) != len(changes):
        raise ValueError("duplicate change ids")
    changes = [change for change in changes if only == "all" or change["kind"] == only]
    index = out / "inputs.json"
    entries = list(read_benchmark_sequence(index, "benchmark-arm-inputs")) if index.exists() else []
    if len({entry["item_id"] for entry in entries}) != len(entries):
        raise ValueError("duplicate inputs")
    indexed = {entry["item_id"]: entry for entry in entries}
    pending = []
    for change in changes:
        digest = hashlib.sha256(canonical({"change": change, "max_depth": max_depth, "arm": "prover-a"})).hexdigest()
        previous_entry = indexed.get(change["id"])
        if previous_entry and (previous_entry["change"] != change or previous_entry["input_digest"] != digest):
            raise ValueError("existing prover inputs do not match")
        entry = dict(item_id=change["id"], change=change, input_digest=digest)
        if not previous_entry:
            entries.append(entry)
        path = _record_path(out, change["id"])
        previous = read_benchmark_artifact(path, "benchmark-arm-record") if path.exists() else {}
        if previous and previous.get("input_digest") != digest:
            raise ValueError("existing prover inputs do not match")
        if previous.get("done"):
            print(f"{change['id']}: skipped (done)", flush=True)
        else:
            pending.append(entry)
    _write_json(index, entries, kind="benchmark-arm-inputs", replace=True)
    repos = json.loads(repos_path.read_text())
    failed = False
    for entry in pending:
        record = _prove(entry["change"], repos, max_depth=max_depth)
        record["input_digest"] = entry["input_digest"]
        _write_json(_record_path(out, entry["item_id"]), record)
        failed |= record["status"] == "instrument_failure"
    return 2 if failed else 0


def _summary_records(out):
    records = score.load_records(Path(out))
    if not records or len({row["id"] for row in records}) != len(records):
        raise ValueError("missing or duplicate prover records")
    for record in records:
        instrument = record.get("instrument", {})
        if any(size == 0 for field in ("changed_bytes", "file_bytes") for size in instrument.get(field, {}).values()):
            raise ValueError("zero-byte read of a tracked file; summary refused")
    return records


def _target_locations(usable):
    return [
        bool(score.catches({**row, **target, "files": [{"path": target["path"], "head_lines": [target["function_lines"]]}]}))
        for row in usable
        if row["kind"] == "vulnerability"
        for target in row.get("targets") or [row]
    ]


def _coverage_slot():
    fields = ("entry_points", "reached", "unreached", "operations_checked", "operations_discharged")
    return {
        **dict.fromkeys(fields, 0),
        "files_by_language": Counter(),
        "frameworks_recognised": Counter(),
        "unchecked_by_reason": Counter(),
        "changes": 0,
    }


def _add_coverage(target, coverage, fields):
    target["changes"] += 1
    for field in fields:
        target[field] += coverage.get(field, 0)
    for field in ("files_by_language", "frameworks_recognised", "unchecked_by_reason"):
        if field in target:
            target[field].update(coverage.get(field, {}))


def _coverage_tables(records):
    fields = ("entry_points", "reached", "unreached", "operations_checked", "operations_discharged")
    repositories, framework_rows = {}, {}
    for row in records:
        coverage = row.get("coverage", {})
        _add_coverage(repositories.setdefault(row["repo"], _coverage_slot()), coverage, fields)
        for framework, count in coverage.get("frameworks_recognised", {}).items():
            table = framework_rows.setdefault(
                framework, {"files": 0, "changes": 0, **dict.fromkeys(fields, 0), "unchecked_by_reason": Counter()}
            )
            table["files"] += count
            _add_coverage(table, coverage, fields)
    return repositories, framework_rows


def _annotate_repositories(repositories, usable):
    uncovered = sorted(repo for repo, coverage in repositories.items() if not coverage["frameworks_recognised"])
    for repo, coverage in repositories.items():
        coverage["status"] = "uncovered" if repo in uncovered else "covered"
        positives = [row for row in usable if row["repo"] == repo and row["kind"] == "vulnerability"]
        coverage["recall"] = score.rate([bool(score.catches(row)) for row in metrics.scoring_records(positives)])
    return uncovered


def _scope_recall(usable, uncovered, *, include_uncovered):
    values = [
        bool(score.catches(row))
        for row in metrics.scoring_records(usable)
        if row["kind"] == "vulnerability" and ((row["repo"] in uncovered) is include_uncovered)
    ]
    return score.rate(values)


def summarize(out, *, seed=0):
    out = Path(out)
    records = _summary_records(out)
    result = metrics.summarize(records, seed=seed)
    usable = [row for row in records if score.coverage(row) != "instrument_failure"]
    repositories, framework_rows = _coverage_tables(records)
    uncovered = _annotate_repositories(repositories, usable)
    seconds = [row["seconds"] for row in records]
    result.update(
        location_recall=score.rate(_target_locations(usable)),
        per_repository=repositories,
        per_framework=framework_rows,
        uncovered_repositories=uncovered,
        covered_recall=_scope_recall(usable, uncovered, include_uncovered=False),
        uncovered_recall=_scope_recall(usable, uncovered, include_uncovered=True),
        instrument_failures=sum(score.coverage(row) == "instrument_failure" for row in records),
        seconds_per_change=dict(n=len(seconds), mean=statistics.mean(seconds), median=statistics.median(seconds), maximum=max(seconds)),
        changes_over_120_seconds=[{"id": row["id"], "seconds": row["seconds"]} for row in records if row["seconds"] > 120],
    )
    # This is a training pool; the generic metric headline refers to unseen repositories.
    result["headline"] = result["headline"].replace("On unseen app repos:", "On the training pool:")
    _write_json(out / "summary.json", result, kind="benchmark-arm-summary", replace=True)
    return result


def main(argv=None):
    parser = argparse.ArgumentParser(description=__doc__)
    commands = parser.add_subparsers(dest="command", required=True)
    runner = commands.add_parser("run")
    runner.add_argument("pool", type=Path)
    runner.add_argument("--repos", type=Path, required=True)
    runner.add_argument("--out", type=Path, required=True)
    runner.add_argument("--slice", default="1")
    runner.add_argument("--only", choices=("all", "ordinary", "vulnerability"), default="all")
    runner.add_argument("--max-depth", type=int, default=3)
    summarizer = commands.add_parser("summarize")
    summarizer.add_argument("out", type=Path)
    args = parser.parse_args(argv)
    if args.command == "run":
        status = run(args.pool, args.repos, args.out, slice_id=args.slice, only=args.only, max_depth=args.max_depth)
        if status:
            return status
    result = summarize(args.out)
    print(result["headline"])
    return 0


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