#!/usr/bin/env python3
"""Regression checks for the Python Pwned v2 jail.

Run locally with ``python3 regression_harness.py`` or against the challenge
image with ``python3 regression_harness.py --docker-image IMAGE``.
"""

from __future__ import annotations

import argparse
import ast
from dataclasses import dataclass
from pathlib import Path
import subprocess
import sys

import solve_v2


ROOT = Path(__file__).resolve().parent
RECOVER_BUILTINS = (
    "[x for x in ().__class__.__base__.__subclasses__() "
    "if x.__name__=='BuiltinImporter'][0].load_module('builtins')"
)


@dataclass(frozen=True)
class Case:
    name: str
    payload: bytes
    should_succeed: bool
    should_leak_flag: bool = False


def official_payload(path: Path) -> bytes:
    """Read the initial bytes literal from solver/solver.py without running it."""
    module = ast.parse(path.read_text(), filename=str(path))
    for statement in module.body:
        if not isinstance(statement, ast.Assign):
            continue
        if not any(isinstance(target, ast.Name) and target.id == "payload"
                   for target in statement.targets):
            continue
        if isinstance(statement.value, ast.Constant) and isinstance(statement.value.value, bytes):
            return statement.value.value.replace(b"\n", b"\r")
    raise RuntimeError(f"initial payload bytes literal not found in {path}")


def run_payload(command: list[str], payload: bytes, timeout: float) -> subprocess.CompletedProcess[bytes]:
    return subprocess.run(
        command,
        input=payload + b"\n",
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
        timeout=timeout,
        check=False,
    )


def display(data: bytes, limit: int = 260) -> str:
    if len(data) > limit:
        data = data[:limit] + b"...<truncated>"
    return repr(data.decode("utf-8", "backslashreplace"))


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--app", type=Path, default=ROOT / "app.py")
    parser.add_argument("--solver", type=Path, default=ROOT / "solver" / "solver.py")
    parser.add_argument("--flag-file", type=Path, default=ROOT / "flag.txt")
    parser.add_argument("--docker-image", help="run each case in this challenge image")
    parser.add_argument("--timeout", type=float, default=15.0)
    args = parser.parse_args()

    expected_flag = args.flag_file.read_bytes().strip()
    if args.docker_image:
        target = "/flag.txt"
        command = [
            "docker", "run", "--rm", "-i", args.docker_image,
            "python3", "/opt/chall/app.py",
        ]
    else:
        target = str(args.flag_file.resolve())
        command = [sys.executable, str(args.app.resolve())]

    quitter = (
        f"b={RECOVER_BUILTINS};"
        "b.exit.__class__.__call__=lambda *a:None;"
        f"b.print(b.open({target!r}).read())"
    ).encode()
    direct_open = (
        f"b={RECOVER_BUILTINS};b.print(b.open({target!r}).read())"
    ).encode()
    direct_system = (
        f"b={RECOVER_BUILTINS};"
        f"b.__import__('os').system({'cat ' + target!r})"
    ).encode()
    extension_loader = solve_v2.build_payload(target).encode()

    cases = [
        Case("safe_assignment", b"answer=6*7", should_succeed=True),
        Case("old_quitter_mutation", quitter, should_succeed=False),
        Case("solve_v2_extension_loader", extension_loader, should_succeed=False),
        Case("direct_open", direct_open, should_succeed=False),
        Case("direct_os_system", direct_system, should_succeed=False),
        Case(
            "official_memory_corruption_solver",
            official_payload(args.solver),
            should_succeed=True,
            should_leak_flag=True,
        ),
    ]

    print("command:", " ".join(command))
    print("target:", target)
    failures = 0
    for case in cases:
        try:
            result = run_payload(command, case.payload, args.timeout)
        except subprocess.TimeoutExpired as exc:
            failures += 1
            print(f"FAIL {case.name}: TIMEOUT after {args.timeout}s")
            if exc.stdout:
                print("  stdout:", display(exc.stdout))
            if exc.stderr:
                print("  stderr:", display(exc.stderr))
            continue

        combined = result.stdout + result.stderr
        leaked = expected_flag in combined
        passed = (
            (result.returncode == 0) == case.should_succeed
            and leaked == case.should_leak_flag
        )
        if not passed:
            failures += 1
        print(
            f"{'PASS' if passed else 'FAIL'} {case.name}: "
            f"rc={result.returncode} flag={leaked}"
        )
        print("  stdout:", display(result.stdout))
        if result.stderr:
            print("  stderr:", display(result.stderr))

    print(f"summary: {len(cases) - failures}/{len(cases)} passed")
    return 1 if failures else 0


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