#!/usr/bin/env python3
"""
mock_jwks.py — minimal JWKS server for pg_knoe_auth integration tests.

Generates an RSA-2048 key pair on startup, serves:
  GET /jwks.json          → {"keys":[{...RSA public key JWK...}]}
  GET /token              → signed JWT for the test user
  GET /token/expired      → JWT with exp in the past
  GET /token/wrong_iss    → JWT with wrong issuer
  GET /token/wrong_aud    → JWT with wrong audience
  GET /token/tampered     → JWT with last byte of signature flipped
  GET /token/unknown_kid  → JWT with kid not in JWKS
  POST /control/down      → stop serving JWKS (simulate unreachable issuer)
  POST /control/up        → resume serving JWKS

Usage:
  python3 mock_jwks.py [--port 9999]

The server writes the RSA private key and kid to stdout on startup so
the test harness can generate tokens independently if needed.

See docs/plans/junie/pg_knoe_auth-rename-harden-modularize.md §2.5
"""

import argparse
import base64
import json
import struct
import time
import threading
from http.server import BaseHTTPRequestHandler, HTTPServer

try:
    from cryptography.hazmat.primitives.asymmetric import rsa, padding
    from cryptography.hazmat.primitives import hashes, serialization
    from cryptography.hazmat.backends import default_backend
except ImportError:
    raise SystemExit("pip install cryptography")

# ── Key generation ────────────────────────────────────────────────────────────

PRIVATE_KEY = rsa.generate_private_key(
    public_exponent=65537,
    key_size=2048,
    backend=default_backend(),
)
PUBLIC_KEY = PRIVATE_KEY.public_key()
KID = "test-key-001"
ISSUER = "http://localhost:9999"
AUDIENCE = "pg.0.knoe.dev"
TEST_USER = "chrisfu"

_jwks_down = False  # controlled via /control/down and /control/up


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


def _int_to_b64url(n: int) -> str:
    length = (n.bit_length() + 7) // 8
    return _b64url(n.to_bytes(length, "big"))


def _make_jwks() -> dict:
    pub = PUBLIC_KEY.public_numbers()
    return {
        "keys": [
            {
                "kty": "RSA",
                "use": "sig",
                "alg": "RS256",
                "kid": KID,
                "n": _int_to_b64url(pub.n),
                "e": _int_to_b64url(pub.e),
            }
        ]
    }


def _make_jwt(claims: dict, kid: str = KID, tamper: bool = False) -> str:
    header = {"alg": "RS256", "typ": "JWT", "kid": kid}
    h = _b64url(json.dumps(header, separators=(",", ":")).encode())
    p = _b64url(json.dumps(claims, separators=(",", ":")).encode())
    signing_input = f"{h}.{p}".encode()
    sig_bytes = PRIVATE_KEY.sign(signing_input, padding.PKCS1v15(), hashes.SHA256())
    if tamper:
        sig_list = bytearray(sig_bytes)
        sig_list[-1] ^= 0xFF
        sig_bytes = bytes(sig_list)
    return f"{h}.{p}.{_b64url(sig_bytes)}"


def _base_claims(extra: dict = None) -> dict:
    now = int(time.time())
    claims = {
        "iss": ISSUER,
        "aud": AUDIENCE,
        "sub": TEST_USER,
        "preferred_username": TEST_USER,
        "iat": now,
        "exp": now + 3600,
    }
    if extra:
        claims.update(extra)
    return claims


# Pre-generate tokens
TOKENS = {
    "valid": _make_jwt(_base_claims()),
    "expired": _make_jwt(_base_claims({"exp": int(time.time()) - 10})),
    "wrong_iss": _make_jwt(_base_claims({"iss": "https://evil.example.com"})),
    "wrong_aud_string": _make_jwt(_base_claims({"aud": "other.service"})),
    "wrong_aud_array": _make_jwt(_base_claims({"aud": ["other.service"]})),
    "right_aud_array": _make_jwt(_base_claims({"aud": [AUDIENCE]})),
    "tampered": _make_jwt(_base_claims(), tamper=True),
    "unknown_kid": _make_jwt(_base_claims(), kid="no-such-kid"),
}

JWKS_JSON = json.dumps(_make_jwks()).encode()


# ── HTTP handler ──────────────────────────────────────────────────────────────


class Handler(BaseHTTPRequestHandler):
    def log_message(self, fmt, *args):
        pass  # suppress default access log

    def _send(self, code: int, body: bytes, content_type: str = "application/json"):
        self.send_response(code)
        self.send_header("Content-Type", content_type)
        self.send_header("Content-Length", str(len(body)))
        self.end_headers()
        self.wfile.write(body)

    def do_GET(self):
        global _jwks_down
        if self.path == "/jwks.json" or self.path == "/.well-known/jwks.json":
            if _jwks_down:
                self._send(503, b'{"error":"simulated outage"}')
            else:
                self._send(200, JWKS_JSON)
        elif self.path.startswith("/token"):
            key = self.path.split("/token")[-1].lstrip("/") or "valid"
            token = TOKENS.get(key)
            if token:
                self._send(200, token.encode(), "text/plain")
            else:
                self._send(404, b"unknown token type")
        else:
            self._send(404, b"not found")

    def do_POST(self):
        global _jwks_down
        if self.path == "/control/down":
            _jwks_down = True
            self._send(200, b'{"status":"down"}')
        elif self.path == "/control/up":
            _jwks_down = False
            self._send(200, b'{"status":"up"}')
        else:
            self._send(404, b"not found")


# ── Entry point ───────────────────────────────────────────────────────────────

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--port", type=int, default=9999)
    args = parser.parse_args()

    server = HTTPServer(("127.0.0.1", args.port), Handler)
    print(f"mock_jwks: listening on http://127.0.0.1:{args.port}", flush=True)
    print(f"mock_jwks: kid={KID}", flush=True)
    print(f"mock_jwks: issuer={ISSUER}", flush=True)
    try:
        server.serve_forever()
    except KeyboardInterrupt:
        pass