ZKAC/demo/server.py

#!/usr/bin/env python3
"""
HTTP static site + ZKAC TCP service. Authenticated /api is accessed over TCP with zkac.tcp
after setup_demo.py has created creds/.
"""

from __future__ import annotations

import argparse
import base64
import json
import os
import threading
import http.server
import socket
import traceback
from pathlib import Path

import zkac
from zkac.tcp import FramedSession, server_handshake


def load_registry(creds_dir: Path, epoch: int) -> zkac.RoleRegistry:
    """Load issuer public key and register every demo role at the same epoch."""
    iss = json.loads((creds_dir / "issuer.json").read_text(encoding="utf-8"))
    issuer_pk = zkac.BbsPublicKey.from_bytes(
        base64.b64decode(iss["issuer_public_key_b64"])
    )
    reg = zkac.RoleRegistry()
    for name in ("analyst", "operator"):
        reg.register_role(zkac.role_id(name), issuer_pk, epoch)
    return reg


def _role_debug_label(role_id: bytes) -> str:
    """Map verified role_id bytes to a short label for logs (demo only)."""
    for name in ("analyst", "operator"):
        if role_id == zkac.role_id(name):
            return name
    return "unknown"


def api_body_for_role(role_id: bytes) -> dict:
    """JSON returned for the logical /api resource after ZKAC auth; varies by credential role."""
    analyst = zkac.role_id("analyst")
    operator = zkac.role_id("operator")
    if role_id == analyst:
        return {
            "path": "/api",
            "role": "analyst",
            "datasets": ["summary", "aggregated_metrics"],
            "note": "Analyst tier: aggregated data only.",
        }
    if role_id == operator:
        return {
            "path": "/api",
            "role": "operator",
            "datasets": ["summary", "aggregated_metrics", "raw_logs", "pii"],
            "note": "Operator tier: full API slice including raw logs.",
        }
    return {"error": "unknown role", "path": "/api"}


def handle_zkac_client(
    conn: socket.socket,
    client_addr: tuple,
    creds_dir: Path,
    registry: zkac.RoleRegistry,
) -> None:
    """
    One TCP connection: ZKAC handshake + BBS+ auth, then one framed JSON request and response.
    Each handler rebuilds the server Node from persisted secret (Keypair is consumed by Node).
    """
    peer = f"{client_addr[0]}:{client_addr[1]}"
    print(f"[zkac] connect peer={peer}")

    try:
        # Same long-term server identity every time; from_secret_key because Node consumes Keypair.
        t = json.loads((creds_dir / "transport.json").read_text(encoding="utf-8"))
        sk = base64.b64decode(t["server_secret_key_b64"])
        server_kp = zkac.Keypair.from_secret_key(sk)
        node = zkac.Node(server_kp)

        session, role_id = server_handshake(conn, node, registry)
        label = _role_debug_label(role_id)
        print(
            f"[zkac] handshake_ok peer={peer} role_id={role_id.hex()} role={label!r}"
        )

        framed = FramedSession(conn, session)
        raw = framed.recv()
        print(
            f"[zkac] request peer={peer} plaintext_bytes={len(raw)} raw={raw!r}"
        )

        req = json.loads(raw.decode("utf-8"))
        print(f"[zkac] request_json peer={peer} parsed={req!r}")

        path = req.get("path")
        if path != "/api":
            err_body = {"error": "unsupported path", "allowed": ["/api"], "got": path}
            out = json.dumps(err_body).encode()
            framed.send(out)
            print(
                f"[zkac] response peer={peer} status=reject path={path!r} response_bytes={len(out)}"
            )
            return

        body = api_body_for_role(role_id)
        out_bytes = json.dumps(body).encode()
        framed.send(out_bytes)
        print(
            f"[zkac] response peer={peer} status=ok path=/api role={label!r} "
            f"response_bytes={len(out_bytes)} body_keys={list(body.keys())}"
        )
    except (ConnectionError, BrokenPipeError, OSError) as e:
        print(f"[zkac] peer={peer} connection_error: {e!r}")
    except (json.JSONDecodeError, ValueError) as e:
        print(f"[zkac] peer={peer} protocol_error: {e!r}")
    except Exception as e:
        print(f"[zkac] peer={peer} unexpected_error: {e!r}")
        traceback.print_exc()
    finally:
        conn.close()
        print(f"[zkac] closed peer={peer}")


def run_http(host: str, port: int, static_root: Path) -> None:
    # Process-wide CWD: only this thread should rely on relative paths after chdir.
    os.chdir(static_root)

    class Handler(http.server.SimpleHTTPRequestHandler):
        def log_message(self, fmt: str, *args) -> None:
            # Default fmt is like '%s - - [%s] %s' — include client address for debugging.
            try:
                line = fmt % args if args else fmt
            except (TypeError, ValueError):
                line = f"{fmt} {args}"
            peer_ip = self.client_address[0] if self.client_address else "?"
            peer_port = self.client_address[1] if len(self.client_address) > 1 else "?"
            print(f"[http] peer={peer_ip}:{peer_port} | {line.strip()}")

    http.server.HTTPServer((host, port), Handler).serve_forever()


def main() -> None:
    ap = argparse.ArgumentParser(description="ZKAC demo HTTP + TCP server")
    ap.add_argument(
        "--creds-dir",
        type=Path,
        default=Path(__file__).resolve().parent / "creds",
    )
    ap.add_argument("--http-host", default="127.0.0.1")
    ap.add_argument("--http-port", type=int, default=8765)
    ap.add_argument("--zkac-host", default="127.0.0.1")
    ap.add_argument("--zkac-port", type=int, default=9876)
    args = ap.parse_args()
    creds_dir: Path = args.creds_dir
    if not (creds_dir / "transport.json").is_file():
        raise SystemExit(f"Missing {creds_dir}/transport.json — run setup_demo.py first.")

    # Epoch must match the member files issued at setup (any member file is enough).
    member = json.loads((creds_dir / "member_analyst.json").read_text(encoding="utf-8"))
    epoch = int(member["epoch"])
    registry = load_registry(creds_dir, epoch)

    static_root = Path(__file__).resolve().parent / "static"
    if not static_root.is_dir():
        raise SystemExit(f"Missing static directory: {static_root}")

    http_thread = threading.Thread(
        target=run_http,
        args=(args.http_host, args.http_port, static_root),
        daemon=True,
    )
    http_thread.start()
    print(
        f"HTTP  http://{args.http_host}:{args.http_port}/  (static demo page)\n"
        f"ZKAC  {args.zkac_host}:{args.zkac_port}  (authenticated /api over TCP)"
    )

    zkac_sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    zkac_sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    zkac_sock.bind((args.zkac_host, args.zkac_port))
    zkac_sock.listen(8)
    while True:
        conn, addr = zkac_sock.accept()
        threading.Thread(
            target=handle_zkac_client,
            args=(conn, addr, creds_dir, registry),
            daemon=True,
        ).start()


if __name__ == "__main__":
    main()