#!/usr/bin/env python3
"""hush-sealed-client.py - reference client for Hush AI's SEALED chat lanes.

Cloudflare (and every other hop) only ever sees ciphertext: your prompt, your API key and
the model's answer are encrypted in THIS process to a key held only on Hush's own hardware.
Construction (byte-identical to the server's chat_sealing.py and the site's hush-chat-seal.js):

    ikm  = ECDH(ephemeral P-256, origin P-256)  ||  ML-KEM-768 shared secret   (hybrid, v2)
    keys = HKDF-SHA256(ikm, salt, info per direction)  ->  AES-256-GCM request / reply / stream

Requires:  python3 >= 3.9,  pip install "cryptography>=45"      (HTTP is stdlib urllib)
           (cryptography 45+ ships ML-KEM; older builds fall back to classical v1, which the
            origin REFUSES when it mandates post-quantum - you will see a clear error.)

Usage:
    python3 hush-sealed-client.py --key YOUR_API_KEY "Summarise this clause ..."
    python3 hush-sealed-client.py --key YOUR_API_KEY --model sovereign-omega --stream "..."
    python3 hush-sealed-client.py --try "Hello"                   # signed-out demo lane
    python3 hush-sealed-client.py --pin <spk> ...                  # pin the signing key (recommended)

The pin: Hush publishes its long-term signing key in the site repo (deploy/seal/PINS.txt) and
at https://hush-ai.uk/v1/crypto/pubkey as "spk". Pass it with --pin and this client refuses to
send anything unless the sealing keys arrive signed by it - a proxy that swaps keys gets an
error, never your plaintext. Without --pin the signature is still verified against whatever
spk is served (trust-on-first-use); pin it for the strong guarantee.

Nothing here is secret: publish, audit, diff it against the site copy at
https://hush-ai.uk/static/hush-sealed-client.py at any time.
"""
from __future__ import annotations

import argparse
import base64
import json
import os
import sys
import time

try:
    from cryptography.hazmat.primitives import hashes, serialization
    from cryptography.hazmat.primitives.asymmetric import ec
    from cryptography.hazmat.primitives.asymmetric.utils import encode_dss_signature
    from cryptography.hazmat.primitives.ciphers.aead import AESGCM
    from cryptography.hazmat.primitives.kdf.hkdf import HKDF
except ImportError:
    sys.exit("pip install 'cryptography>=45'")
try:
    from cryptography.hazmat.primitives.asymmetric import mlkem as _mlkem
except ImportError:                      # cryptography < 45: no post-quantum layer available
    _mlkem = None
import urllib.error
import urllib.request

BASE = os.environ.get("HUSH_BASE", "https://hush-ai.uk")
SALT = b"hush-chat-seal-salt-v1"
INFO = {
    (1, "c2s"): b"hush-chat-seal-v1-c2s", (1, "s2c"): b"hush-chat-seal-v1-s2c",
    (1, "stream"): b"hush-chat-seal-v1-s2c-stream",
    (2, "c2s"): b"hush-chat-seal-v2-hybrid-c2s", (2, "s2c"): b"hush-chat-seal-v2-hybrid-s2c",
    (2, "stream"): b"hush-chat-seal-v2-hybrid-s2c-stream",
}
STREAM_AAD = b"hush-chat-s2c"
FIN = b"HUSH-FIN:"
PIN_DOMAIN = b"hush-chat-seal-pin-v1|"
LANES = {"api": ("/v1/sealed/chat/completions", True), "try": ("/try/sealed/chat", False)}


def b64u(b: bytes) -> str:
    return base64.urlsafe_b64encode(b).rstrip(b"=").decode()


def unb64u(s: str) -> bytes:
    s = s.replace("-", "+").replace("_", "/")
    return base64.b64decode(s + "=" * (-len(s) % 4))


def derive(ikm: bytes, ver: int, direction: str) -> bytes:
    return HKDF(algorithm=hashes.SHA256(), length=32, salt=SALT,
                info=INFO[(ver, direction)]).derive(ikm)


def verify_signed_keys(k: dict, pin: str) -> None:
    """Refuse the key document unless spk signs (ec, mlkem) and, when pinned, spk == pin."""
    spk, sig = k.get("spk"), k.get("sig")
    if not spk or not sig:
        sys.exit("refusing: origin served UNSIGNED sealing keys (no spk/sig)")
    if pin and spk != pin:
        sys.exit("refusing: served signing key does not match your --pin (possible substitution)")
    pub = ec.EllipticCurvePublicKey.from_encoded_point(ec.SECP256R1(), unb64u(spk))
    raw = unb64u(sig)
    if len(raw) != 64:
        sys.exit("refusing: malformed signature")
    der = encode_dss_signature(int.from_bytes(raw[:32], "big"), int.from_bytes(raw[32:], "big"))
    msg = PIN_DOMAIN + k["ec"].encode() + b"|" + (k.get("mlkem") or "").encode()
    try:
        pub.verify(der, msg, ec.ECDSA(hashes.SHA256()))
    except Exception:
        sys.exit("refusing: sealing keys FAILED signature verification - do not trust this endpoint")


def seal(keys: dict, payload: dict) -> tuple[dict, int, bytes]:
    eph = ec.generate_private_key(ec.SECP256R1())
    origin = ec.EllipticCurvePublicKey.from_encoded_point(ec.SECP256R1(), unb64u(keys["ec"]))
    ikm = eph.exchange(ec.ECDH(), origin)
    ver, env = 1, {}
    if keys.get("mlkem"):
        if _mlkem is None:
            if keys.get("pq_required"):
                sys.exit("origin requires post-quantum sealing: pip install 'cryptography>=45'")
        else:
            ss, ct = _mlkem.MLKEM768PublicKey.from_public_bytes(unb64u(keys["mlkem"])).encapsulate()
            ikm += ss
            ver, env["kem"] = 2, b64u(ct)
    elif keys.get("pq_required"):
        sys.exit("origin requires post-quantum sealing but advertises no ML-KEM key - refusing")
    # pad to 1 KiB buckets so ciphertext length does not reveal prompt length
    body = json.dumps(payload).encode()
    payload["pad"] = "0" * ((-(len(body) + 9)) % 1024)
    iv = os.urandom(12)
    ct = AESGCM(derive(ikm, ver, "c2s")).encrypt(iv, json.dumps(payload).encode(), None)
    env.update({"v": ver, "epk": b64u(eph.public_key().public_bytes(
        serialization.Encoding.X962, serialization.PublicFormat.UncompressedPoint)),
        "iv": b64u(iv), "ct": b64u(ct)})
    return env, ver, ikm


def open_json(env: dict, ver: int, ikm: bytes) -> dict:
    pt = AESGCM(derive(ikm, ver, "s2c")).decrypt(unb64u(env["iv"]), unb64u(env["ct"]), None)
    return json.loads(pt)


def http(method: str, path: str, body: dict | None = None):
    """stdlib HTTP; returns the response object (never raises on 4xx/5xx - caller decides)."""
    data = json.dumps(body).encode() if body is not None else None
    req = urllib.request.Request(BASE + path, data=data, method=method,
                                 headers={"Content-Type": "application/json", "Accept": "*/*",
                                          "User-Agent": "hush-sealed-client/1.0"})
    try:
        return urllib.request.urlopen(req, timeout=300)
    except urllib.error.HTTPError as e:
        return e                                              # has .status/.headers/.read()


def open_stream(resp, ver: int, ikm: bytes):
    """Yield plaintext SSE bytes; verify order + the authenticated FIN (truncation-proof)."""
    aes, nxt, done = AESGCM(derive(ikm, ver, "stream")), 0, False
    for line in resp:
        line = line.decode(errors="replace").rstrip("\r\n")
        if not line.startswith("data: "):
            continue
        if line == "data: [SEALED-DONE]":
            if not done:
                sys.exit("sealed stream truncated (no FIN)")
            continue
        ev = json.loads(line[6:])
        if ev["n"] != nxt:
            sys.exit("sealed stream out of order - tampered or broken")
        pt = aes.decrypt(ev["n"].to_bytes(12, "big"), unb64u(ev["ct"]), STREAM_AAD)
        if ev.get("fin"):
            if pt != FIN + str(ev["n"]).encode():
                sys.exit("sealed stream FIN mismatch - truncated or tampered")
            done = True
            continue
        nxt += 1
        yield pt
    if not done:
        sys.exit("sealed stream ended without FIN - truncated")


def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("prompt")
    ap.add_argument("--key", help="your Hush API key (rides INSIDE the envelope; Cloudflare never sees it)")
    ap.add_argument("--try", dest="try_lane", action="store_true", help="signed-out demo lane (no key)")
    ap.add_argument("--model", default="sovereign-prism")
    ap.add_argument("--stream", action="store_true")
    ap.add_argument("--pin", default=os.environ.get("HUSH_PIN", ""), help="expected signing key (spk)")
    ap.add_argument("--max-tokens", type=int, default=512)
    a = ap.parse_args()
    lane, needs_key = LANES["try" if a.try_lane else "api"]
    if needs_key and not a.key:
        ap.error("--key is required for the API lane (or use --try for the demo lane)")

    kr = http("GET", "/v1/crypto/pubkey")
    if getattr(kr, "status", 0) != 200:
        sys.exit(f"cannot fetch {BASE}/v1/crypto/pubkey (HTTP {getattr(kr, 'status', '?')})")
    keys = json.loads(kr.read())
    if not keys.get("enabled"):
        sys.exit("origin reports sealing disabled - refusing to send plaintext")
    verify_signed_keys(keys, a.pin)
    offset = keys.get("now", time.time()) - time.time()      # tolerate a skewed local clock
    body = {"model": a.model, "stream": a.stream, "max_tokens": a.max_tokens,
            "messages": [{"role": "user", "content": a.prompt}]}
    inner = {"authorization": "Bearer " + a.key} if a.key else {}
    env, ver, ikm = seal(keys, {"ts": time.time() + offset, "headers": inner, "body": body})
    print(f"[sealed v={ver}{' hybrid post-quantum' if ver == 2 else ' classical'} -> {lane}]",
          file=sys.stderr)

    r = http("POST", lane, env)
    status = getattr(r, "status", 0)
    sealed = r.headers.get("X-Hush-Sealed")
    if 200 <= status < 300 and not sealed:
        sys.exit("unsealed 2xx on the sealed lane - refusing to trust it (on-path forgery?)")
    if not sealed:                                           # seal-layer error: plain, no content
        sys.exit(f"HTTP {status}: {r.read()[:300].decode(errors='replace')}")
    if a.stream:
        buf = b""
        for pt in open_stream(r, ver, ikm):
            buf += pt
            while b"\n\n" in buf:
                block, buf = buf.split(b"\n\n", 1)
                block = block.decode(errors="replace").strip()
                if block.startswith("data: ") and block != "data: [DONE]":
                    d = json.loads(block[6:])
                    sys.stdout.write(d.get("choices", [{}])[0].get("delta", {}).get("content") or "")
                    sys.stdout.flush()
        print()
        return 0
    data = open_json(json.loads(r.read()), ver, ikm)
    if status != 200:
        sys.exit(f"HTTP {status}: {json.dumps(data)[:300]}")
    print(data["choices"][0]["message"]["content"])
    return 0


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