# pyright: reportAttributeAccessIssue=false, reportIncompatibleMethodOverride=false, reportArgumentType=false
#!/usr/bin/env python3
"""Four backend hardening fixes, each with a check that fails without it.

- llm.complete does not follow a redirect (SSRF guard only vets the first URL);
- an oversized request body gets 413 and closes the connection (no desync);
- two concurrent files.write calls never share a temp file;
- a failed create_session removes its worktree and kills its tmux session.

No network beyond loopback, no tmux, no browser.

    python3 tools/hardening_check.py

Exit status is 0 on pass, 1 on fail.
"""

from __future__ import annotations

import http.server
import shutil
import sys
import tempfile
import threading
from pathlib import Path
from types import SimpleNamespace

ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT))
from clique import app, attention, files, gitinfo, llm, store, tmux, usage


def check_llm_redirect() -> bool:
    hits: list[str] = []
    raised = False

    class H(http.server.BaseHTTPRequestHandler):
        def do_POST(self):
            hits.append(self.path)
            if self.path == "/start":
                self.send_response(302)
                self.send_header("Location", "/internal")
                self.end_headers()
            else:
                self.send_response(200)
                self.end_headers()
                self.wfile.write(b"{}")

        do_GET = do_POST

        def log_message(self, *a):
            pass

    srv = http.server.HTTPServer(("127.0.0.1", 0), H)
    threading.Thread(target=srv.serve_forever, daemon=True).start()
    guard, llm._guard_url = llm._guard_url, lambda url: None  # loopback is the test fixture
    profile = {"kind": "openai", "base_url": f"http://127.0.0.1:{srv.server_port}/start", "model": "m", "key": "k"}
    orig = llm._request_for
    llm._request_for = lambda p, m, t: (profile["base_url"], {}, {})
    try:
        llm.complete(profile, [{"role": "user", "content": "x"}])
    except llm.LLMError:
        raised = True
    finally:
        llm._guard_url, llm._request_for = guard, orig
        srv.shutdown()
    return raised and hits == ["/start"]


def check_body_413() -> bool:
    """Oversize POST and PATCH against a real Handler over a socket: one 413,
    then the connection is closed, and the handler never carries on."""
    import socket

    errors: list[object] = []

    class S(app.Server):
        def handle_error(self, request, client_address):
            errors.append(client_address)

    saved = (app.Handler.panel if hasattr(app.Handler, "panel") else None, app.Handler._host_ok, app.Handler._may_write)
    app.Handler.panel = SimpleNamespace()
    app.Handler._host_ok = lambda self: True
    app.Handler._may_write = lambda self: (True, "")
    srv = S(("127.0.0.1", 0), app.Handler)
    threading.Thread(target=srv.serve_forever, daemon=True).start()
    good = True
    try:
        for method, path in (("POST", "/api/sessions"), ("PATCH", "/api/settings")):
            c = socket.create_connection(("127.0.0.1", srv.server_address[1]), timeout=5)
            c.sendall(f"{method} {path} HTTP/1.1\r\nHost: 127.0.0.1\r\n"
                      f"Content-Length: {app.MAX_BODY_BYTES + 1}\r\n\r\n".encode())
            data = b""
            while True:  # EOF means the server closed; a timeout means it did not
                chunk = c.recv(65536)
                if not chunk:
                    break
                data += chunk
            c.close()
            good &= b" 413 " in data.split(b"\r\n")[0]
            good &= data.count(b"HTTP/1.") == 1
    finally:
        srv.shutdown()
        app.Handler.panel, app.Handler._host_ok, app.Handler._may_write = saved
    return good and not errors


def check_write_temp_unique() -> bool:
    d = Path(tempfile.mkdtemp())
    (d / "f.txt").write_text("a")
    seen: list[str] = []
    real = files.os.replace
    files.os.replace = lambda src, dst: (seen.append(str(src)), real(src, dst))[1]
    try:
        files.write(str(d), "f.txt", "one")
        # a stale temp from a "concurrent" writer must not be reused
        files.write(str(d), "f.txt", "two")
    finally:
        files.os.replace = real
    return len(set(seen)) == 2 and (d / "f.txt").read_text() == "two" and not list(d.glob("*.clique-save"))


def _rollback_run(fail: str) -> list[str]:
    log: list[str] = []
    saved = (gitinfo.repo_root, gitinfo.add_worktree, gitinfo.worktree_setup, gitinfo.remove_worktree,
             tmux.bootstrap, tmux.create, tmux.kill)
    wt = tempfile.mkdtemp()

    def boom(*a, **k):
        raise OSError(fail)

    def make_tmux(*a, **k):
        log.append("tmux")
        if fail == "tmux":  # the session exists, then create fails
            boom()

    gitinfo.repo_root = lambda c: c
    gitinfo.add_worktree = lambda _r, _b: (wt, "")
    gitinfo.worktree_setup = boom if fail == "setup" else (lambda r, m: None)
    gitinfo.remove_worktree = lambda p: log.append("rmworktree")
    tmux.bootstrap = lambda: None
    tmux.create = make_tmux
    tmux.kill = lambda m, *a, **k: log.append("kill")
    p = object.__new__(app.Panel)
    cli = SimpleNamespace(installed=True, default_mode=None, hooks=False, label="x", command="x")
    p.registry = SimpleNamespace(launch_argv=lambda *a, **k: ["true"], get=lambda c: cli)
    p._pane_env = lambda s: {}
    p._hooks_argv = lambda c: []
    p.store = SimpleNamespace(add_session=boom if fail == "store" else None)
    try:
        p.create_session({"cwd": wt, "worktree": True, "branch": "b"})
    except OSError:
        pass
    finally:
        (gitinfo.repo_root, gitinfo.add_worktree, gitinfo.worktree_setup, gitinfo.remove_worktree,
         tmux.bootstrap, tmux.create, tmux.kill) = saved
    return log


def check_create_rollback() -> bool:
    setup = _rollback_run("setup")
    tm = _rollback_run("tmux")
    store = _rollback_run("store")
    return ("rmworktree" in setup and "tmux" not in setup
            and "kill" in tm and "rmworktree" in tm
            and "kill" in store and "rmworktree" in store)


def check_write_chmod_failure_cleans_up() -> bool:
    d = Path(tempfile.mkdtemp())
    (d / "f.txt").write_text("a")
    real = files.os.fchmod

    def boom(fd, mode):
        raise OSError("no chmod")

    files.os.fchmod = boom
    try:
        files.write(str(d), "f.txt", "x")
    except OSError:
        pass
    finally:
        files.os.fchmod = real
    return sorted(p.name for p in d.iterdir()) == ["f.txt"]


def check_upload_symlink_swap() -> bool:
    """A drop dir swapped for a symlink between the fence check and the create
    must not let the file land outside cwd."""
    cwd, outside = Path(tempfile.mkdtemp()), Path(tempfile.mkdtemp())
    real_mkdir, fence = Path.mkdir, files._FENCE

    def swap(self, *a, **k):
        if self.name == files.DROP_DIR:
            real_mkdir(self, *a, **k)
            self.rmdir()
            self.symlink_to(outside)
            return
        return real_mkdir(self, *a, **k)

    Path.mkdir, files._FENCE = swap, True
    msg = ""
    try:
        files.store_upload(str(cwd), "a.txt", b"x")
    except ValueError as exc:
        msg = str(exc)
    finally:
        Path.mkdir, files._FENCE = real_mkdir, fence
    return msg == "outside the session directory" and not list(outside.iterdir())


def check_share_cleanup_contained() -> bool:
    """Scratch cleanup touches only flat regular files in a real scratch folder:
    a linked scratch folder (either name), a child symlink, and a folder swapped
    for a link mid-sweep must never lead it to delete outside the project."""
    import os

    base = Path(tempfile.mkdtemp())
    proj, out = base / "project", base / "outside"
    proj.mkdir()
    out.mkdir()
    old = 0  # mtime of the epoch: past any age cutoff

    def keep(name: str) -> Path:
        f = out / name
        f.write_text("keep")
        os.utime(f, (old, old))
        return f

    results = []
    # Both SHARE_DIRS as links to outside: usage, purge and prune all refuse.
    for name in files.SHARE_DIRS:
        f = keep("keep.txt")
        (proj / name).symlink_to(out)
        results.append(files.shares_usage([str(proj)])["files"] == 0)
        results.append(files.prune_shares([str(proj)], 1) == 0 and f.exists())
        results.append(files.purge_shares([str(proj)])["files"] == 0 and f.exists())
        (proj / name).unlink()
    # A normal scratch folder: a child symlink is left alone; age cutoff holds.
    drops = proj / files.DROP_DIR
    drops.mkdir()
    target = keep("target.txt")
    (drops / "link").symlink_to(target)
    (drops / "new.txt").write_text("n")
    stale = drops / "old.txt"
    stale.write_text("o")
    os.utime(stale, (old, old))
    results.append(files.shares_usage([str(proj)])["files"] == 2)
    results.append(
        files.prune_shares([str(proj)], 1) == 1
        and not stale.exists()
        and (drops / "new.txt").exists()
    )
    results.append(
        files.purge_shares([str(proj)])["files"] == 1
        and target.exists()
        and (drops / "link").is_symlink()
    )
    # Swap the folder for a link to outside after it is listed, before deletes.
    (drops / "link").unlink()
    (drops / "victim.txt").write_text("v")
    victim = keep("victim.txt")
    moved = proj / "moved"
    real_scandir = os.scandir

    class Swap:
        """scandir that lists the real folder, then swaps it for a link."""

        def __init__(self, *a):
            self.entries = list(real_scandir(*a))
            os.rename(drops, moved)
            drops.symlink_to(out)

        def __enter__(self):
            return iter(self.entries)

        def __exit__(self, *a):
            return False

    files.os.scandir = Swap
    try:
        files.purge_shares([str(proj)])
    finally:
        files.os.scandir = real_scandir
    results.append(victim.exists() and not (moved / "victim.txt").exists())
    shutil.rmtree(base)
    return all(results)


def check_login_throttle() -> bool:
    """N threads attempt at once (each sleeps in the check, as a real one does
    between check and failure): at most MAX_FAILURES may pass, and a success
    clearing the dict while others iterate must not raise."""
    h = SimpleNamespace(panel=SimpleNamespace(failures={}, failures_lock=threading.Lock()),
                        FAILURE_WINDOW=300, MAX_FAILURES=8)
    h.panel.failures["old"] = [1.0]
    passed: list[bool] = []
    barrier = threading.Barrier(40)

    def attempt():
        barrier.wait()
        passed.append(not app.Handler._reserve_attempt(h, "x"))

    threads = [threading.Thread(target=attempt) for _ in range(40)]
    for t in threads:
        t.start()
    for t in threads:
        t.join()
    with h.panel.failures_lock:
        h.panel.failures.pop("x", None)
    return sum(passed) <= h.MAX_FAILURES and "old" not in h.panel.failures


def check_llm_response_cap() -> bool:
    class H(http.server.BaseHTTPRequestHandler):
        def do_POST(self):
            # Drain the request: closing with it unread makes the kernel send RST,
            # and the client sees "connection reset" instead of the oversized body.
            self.rfile.read(int(self.headers.get("Content-Length", 0)))
            self.send_response(200)
            self.end_headers()
            self.wfile.write(b'{"choices":[{"message":{"content":"' + b"a" * (llm.MAX_RESPONSE_BYTES + 10) + b'"}}]}')

        def log_message(self, *a):
            pass

    srv = http.server.HTTPServer(("127.0.0.1", 0), H)
    threading.Thread(target=srv.serve_forever, daemon=True).start()
    guard, orig = llm._guard_url, llm._request_for
    llm._guard_url = lambda url: None
    llm._request_for = lambda p, m, t: (f"http://127.0.0.1:{srv.server_port}/", {}, {})
    try:
        llm.complete({"kind": "openai"}, [{"role": "user", "content": "x"}])
        return False
    except llm.LLMError as exc:
        return "too large" in str(exc)
    finally:
        llm._guard_url, llm._request_for = guard, orig
        srv.shutdown()


def check_body_must_be_object() -> bool:
    import socket

    class S(app.Server):
        pass

    saved = (getattr(app.Handler, "panel", None), app.Handler._host_ok, app.Handler._may_write)
    app.Handler.panel = SimpleNamespace()
    app.Handler._host_ok = lambda self: True
    app.Handler._may_write = lambda self: (True, "")
    srv = S(("127.0.0.1", 0), app.Handler)
    threading.Thread(target=srv.serve_forever, daemon=True).start()
    good = True
    try:
        for payload in (b"[1]", b"not json", b'"s"'):
            c = socket.create_connection(("127.0.0.1", srv.server_address[1]), timeout=5)
            c.sendall(b"PATCH /api/settings HTTP/1.1\r\nHost: 127.0.0.1\r\nConnection: close\r\n"
                      b"Content-Length: %d\r\n\r\n" % len(payload) + payload)
            data = b""
            while chunk := c.recv(65536):
                data += chunk
            c.close()
            good &= b" 400 " in data.split(b"\r\n")[0]
    finally:
        srv.shutdown()
        app.Handler.panel, app.Handler._host_ok, app.Handler._may_write = saved
    return good


def check_store_bad_shape_uses_backup() -> bool:
    import json

    d = Path(tempfile.mkdtemp())
    (d / "state.json").write_text(json.dumps({"sessions": [1]}))
    (d / "state.json.bak").write_text(json.dumps({"folders": [{"id": "f", "name": "Bak", "color": "#000", "match": []}]}))
    try:
        st = store.Store(d / "state.json")
    except Exception:  # noqa: BLE001
        return False
    return [f.name for f in st.folders] == ["Bak"] and st.sessions == []


def check_store_bad_settings_uses_backup() -> bool:
    import json

    d = Path(tempfile.mkdtemp())
    (d / "state.json").write_text(json.dumps({"settings": []}))
    (d / "state.json.bak").write_text(json.dumps({"folders": [{"id": "f", "name": "Bak", "color": "#000", "match": []}]}))
    try:
        st = store.Store(d / "state.json")
    except Exception:  # noqa: BLE001
        return False
    return [f.name for f in st.folders] == ["Bak"]


def check_auto_resume_bad_timestamp() -> bool:
    mk = lambda i: SimpleNamespace(id=i, folder="f", cli="c", mux=i, socket=None)  # noqa: E731
    cli = SimpleNamespace(usage={"x": 1}, limited_patterns=[], after_limit="go")
    sent: list[str] = []
    panel = SimpleNamespace(
        store=SimpleNamespace(folders=[SimpleNamespace(id="f", auto_resume=True)], sessions=[mk("a"), mk("b")]),
        registry=SimpleNamespace(types=lambda: {"c": cli}),
        session_state=lambda s: "idle", _limit_at={})
    reads = iter([{"windows": [{"percent": 100, "resets_at": "garbage"}]},
                  {"windows": [{"percent": 100, "resets_at": "2999-01-01T00:00:00+00:00"}]}])
    saved = (tmux.capture, attention.limited_text, usage.read, tmux.send_text)
    tmux.capture, attention.limited_text = lambda *a: "", lambda *a: True
    usage.read, tmux.send_text = lambda *a: next(reads), lambda *a: sent.append(a[0])
    try:
        app.Panel.auto_resume_tick(panel, now=1.0)
    finally:
        tmux.capture, attention.limited_text, usage.read, tmux.send_text = saved
    return "b" in panel._limit_at


def check_diff_reads_bounded() -> bool:
    import subprocess

    repo = Path(tempfile.mkdtemp())
    subprocess.run(["git", "init", "-q", str(repo)], check=True)
    big = b"large line contents\n" * 150_000  # ~3 MB each, far past MAX_DIFF_BYTES
    for name in ("a.txt", "b.txt", "c.txt"):
        (repo / name).write_bytes(big)
    reads: list[int] = []
    orig = gitinfo._git_diff

    def counted(*a, **k):
        out = orig(*a, **k)
        reads.append(len(out))
        return out

    gitinfo._git_diff = counted
    try:
        info = gitinfo.diff(str(repo)) or {}
        diff_reads = list(reads)
        cp = gitinfo.checkpoint(str(repo)) or {}
    finally:
        gitinfo._git_diff = orig
    # Bounded reading, not just a short reply: diff() pulls one budget's worth
    # from git in total, and checkpoint() reuses the same capped reads.
    return (info.get("truncated") and cp.get("truncated") and "large line" in info.get("diff", "")
            and sum(diff_reads) <= gitinfo.MAX_DIFF_BYTES + 1
            and sum(reads) <= 2 * gitinfo.MAX_DIFF_BYTES + 10_000)


def main() -> int:
    ok = True
    for fn in (check_llm_redirect, check_body_413, check_write_temp_unique, check_write_chmod_failure_cleans_up, check_create_rollback,
               check_upload_symlink_swap, check_share_cleanup_contained, check_login_throttle, check_llm_response_cap,
               check_body_must_be_object, check_store_bad_shape_uses_backup, check_store_bad_settings_uses_backup, check_auto_resume_bad_timestamp,
               check_diff_reads_bounded):
        try:
            good = fn()
        except Exception as exc:  # noqa: BLE001
            print(f"FAIL {fn.__name__}: {exc!r}")
            ok = False
            continue
        print(("ok   " if good else "FAIL ") + fn.__name__)
        ok &= good
    return 0 if ok else 1


if __name__ == "__main__":
    sys.exit(main())