browser-cli/tests/test_serve.py

"""Unit tests for the TCP serve layer (challenge-response auth, framing, rejection paths)."""
import json
import socket
import struct
import threading

import pytest

from browser_cli import transport
from browser_cli.auth import generate_keypair, load_private_key, new_nonce, pq_decrypt, pq_encrypt, sign
from browser_cli.client import BrowserNotConnected
from browser_cli.commands.serve import _handle_client

FAKE_UA = "browser-cli/0.9.3"


# ── helpers ────────────────────────────────────────────────────────────────────

def _send_framed(sock: socket.socket, data: bytes) -> None:
    sock.sendall(struct.pack("<I", len(data)) + data)

def _recv_framed(sock: socket.socket) -> dict:
    raw = b""
    while len(raw) < 4:
        chunk = sock.recv(4 - len(raw))
        if not chunk:
            raise ConnectionError("socket closed before response header")
        raw += chunk
    n = struct.unpack("<I", raw)[0]
    data = b""
    while len(data) < n:
        chunk = sock.recv(n - len(data))
        if not chunk:
            raise ConnectionError("socket closed mid-response")
        data += chunk
    return json.loads(data)

def _spawn(server_sock: socket.socket, auth_keys_path) -> threading.Thread:
    t = threading.Thread(
        target=_handle_client,
        args=(server_sock, ("127.0.0.1", 9999), None, auth_keys_path),
        daemon=True,
    )
    t.start()
    return t

def _pair():
    return socket.socketpair()

def _mock_no_browser(*_args, **_kwargs):
    raise BrowserNotConnected("no browser")

# ── challenge frame ────────────────────────────────────────────────────────────

class TestChallenge:
    def test_challenge_sent_on_connect(self):
        client, server = _pair()
        t = _spawn(server, None)
        challenge = _recv_framed(client)
        assert challenge["type"] == "challenge"
        assert "nonce" in challenge
        client.close()
        t.join(timeout=2)

    def test_challenge_includes_version_fields(self):
        client, server = _pair()
        t = _spawn(server, None)
        challenge = _recv_framed(client)
        assert "server_version" in challenge
        assert "min_client_version" in challenge
        client.close()
        t.join(timeout=2)

    def test_challenge_advertises_post_quantum_kex_when_available(self, tmp_path, monkeypatch):
        monkeypatch.setattr("browser_cli.auth.pq_kex_server_keypair", lambda: ("fake-private", b"fake-public"))
        path = tmp_path / "authorized_keys"
        _, pub = generate_keypair()
        path.write_text(pub + "\n")

        client, server = _pair()
        t = _spawn(server, path)
        challenge = _recv_framed(client)

        assert challenge["pq_kex"] == {"alg": "ML-KEM-768", "public_key": b"fake-public".hex()}
        client.close()
        t.join(timeout=2)

# ── rejection paths ────────────────────────────────────────────────────────────

class TestRejection:
    def _connect(self, auth_keys_path):
        client, server = _pair()
        t = _spawn(server, auth_keys_path)
        challenge = _recv_framed(client)
        return client, t, challenge

    def test_bad_user_agent_rejected(self):
        client, t, _ = self._connect(None)
        msg = {"id": "x", "command": "tabs.list", "args": {}, "user_agent": "curl/7.88"}
        _send_framed(client, json.dumps(msg).encode())
        resp = _recv_framed(client)
        assert resp["success"] is False
        assert "forbidden" in resp["error"].lower() or "client" in resp["error"].lower()
        client.close()
        t.join(timeout=2)

    def test_missing_pubkey_sig_rejected(self, tmp_path):
        path = tmp_path / "authorized_keys"
        _, pub = generate_keypair()
        path.write_text(pub + "\n")
        client, t, _ = self._connect(path)
        msg = {"id": "x", "command": "tabs.list", "args": {}, "user_agent": FAKE_UA}
        _send_framed(client, json.dumps(msg).encode())
        resp = _recv_framed(client)
        assert resp["success"] is False
        assert "unauthorized" in resp["error"].lower()
        client.close()
        t.join(timeout=2)

    def test_untrusted_pubkey_rejected(self, tmp_path):
        path = tmp_path / "authorized_keys"
        _, trusted_pub = generate_keypair()
        path.write_text(trusted_pub + "\n")

        pem, untrusted_pub = generate_keypair()
        key_path = tmp_path / "other.pem"
        key_path.write_bytes(pem)
        priv = load_private_key(key_path)

        client, t, challenge = self._connect(path)
        nonce = bytes.fromhex(challenge["nonce"])
        msg = {"id": "x", "command": "tabs.list", "args": {}, "user_agent": FAKE_UA, "pubkey": untrusted_pub}
        msg["sig"] = sign(priv, nonce, msg).hex()
        _send_framed(client, json.dumps(msg).encode())
        resp = _recv_framed(client)
        assert resp["success"] is False
        assert "untrusted" in resp["error"].lower()
        client.close()
        t.join(timeout=2)

    def test_bad_signature_rejected(self, tmp_path):
        path = tmp_path / "authorized_keys"
        _, pub = generate_keypair()
        path.write_text(pub + "\n")
        client, t, _ = self._connect(path)
        msg = {"id": "x", "command": "tabs.list", "args": {}, "user_agent": FAKE_UA, "pubkey": pub, "sig": "00" * 64}
        _send_framed(client, json.dumps(msg).encode())
        resp = _recv_framed(client)
        assert resp["success"] is False
        assert "signature" in resp["error"].lower() or "invalid" in resp["error"].lower()
        client.close()
        t.join(timeout=2)

    def test_missing_post_quantum_kex_rejected_when_required(self, tmp_path, monkeypatch):
        monkeypatch.setattr("browser_cli.auth.pq_kex_server_keypair", lambda: ("fake-private", b"fake-public"))
        path = tmp_path / "authorized_keys"
        pem, pub = generate_keypair()
        path.write_text(pub + "\n")
        priv_path = tmp_path / "client.pem"
        priv_path.write_bytes(pem)
        priv = load_private_key(priv_path)

        client, t, challenge = self._connect(path)
        nonce = bytes.fromhex(challenge["nonce"])
        msg = {"id": "x", "command": "tabs.list", "args": {}, "user_agent": "browser-cli/0.9.5", "pubkey": pub}
        msg["sig"] = sign(priv, nonce, msg).hex()
        _send_framed(client, json.dumps(msg).encode())
        resp = _recv_framed(client)

        assert resp["success"] is False
        assert "post-quantum" in resp["error"].lower()
        client.close()
        t.join(timeout=2)

    def test_oversized_message_rejected(self):
        client, server = _pair()
        t = _spawn(server, None)
        _recv_framed(client)  # consume challenge
        client.sendall(struct.pack("<I", 33 * 1024 * 1024))
        resp = _recv_framed(client)
        assert resp["success"] is False
        assert "too large" in resp["error"].lower()
        client.close()
        t.join(timeout=2)

    def test_invalid_json_rejected(self):
        client, server = _pair()
        t = _spawn(server, None)
        _recv_framed(client)  # consume challenge
        bad = b"this is not json {"
        _send_framed(client, bad)
        resp = _recv_framed(client)
        assert resp["success"] is False
        client.close()
        t.join(timeout=2)

# ── auth success paths ─────────────────────────────────────────────────────────

class TestAuthSuccess:
    def test_valid_auth_reaches_proxy(self, tmp_path, monkeypatch):
        """Correct signature → error must be 'browser not connected', not 'unauthorized'."""
        path = tmp_path / "authorized_keys"
        pem, pub = generate_keypair()
        path.write_text(pub + "\n")
        key_path = tmp_path / "client.key.pem"
        key_path.write_bytes(pem)
        priv = load_private_key(key_path)

        monkeypatch.setattr("browser_cli.client._resolve_socket", _mock_no_browser)

        client, server = _pair()
        t = threading.Thread(
            target=_handle_client,
            args=(server, ("127.0.0.1", 9999), None, path),
            daemon=True,
        )
        t.start()

        challenge = _recv_framed(client)
        nonce = bytes.fromhex(challenge["nonce"])
        msg = {"id": "x", "command": "tabs.list", "args": {}, "user_agent": FAKE_UA, "pubkey": pub}
        msg["sig"] = sign(priv, nonce, msg).hex()
        _send_framed(client, json.dumps(msg).encode())
        resp = _recv_framed(client)

        assert resp["success"] is False
        assert "unauthorized" not in resp["error"].lower()
        assert "browser" in resp["error"].lower() or "connected" in resp["error"].lower()
        client.close()
        t.join(timeout=2)

    def test_uppercase_pubkey_normalized_by_compat(self, tmp_path, monkeypatch):
        """Clients < 0.9.3 may send uppercase pubkeys; compat layer normalises before auth."""
        path = tmp_path / "authorized_keys"
        pem, pub = generate_keypair()  # pub is lowercase hex
        path.write_text(pub + "\n")
        key_path = tmp_path / "client.key.pem"
        key_path.write_bytes(pem)
        priv = load_private_key(key_path)

        monkeypatch.setattr("browser_cli.client._resolve_socket", _mock_no_browser)

        client, server = _pair()
        t = _spawn(server, path)

        challenge = _recv_framed(client)
        nonce = bytes.fromhex(challenge["nonce"])
        # old client sends uppercase pubkey
        msg = {"id": "x", "command": "tabs.list", "args": {}, "user_agent": "browser-cli/0.9.2", "pubkey": pub.upper()}
        msg["sig"] = sign(priv, nonce, msg).hex()
        _send_framed(client, json.dumps(msg).encode())
        resp = _recv_framed(client)

        assert "unauthorized" not in resp.get("error", "").lower()
        assert "browser" in resp.get("error", "").lower() or "connected" in resp.get("error", "").lower()
        client.close()
        t.join(timeout=2)

    def test_post_quantum_kex_auth_reaches_proxy(self, tmp_path, monkeypatch):
        """ML-KEM shared secret is decapsulated and bound to the auth signature."""
        monkeypatch.setattr("browser_cli.client._resolve_socket", _mock_no_browser)
        monkeypatch.setattr("browser_cli.auth.pq_kex_server_keypair", lambda: ("fake-private", b"fake-public"))
        monkeypatch.setattr("browser_cli.auth.pq_kex_server_decapsulate", lambda priv, ct: b"pq-secret")

        path = tmp_path / "authorized_keys"
        pem, pub = generate_keypair()
        path.write_text(pub + "\n")
        key_path = tmp_path / "client.key.pem"
        key_path.write_bytes(pem)
        priv = load_private_key(key_path)

        client, server = _pair()
        t = threading.Thread(
            target=_handle_client,
            args=(server, ("127.0.0.1", 9999), None, path),
            daemon=True,
        )
        t.start()

        challenge = _recv_framed(client)
        nonce = bytes.fromhex(challenge["nonce"])
        msg = {
            "id": "x",
            "command": "tabs.list",
            "args": {},
            "user_agent": FAKE_UA,
            "pubkey": pub,
            "pq_kex": {"alg": "ML-KEM-768", "ciphertext": "cafe"},
        }
        msg["sig"] = sign(priv, nonce, msg, b"pq-secret").hex()
        _send_framed(client, json.dumps(msg).encode())
        resp = _recv_framed(client)

        assert "unauthorized" not in resp.get("error", "").lower()
        assert "browser" in resp.get("error", "").lower() or "connected" in resp.get("error", "").lower()
        client.close()
        t.join(timeout=2)

    def test_post_quantum_encrypted_transport_reaches_proxy(self, tmp_path, monkeypatch):
        """New clients encrypt the command payload and receive encrypted responses."""
        monkeypatch.setattr("browser_cli.client._resolve_socket", _mock_no_browser)
        monkeypatch.setattr("browser_cli.auth.pq_kex_server_keypair", lambda: ("fake-private", b"fake-public"))
        monkeypatch.setattr("browser_cli.auth.pq_kex_server_decapsulate", lambda priv, ct: b"pq-secret")

        path = tmp_path / "authorized_keys"
        pem, pub = generate_keypair()
        path.write_text(pub + "\n")
        key_path = tmp_path / "client.key.pem"
        key_path.write_bytes(pem)
        priv = load_private_key(key_path)

        client, server = _pair()
        t = threading.Thread(
            target=_handle_client,
            args=(server, ("127.0.0.1", 9999), None, path),
            daemon=True,
        )
        t.start()

        challenge = _recv_framed(client)
        nonce = bytes.fromhex(challenge["nonce"])
        clean_msg = {
            "id": "x",
            "command": "tabs.list",
            "args": {},
            "user_agent": "browser-cli/0.9.5",
            "pq_kex": {"alg": "ML-KEM-768", "ciphertext": "cafe"},
        }
        sig = sign(priv, nonce, clean_msg, b"pq-secret").hex()
        envelope = {
            "id": "x",
            "user_agent": "browser-cli/0.9.5",
            "pubkey": pub,
            "sig": sig,
            "pq_kex": clean_msg["pq_kex"],
            "encrypted": pq_encrypt(b"pq-secret", "request", json.dumps(clean_msg).encode()),
        }
        _send_framed(client, json.dumps(envelope).encode())
        encrypted_resp = _recv_framed(client)

        assert "encrypted" in encrypted_resp
        resp = json.loads(pq_decrypt(b"pq-secret", "response", encrypted_resp["encrypted"]))
        assert "unauthorized" not in resp.get("error", "").lower()
        assert "browser" in resp.get("error", "").lower() or "connected" in resp.get("error", "").lower()
        client.close()
        t.join(timeout=2)

    def test_no_auth_mode_reaches_proxy(self, monkeypatch):
        """auth_keys_path=None (--no-auth): no pubkey required, reaches proxy layer."""
        monkeypatch.setattr("browser_cli.client._resolve_socket", _mock_no_browser)

        client, server = _pair()
        t = threading.Thread(
            target=_handle_client,
            args=(server, ("127.0.0.1", 9999), None, None),
            daemon=True,
        )
        t.start()

        _recv_framed(client)  # challenge
        msg = {"id": "x", "command": "tabs.list", "args": {}, "user_agent": FAKE_UA}
        _send_framed(client, json.dumps(msg).encode())
        resp = _recv_framed(client)

        assert "unauthorized" not in resp.get("error", "").lower()
        client.close()
        t.join(timeout=2)

# ── response encoding (compression / msgpack) ───────────────────────────────────

def _recv_framed_raw(sock: socket.socket) -> bytes:
    raw = b""
    while len(raw) < 4:
        chunk = sock.recv(4 - len(raw))
        if not chunk:
            raise ConnectionError("socket closed before response header")
        raw += chunk
    n = struct.unpack("<I", raw)[0]
    data = b""
    while len(data) < n:
        chunk = sock.recv(n - len(data))
        if not chunk:
            raise ConnectionError("socket closed mid-response")
        data += chunk
    return data

class _FakeNativeHost:
    """Minimal AF_UNIX server speaking the 4-byte framed protocol the proxy expects."""

    def __init__(self, sock_path, response_obj):
        self.sock_path = str(sock_path)
        self.response = json.dumps(response_obj).encode()
        self.srv = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
        self.srv.bind(self.sock_path)
        self.srv.listen(1)
        self.thread = threading.Thread(target=self._serve, daemon=True)
        self.thread.start()

    def _serve(self):
        try:
            conn, _ = self.srv.accept()
            with conn:
                hdr = conn.recv(4)
                n = struct.unpack("<I", hdr)[0]
                got = b""
                while len(got) < n:
                    got += conn.recv(n - len(got))
                conn.sendall(struct.pack("<I", len(self.response)) + self.response)
        except OSError:
            pass

    def close(self):
        self.srv.close()

class TestResponseEncoding:
    def test_client_accept_encoding_yields_decodable_tagged_response(self, tmp_path, monkeypatch):
        if socket.AF_UNIX is None:  # pragma: no cover
            pytest.skip("AF_UNIX unavailable")

        big = {"id": "x", "success": True,
               "data": {"items": [{"url": f"https://example.com/{i}", "title": f"Tab {i}"} for i in range(300)]}}
        host_path = tmp_path / "native.sock"
        host = _FakeNativeHost(host_path, big)
        monkeypatch.setattr("browser_cli.client._resolve_socket", lambda *_a, **_k: str(host_path))

        client, server = _pair()
        t = _spawn(server, None)  # no auth
        _recv_framed(client)  # challenge

        msg = {
            "id": "x", "command": "tabs.list", "args": {},
            "user_agent": "browser-cli/0.9.5",
            "accept_encoding": {"ser": ["json"], "comp": ["gzip"]},
        }
        _send_framed(client, json.dumps(msg).encode())
        raw = _recv_framed_raw(client)

        assert raw[:1] not in (b"{", b"[")  # tagged, not plain JSON
        assert len(raw) < len(json.dumps(big))  # compressed
        assert transport.decode_response(raw) == big
        client.close()
        host.close()
        t.join(timeout=2)

    def test_no_accept_encoding_stays_plain_json(self, tmp_path, monkeypatch):
        big = {"id": "y", "success": True, "data": {"items": list(range(500))}}
        host_path = tmp_path / "native2.sock"
        host = _FakeNativeHost(host_path, big)
        monkeypatch.setattr("browser_cli.client._resolve_socket", lambda *_a, **_k: str(host_path))

        client, server = _pair()
        t = _spawn(server, None)
        _recv_framed(client)

        msg = {"id": "y", "command": "tabs.list", "args": {}, "user_agent": "browser-cli/0.9.5"}
        _send_framed(client, json.dumps(msg).encode())
        raw = _recv_framed_raw(client)

        assert raw[:1] in (b"{", b"[")  # old client → plain JSON
        assert json.loads(raw) == big
        client.close()
        host.close()
        t.join(timeout=2)

    def test_no_compress_flag_forces_plain_json(self, tmp_path, monkeypatch):
        big = {"id": "z", "success": True,
               "data": {"items": [{"url": f"https://e/{i}"} for i in range(300)]}}
        host_path = tmp_path / "native3.sock"
        host = _FakeNativeHost(host_path, big)
        monkeypatch.setattr("browser_cli.client._resolve_socket", lambda *_a, **_k: str(host_path))

        client, server = _pair()
        t = threading.Thread(
            target=_handle_client,
            args=(server, ("127.0.0.1", 9999), None, None, False),  # compress=False
            daemon=True,
        )
        t.start()
        _recv_framed(client)

        msg = {"id": "z", "command": "tabs.list", "args": {},
               "user_agent": "browser-cli/0.9.5",
               "accept_encoding": {"ser": ["msgpack"], "comp": ["zstd"]}}
        _send_framed(client, json.dumps(msg).encode())
        raw = _recv_framed_raw(client)

        assert raw[:1] in (b"{", b"[")  # server disabled encoding
        assert json.loads(raw) == big
        client.close()
        host.close()
        t.join(timeout=2)