#!/usr/bin/env python3
"""CASAN Plan-16 SEC-18 (ARCH-02) — test-integrity manifest.

The test suites live in the same repo an attacker can edit, so "280/0 green" proves
nothing if a fail-able adversarial check was quietly deleted or weakened. This tool
records, per suite, its content hash and its number of fail-able checks, into a
manifest whose head is signed with an OFF-REPO key. CI then re-verifies:

  * a manifested suite that DISAPPEARED  -> FAIL (suite removed)
  * a suite whose fail-able-check count DROPPED -> FAIL (coverage regression)
  * (enforced mode) an invalid/absent manifest signature -> FAIL (manifest tampered)

Adding checks/suites is fine (regenerate the manifest); only REMOVING coverage fails.

Usage:
  test-integrity.py generate            # write + sign the manifest
  test-integrity.py verify [--strict]   # re-check; exit 1 on regression/tamper

Env: CASAN_TESTS_DIR, CASAN_TEST_MANIFEST, CASAN_TI_KEY_DIR, CASAN_TI_PUB,
     CASAN_PROFILE=prod / CASAN_VERIFY_STRICT=1 (enforce signature).
"""
import argparse
import glob
import hashlib
import json
import os
import re
import shutil
import subprocess
import sys

# A "fail-able check" is any assertion call site. Phase suites use `pass "..."`;
# the adversarial suite uses `expect_rc <n> "..."`. Counting these makes deleting
# or short-circuiting a check reduce the number.
CHECK_RE = re.compile(r'(?:(?<![\w])pass\s+")|(?:\bexpect_rc\s)')


def project_root() -> str:
    d = os.path.abspath(os.path.dirname(__file__))
    p = d
    while p != os.path.dirname(p):
        if os.path.isdir(os.path.join(p, ".specify")) or os.path.isdir(os.path.join(p, "packages/casan-harness")):
            return p
        p = os.path.dirname(p)
    return os.path.abspath(os.path.join(d, "..", "..", ".."))


def tests_dir() -> str:
    return os.environ.get("CASAN_TESTS_DIR") or os.path.join(project_root(), ".specify/tests")


def manifest_path() -> str:
    return os.environ.get("CASAN_TEST_MANIFEST") or os.path.join(tests_dir(), "test-integrity-manifest.json")


def _enforced() -> bool:
    return os.environ.get("CASAN_PROFILE") == "prod" or os.environ.get("CASAN_VERIFY_STRICT") == "1"


def suite_files():
    d = tests_dir()
    files = []
    for pat in ("phase*.sh", "adversarial*.sh", "run-casan*.sh"):
        files.extend(glob.glob(os.path.join(d, pat)))
    # Exclude the integrity suite itself to avoid a self-reference cycle.
    return sorted(f for f in files if os.path.basename(f) != "phase-sec18-tests.sh")


def scan(path):
    data = open(path, "rb").read()
    text = data.decode("utf-8", errors="replace")
    checks = len(CHECK_RE.findall(text))
    return {"sha256": hashlib.sha256(data).hexdigest(), "checks": checks}


def build_manifest():
    suites = {os.path.basename(f): scan(f) for f in suite_files()}
    total = sum(s["checks"] for s in suites.values())
    return {"suites": suites, "total_checks": total, "suite_count": len(suites)}


def manifest_head(manifest) -> str:
    core = json.dumps(manifest["suites"], sort_keys=True, separators=(",", ":"))
    return hashlib.sha256(core.encode()).hexdigest()


# --- signing (off-repo key; pubkey provisioned out-of-band / KMS in prod) ------
def _priv():
    d = os.environ.get("CASAN_TI_KEY_DIR") or os.path.join(os.path.expanduser("~"), ".casan", "audit-keys")
    return os.path.join(d, "test-integrity-private.pem")


def _pub():
    return os.environ.get("CASAN_TI_PUB") or (manifest_path() + ".pub")


def _sig():
    return manifest_path() + ".sig"


def _head_file():
    return manifest_path() + ".head"


def sign(manifest):
    head = manifest_head(manifest)
    open(_head_file(), "w", encoding="utf-8").write(head)
    ossl = shutil.which("openssl")
    if not ossl:
        return
    priv, pub = _priv(), _pub()
    os.makedirs(os.path.dirname(priv), exist_ok=True)
    if not os.path.isfile(priv):
        subprocess.run([ossl, "genpkey", "-algorithm", "RSA", "-pkeyopt", "rsa_keygen_bits:2048", "-out", priv],
                       capture_output=True)
        try:
            os.chmod(priv, 0o600)
        except OSError:
            pass
    os.makedirs(os.path.dirname(pub) or ".", exist_ok=True)
    subprocess.run([ossl, "rsa", "-in", priv, "-pubout", "-out", pub], capture_output=True)
    subprocess.run([ossl, "dgst", "-sha256", "-sign", priv, "-out", _sig(), _head_file()], capture_output=True)


def check_signature(manifest):
    ossl = shutil.which("openssl")
    if not (ossl and os.path.isfile(_head_file()) and os.path.isfile(_sig()) and os.path.isfile(_pub())):
        return "unsigned"
    if open(_head_file(), encoding="utf-8").read().strip() != manifest_head(manifest):
        return "invalid"
    res = subprocess.run([ossl, "dgst", "-sha256", "-verify", _pub(), "-signature", _sig(), _head_file()],
                         capture_output=True)
    return "signed" if res.returncode == 0 else "invalid"


def do_generate():
    manifest = build_manifest()
    with open(manifest_path(), "w", encoding="utf-8") as fh:
        json.dump(manifest, fh, indent=2)
        fh.write("\n")
    sign(manifest)
    print(f"TEST_INTEGRITY_GENERATED suites={manifest['suite_count']} total_checks={manifest['total_checks']}")
    return 0


def do_verify(strict):
    mp = manifest_path()
    if not os.path.isfile(mp):
        # No manifest provisioned. Enforced mode fails closed; dev/CI skips cleanly.
        if strict or _enforced():
            print("TEST_INTEGRITY_FAIL no_manifest_in_enforced_mode", file=sys.stderr)
            return 1
        print("TEST_INTEGRITY_SKIP no manifest (run: test-integrity.py generate)")
        return 0
    manifest = json.load(open(mp, encoding="utf-8"))
    current = {os.path.basename(f): scan(f) for f in suite_files()}

    regressions = []
    for name, rec in manifest.get("suites", {}).items():
        if name not in current:
            regressions.append(f"{name}:suite_removed")
        elif current[name]["checks"] < rec["checks"]:
            regressions.append(f"{name}:checks_dropped({rec['checks']}->{current[name]['checks']})")

    sig_state = check_signature(manifest)
    if sig_state == "invalid":
        print("TEST_INTEGRITY_FAIL manifest_signature_invalid", file=sys.stderr)
        return 1
    if sig_state == "unsigned" and (strict or _enforced()):
        print("TEST_INTEGRITY_FAIL manifest_unsigned_in_strict_mode", file=sys.stderr)
        return 1

    if regressions:
        print("TEST_INTEGRITY_FAIL coverage_regression " + " ".join(regressions), file=sys.stderr)
        return 1

    cur_total = sum(v["checks"] for k, v in current.items() if k in manifest.get("suites", {}))
    print(f"TEST_INTEGRITY_OK suites={len(manifest.get('suites', {}))} "
          f"manifest_checks={manifest.get('total_checks')} current_checks={cur_total} anchor={sig_state}")
    return 0


def main() -> int:
    ap = argparse.ArgumentParser()
    sub = ap.add_subparsers(dest="cmd", required=True)
    sub.add_parser("generate")
    v = sub.add_parser("verify")
    v.add_argument("--strict", action="store_true")
    args = ap.parse_args()
    if args.cmd == "generate":
        return do_generate()
    if args.cmd == "verify":
        return do_verify(args.strict)
    return 2


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