#!/usr/bin/env python3
"""RED-first adversarial Claude B1 gate contracts.

This private harness drives the shipped gate and daemon out of process.  It
never contacts a live broker/runtime and executes a shell payload only after a
regression has already (incorrectly) received the recovery exemption.
"""

from __future__ import annotations

import json
import os
import re
import socket
import subprocess
import sys
import tempfile
import time
import unittest
from pathlib import Path


TOOLS = Path(__file__).parents[2] / "framework/tools/lease-broker"
FRAMEWORK = Path(__file__).parents[2] / "framework"
DAEMON = TOOLS / "daemon.py"
GATE = TOOLS / "mutator-gate.py"
RECOVERY = TOOLS / "recover-context.py"
SKILL = FRAMEWORK / "skills/mosaic-context-refresh/SKILL.md"


def request(socket_path: Path, value: dict[str, object]) -> dict[str, object]:
    with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as connection:
        connection.settimeout(3.0)
        connection.connect(str(socket_path))
        connection.sendall((json.dumps(value, separators=(",", ":")) + "\n").encode())
        connection.shutdown(socket.SHUT_WR)
        response = bytearray()
        while True:
            chunk = connection.recv(4096)
            if not chunk:
                break
            response.extend(chunk)
    if not response.endswith(b"\n") or response.count(b"\n") != 1:
        raise AssertionError(f"unframed broker response: {bytes(response)!r}")
    reply = json.loads(response[:-1])
    if not isinstance(reply, dict):
        raise AssertionError("broker response is not an object")
    return reply


def wait_ready(process: subprocess.Popen[str], socket_path: Path) -> None:
    deadline = time.monotonic() + 5.0
    while time.monotonic() < deadline:
        try:
            with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as probe:
                probe.settimeout(0.25)
                probe.connect(str(socket_path))
            return
        except (ConnectionRefusedError, FileNotFoundError):
            pass
        if process.poll() is not None:
            output = process.stdout.read() if process.stdout is not None else ""
            raise RuntimeError(f"daemon exited before READY: {output}")
        time.sleep(0.02)
    raise TimeoutError("daemon did not create private socket")


class ClaudeRecoveryGateAdversarialTest(unittest.TestCase):
    def setUp(self) -> None:
        self.temporary = tempfile.TemporaryDirectory()
        self.root = Path(self.temporary.name)
        os.chmod(self.root, 0o700)
        self.socket = self.root / "broker.sock"
        self.daemon = subprocess.Popen(
            [sys.executable, "-I", "-S", "-B", str(DAEMON), "--socket", str(self.socket), "--state", str(self.root / "state.json")],
            stdin=subprocess.DEVNULL,
            stdout=subprocess.PIPE,
            stderr=subprocess.STDOUT,
            text=True,
        )
        wait_ready(self.daemon, self.socket)
        registered = request(self.socket, {"action": "register_anchor", "runtime_generation": 1})
        self.session_id = registered["session_id"]
        self.environment = {
            **os.environ,
            "MOSAIC_LEASE_BROKER_SOCKET": str(self.socket),
            "MOSAIC_LEASE_SESSION_ID": self.session_id,
            "MOSAIC_RUNTIME_GENERATION": "1",
            "MOSAIC_LEASE_RUNTIME": "claude",
        }
        # This private path is only a gate-classifier identity; if a regression
        # blesses a payload, bash invokes no installed/live recovery command.
        self.recovery_path = self.root / "recover-context.py"
        self.canonical = (
            f"python3 {self.recovery_path} begin "
            f"--construction {self.root / 'construction.json'} --compaction-epoch 0 --request-epoch 0"
        )

    def tearDown(self) -> None:
        if self.daemon.poll() is None:
            self.daemon.terminate()
            try:
                self.daemon.wait(timeout=3.0)
            except subprocess.TimeoutExpired:
                self.daemon.kill()
                self.daemon.wait()
        if self.daemon.stdout is not None:
            self.daemon.stdout.close()
        self.temporary.cleanup()

    def gate(
        self, command: str, recovery_command: Path | str | None = None
    ) -> subprocess.CompletedProcess[str]:
        configured_recovery = self.recovery_path if recovery_command is None else recovery_command
        return subprocess.run(
            [
                sys.executable,
                "-I",
                "-S",
                "-B",
                str(GATE),
                "--runtime",
                "claude",
                "--recovery-command",
                str(configured_recovery),
            ],
            input=json.dumps({"tool_name": "Bash", "tool_input": {"command": command}}),
            text=True,
            capture_output=True,
            env=self.environment,
            check=False,
        )

    def command_for(self, position: str, value: str) -> str:
        argv = [
            "python3",
            str(self.recovery_path),
            "begin",
            "--construction",
            str(self.root / "construction.json"),
            "--compaction-epoch",
            "0",
            "--request-epoch",
            "0",
        ]
        positions = {
            "executable": 0,
            "path": 1,
            "phase": 2,
            "construction_flag": 3,
            "construction": 4,
            "compaction_epoch_flag": 5,
            "compaction_epoch": 6,
            "request_epoch_flag": 7,
            "request_epoch": 8,
        }
        try:
            argv[positions[position]] = value
        except KeyError as exc:
            raise AssertionError(f"unknown argv position {position}") from exc
        return " ".join(argv)

    def test_canonical_literal_and_shipped_skill_invocation_are_ungated(self) -> None:
        self.assertEqual(self.gate(self.canonical).returncode, 0)
        source = SKILL.read_text(encoding="utf-8")
        recovery_placeholder = "/absolute/path/to/mosaic/tools/lease-broker/recover-context.py"
        construction_placeholder = "/absolute/path/to/mosaic-context-refresh-construction.json"
        self.assertIn(recovery_placeholder, source)
        self.assertIn(construction_placeholder, source)
        resolved_recovery = "/opt/mosaic/tools/lease-broker/recover-context.py"
        rendered = source.replace(recovery_placeholder, resolved_recovery).replace(
            construction_placeholder, "/opt/mosaic/recovery/construction.json"
        )
        match = re.search(r"```bash\s*\n\s*(.*?)\n\s*```", rendered, flags=re.DOTALL)
        self.assertIsNotNone(match)
        shipped = match.group(1) if match is not None else ""
        self.assertEqual(self.gate(shipped, resolved_recovery).returncode, 0)

    def test_every_shell_active_vector_in_every_argv_position_falls_through_to_bash_deny(self) -> None:
        positions = (
            "executable",
            "path",
            "phase",
            "construction_flag",
            "construction",
            "compaction_epoch_flag",
            "compaction_epoch",
            "request_epoch_flag",
            "request_epoch",
        )
        marker = self.root / "PWNED"
        marker_vector = f"$(touch${{IFS}}{marker})"
        vectors = {
            "command-substitution": marker_vector,
            "backtick": f"`touch${{IFS}}{marker}`",
            "process-substitution": f"<(touch${{IFS}}{marker})",
            "parameter-expansion": "${IFS}",
            "home-expansion": "${HOME}",
            "arithmetic-expansion": "$((1+1))",
            "brace-expansion": "{a,b}",
            "tilde-expansion": "~",
            "glob": "*",
            "redirection": ">",
            "semicolon": f";touch${{IFS}}{marker}",
            "and": f"&&touch${{IFS}}{marker}",
            "pipe": f"|touch${{IFS}}{marker}",
            "embedded-newline": "literal\nnext",
            "quoting-trick": "'literal'",
        }
        for position in positions:
            for kind, vector in vectors.items():
                with self.subTest(position=position, kind=kind):
                    marker.unlink(missing_ok=True)
                    command = self.command_for(position, vector)
                    gated = self.gate(command)
                    if gated.returncode == 0 and kind in {"command-substitution", "backtick", "process-substitution", "semicolon", "and", "pipe"}:
                        subprocess.run(["bash", "-c", command], cwd=self.root, env=self.environment, check=False)
                    self.assertEqual(gated.returncode, 2, f"unexpected recovery exemption: {command!r}")
                    self.assertFalse(marker.exists(), f"shell payload executed: {command!r}")


if __name__ == "__main__":
    unittest.main()
