net-kingdom/local-identity/src/local_identity/serve.py

"""
Minimal OIDC Authorization Code flow server for local-identity.

Binds to 127.0.0.1 only — external binding is explicitly unsupported.
In production mode (run_server) the socket is wrapped with TLS using an
auto-generated self-signed certificate.  In test mode callers can use
make_handler() with scheme="http" and a plain HTTPServer.

Endpoints
---------
  GET  /.well-known/openid-configuration   OIDC discovery document
  GET  /jwks                               JSON Web Key Set
  GET  /auth                               Authorization endpoint — login form
  POST /auth                               Form submission — user selection
  POST /token                              Token endpoint
  GET  /userinfo                           UserInfo endpoint (Bearer token)

Design constraints (from NK-WP-0002)
--------------------------------------
  * No client secret validation (all registered clients trusted — dev only)
  * No refresh tokens (stateless; re-auth required after expiry)
  * Scopes supported: openid, profile, email
  * iss: "local-identity" — intentionally non-routable; production Keycloak
    rejects tokens with this issuer by design
"""

import http.server
import json
import secrets
import ssl
import sys
import time
import urllib.parse
from http import HTTPStatus
from pathlib import Path

from . import audit
from . import revoke as revoke_mod
from . import store
from .jwt_utils import JWTError, create_token, verify_token
from .keys import ensure_signing_key, jwk_public, key_id
from .security import enforce_permissions
from .tls import ensure_tls_cert

_BIND_HOST = "127.0.0.1"
_ISSUER = "local-identity"
_CODE_TTL = 60  # auth codes expire after 60 seconds


# ------------------------------------------------------------------ #
# HTML helpers                                                        #
# ------------------------------------------------------------------ #

def _he(s: str) -> str:
    """Minimal HTML attribute escaping."""
    return (
        s.replace("&", "&")
        .replace("<", "&lt;")
        .replace(">", "&gt;")
        .replace('"', "&quot;")
    )


# ------------------------------------------------------------------ #
# Handler factory                                                     #
# ------------------------------------------------------------------ #

def make_handler(
    private_key,
    kid: str,
    token_ttl: int = 3600,
    scheme: str = "https",
):
    """
    Return a configured OIDCHandler subclass.

    Each call produces a fresh class with its own isolated auth-code store,
    so multiple test servers can run concurrently without sharing state.
    """
    codes: dict = {}
    jwks_response = {"keys": [jwk_public(private_key)]}

    class _Handler(OIDCHandler):
        _private_key = private_key
        _kid = kid
        _token_ttl = token_ttl
        _codes = codes
        _scheme = scheme
        _jwks = jwks_response

    return _Handler


# ------------------------------------------------------------------ #
# HTTP handler                                                        #
# ------------------------------------------------------------------ #

class OIDCHandler(http.server.BaseHTTPRequestHandler):
    # Configured by make_handler() / run_server() via subclass class vars
    _private_key = None
    _kid: str = ""
    _token_ttl: int = 3600
    _codes: dict = {}
    _scheme: str = "https"
    _jwks: dict = {}

    def log_message(self, fmt: str, *args) -> None:
        pass  # silence default Apache-style logging

    # ---------------------------------------------------------------- #
    # Routing                                                           #
    # ---------------------------------------------------------------- #

    def do_GET(self) -> None:
        parsed = urllib.parse.urlparse(self.path)
        path = parsed.path
        if path == "/.well-known/openid-configuration":
            self._handle_discovery()
        elif path == "/jwks":
            self._handle_jwks()
        elif path == "/auth":
            self._handle_auth_get(parsed)
        elif path == "/userinfo":
            self._handle_userinfo()
        else:
            self._send_json({"error": "not_found"}, HTTPStatus.NOT_FOUND)

    def do_POST(self) -> None:
        path = urllib.parse.urlparse(self.path).path
        if path == "/auth":
            self._handle_auth_post()
        elif path == "/token":
            self._handle_token()
        else:
            self._send_json({"error": "not_found"}, HTTPStatus.NOT_FOUND)

    # ---------------------------------------------------------------- #
    # Endpoint: OIDC discovery                                          #
    # ---------------------------------------------------------------- #

    def _handle_discovery(self) -> None:
        base = self._base_url()
        doc = {
            "issuer": _ISSUER,
            "authorization_endpoint": f"{base}/auth",
            "token_endpoint": f"{base}/token",
            "userinfo_endpoint": f"{base}/userinfo",
            "jwks_uri": f"{base}/jwks",
            "response_types_supported": ["code"],
            "subject_types_supported": ["public"],
            "id_token_signing_alg_values_supported": ["RS256"],
            "scopes_supported": ["openid", "profile", "email"],
            "token_endpoint_auth_methods_supported": ["none"],
            "grant_types_supported": ["authorization_code"],
            "claims_supported": [
                "sub", "iss", "aud", "exp", "iat",
                "email", "name", "preferred_username",
            ],
        }
        self._send_json(doc)

    # ---------------------------------------------------------------- #
    # Endpoint: JWKS                                                   #
    # ---------------------------------------------------------------- #

    def _handle_jwks(self) -> None:
        self._send_json(self._jwks)

    # ---------------------------------------------------------------- #
    # Endpoint: GET /auth — display login form                         #
    # ---------------------------------------------------------------- #

    def _handle_auth_get(self, parsed: urllib.parse.ParseResult) -> None:
        params = urllib.parse.parse_qs(parsed.query, keep_blank_values=True)

        if params.get("response_type", [""])[0] != "code":
            self._send_error_redirect(
                params.get("redirect_uri", [""])[0],
                "unsupported_response_type",
                params.get("state", [""])[0],
            )
            return

        users = store.list_users()
        if not users:
            self._send_html(
                "<h1>Error</h1><p>No users in store. Run "
                "<code>local-identity init</code> first.</p>",
                status=HTTPStatus.INTERNAL_SERVER_ERROR,
            )
            return

        client_id = params.get("client_id", [""])[0]
        redirect_uri = params.get("redirect_uri", [""])[0]
        state = params.get("state", [""])[0]
        nonce = params.get("nonce", [""])[0]

        options = "\n".join(
            f'<option value="{_he(u.username)}">'
            f'{_he(u.username)} — {_he(u.fullname)}'
            f'{"  [test]" if u.generated else ""}'
            "</option>"
            for u in users
        )

        form = f"""<!DOCTYPE html>
<html lang="en">
<head>
  <meta charset="utf-8">
  <title>local-identity login</title>
  <style>
    body {{font-family: sans-serif; max-width: 480px; margin: 4em auto; color: #222;}}
    h1 {{color: #333;}}
    .warn {{background:#fff3cd;border:1px solid #ffc107;padding:.75em;border-radius:4px;}}
    select, button {{display:block;margin:.5em 0;padding:.4em;font-size:1em;}}
    button {{background:#0066cc;color:#fff;border:none;border-radius:4px;
             padding:.5em 1.5em;cursor:pointer;}}
  </style>
</head>
<body>
  <h1>local-identity dev login</h1>
  <p class="warn"><strong>Dev/test only.</strong> Tokens carry
     <code>iss: local-identity</code> — production systems reject them.</p>
  <p>Authenticating for <strong>{_he(client_id) or "(unknown client)"}</strong></p>
  <form method="POST" action="/auth">
    <input type="hidden" name="client_id"    value="{_he(client_id)}">
    <input type="hidden" name="redirect_uri" value="{_he(redirect_uri)}">
    <input type="hidden" name="state"        value="{_he(state)}">
    <input type="hidden" name="nonce"        value="{_he(nonce)}">
    <label for="username">Login as:</label>
    <select name="username" id="username">
      {options}
    </select>
    <button type="submit">Login</button>
  </form>
</body>
</html>"""
        self._send_html(form)

    # ---------------------------------------------------------------- #
    # Endpoint: POST /auth — process login form                        #
    # ---------------------------------------------------------------- #

    def _handle_auth_post(self) -> None:
        body = self._read_form_body()
        username = body.get("username", [""])[0]
        client_id = body.get("client_id", [""])[0]
        redirect_uri = body.get("redirect_uri", [""])[0]
        state = body.get("state", [""])[0]
        nonce = body.get("nonce", [""])[0]

        if not redirect_uri:
            self._send_json(
                {"error": "invalid_request", "error_description": "missing redirect_uri"},
                HTTPStatus.BAD_REQUEST,
            )
            return

        try:
            store.read_user(username)
        except FileNotFoundError:
            self._send_error_redirect(redirect_uri, "access_denied", state)
            return

        # Expire stale codes on each write
        now = time.time()
        stale = [c for c, v in self._codes.items() if v["expires_at"] < now]
        for c in stale:
            del self._codes[c]

        code = secrets.token_urlsafe(32)
        self._codes[code] = {
            "username": username,
            "client_id": client_id,
            "redirect_uri": redirect_uri,
            "nonce": nonce or None,
            "expires_at": now + _CODE_TTL,
        }
        audit.log_event("serve/auth", username, "auth_request")

        sep = "&" if "?" in redirect_uri else "?"
        location = (
            f"{redirect_uri}{sep}code={urllib.parse.quote(code)}"
            f"&state={urllib.parse.quote(state)}"
        )
        self.send_response(HTTPStatus.FOUND)
        self.send_header("Location", location)
        self.end_headers()

    # ---------------------------------------------------------------- #
    # Endpoint: POST /token                                            #
    # ---------------------------------------------------------------- #

    def _handle_token(self) -> None:
        body = self._read_form_body()
        grant_type = body.get("grant_type", [""])[0]

        if grant_type != "authorization_code":
            self._send_json(
                {"error": "unsupported_grant_type"},
                HTTPStatus.BAD_REQUEST,
            )
            return

        code = body.get("code", [""])[0]
        client_id = body.get("client_id", [""])[0]

        code_data = self._codes.pop(code, None)
        if code_data is None:
            self._send_json(
                {"error": "invalid_grant", "error_description": "unknown or expired code"},
                HTTPStatus.BAD_REQUEST,
            )
            return

        if time.time() > code_data["expires_at"]:
            self._send_json(
                {"error": "invalid_grant", "error_description": "code expired"},
                HTTPStatus.BAD_REQUEST,
            )
            return

        try:
            user = store.read_user(code_data["username"])
        except FileNotFoundError:
            self._send_json(
                {"error": "invalid_grant", "error_description": "user not found"},
                HTTPStatus.BAD_REQUEST,
            )
            return

        token = create_token(
            private_key=self._private_key,
            kid=self._kid,
            sub=user.username,
            iss=_ISSUER,
            aud=code_data.get("client_id") or client_id or "local-identity",
            email=user.email,
            name=user.fullname,
            preferred_username=user.username,
            ttl=self._token_ttl,
            nonce=code_data.get("nonce"),
        )

        audit.log_event("serve/token", user.username, "token_issued")
        self._send_json({
            "access_token": token,
            "id_token": token,
            "token_type": "Bearer",
            "expires_in": self._token_ttl,
        })

    # ---------------------------------------------------------------- #
    # Endpoint: GET /userinfo                                          #
    # ---------------------------------------------------------------- #

    def _handle_userinfo(self) -> None:
        auth_header = self.headers.get("Authorization", "")
        if not auth_header.startswith("Bearer "):
            self.send_response(HTTPStatus.UNAUTHORIZED)
            self.send_header("WWW-Authenticate", 'Bearer realm="local-identity"')
            self.end_headers()
            return

        token = auth_header[len("Bearer "):]
        try:
            payload = verify_token(token, self._private_key.public_key())
        except JWTError as exc:
            audit.log_event("serve/userinfo", None, f"invalid_token: {exc}")
            self._send_json(
                {"error": "invalid_token", "error_description": str(exc)},
                HTTPStatus.UNAUTHORIZED,
            )
            return

        jti = payload.get("jti", "")
        sub = payload.get("sub")
        if jti and revoke_mod.is_revoked(jti):
            audit.log_event("serve/userinfo", sub, "revoked_token")
            self._send_json(
                {"error": "invalid_token", "error_description": "token has been revoked"},
                HTTPStatus.UNAUTHORIZED,
            )
            return

        audit.log_event("serve/userinfo", sub, "ok")
        self._send_json({
            "sub": sub,
            "email": payload.get("email"),
            "name": payload.get("name"),
            "preferred_username": payload.get("preferred_username"),
        })

    # ---------------------------------------------------------------- #
    # Helpers                                                          #
    # ---------------------------------------------------------------- #

    def _base_url(self) -> str:
        host, port = self.server.server_address
        return f"{self._scheme}://{host}:{port}"

    def _read_form_body(self) -> dict:
        length = int(self.headers.get("Content-Length", 0))
        raw = self.rfile.read(length).decode("utf-8")
        return urllib.parse.parse_qs(raw, keep_blank_values=True)

    def _send_json(self, data: dict, status: HTTPStatus = HTTPStatus.OK) -> None:
        body = json.dumps(data).encode("utf-8")
        self.send_response(status)
        self.send_header("Content-Type", "application/json")
        self.send_header("Content-Length", str(len(body)))
        self.end_headers()
        self.wfile.write(body)

    def _send_html(self, html: str, status: HTTPStatus = HTTPStatus.OK) -> None:
        body = html.encode("utf-8")
        self.send_response(status)
        self.send_header("Content-Type", "text/html; charset=utf-8")
        self.send_header("Content-Length", str(len(body)))
        self.end_headers()
        self.wfile.write(body)

    def _send_error_redirect(
        self, redirect_uri: str, error: str, state: str
    ) -> None:
        if not redirect_uri:
            self._send_json({"error": error}, HTTPStatus.BAD_REQUEST)
            return
        sep = "&" if "?" in redirect_uri else "?"
        location = (
            f"{redirect_uri}{sep}error={urllib.parse.quote(error)}"
            f"&state={urllib.parse.quote(state)}"
        )
        self.send_response(HTTPStatus.FOUND)
        self.send_header("Location", location)
        self.end_headers()


# ------------------------------------------------------------------ #
# Public entry point                                                  #
# ------------------------------------------------------------------ #

def run_server(port: int = 8443, token_ttl: int = 3600) -> None:
    """
    Start the local-identity OIDC server on 127.0.0.1:{port} with TLS.
    Blocks until Ctrl+C.

    The server always binds to 127.0.0.1.  External binding (0.0.0.0) is
    not supported and is not offered as an option.
    """
    if not store.store_exists():
        print(
            "Error: store not initialised. Run 'local-identity init' first.",
            file=sys.stderr,
        )
        sys.exit(1)

    enforce_permissions()
    private_key = ensure_signing_key()
    cert_path, key_path = ensure_tls_cert()

    HandlerClass = make_handler(private_key, key_id(private_key), token_ttl, scheme="https")

    httpd = http.server.HTTPServer((_BIND_HOST, port), HandlerClass)

    ctx = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER)
    ctx.load_cert_chain(str(cert_path), str(key_path))
    httpd.socket = ctx.wrap_socket(httpd.socket, server_side=True)

    base_url = f"https://{_BIND_HOST}:{port}"
    print(f"local-identity OIDC server")
    print(f"  URL:       {base_url}")
    print(f"  Discovery: {base_url}/.well-known/openid-configuration")
    print(f"  Token TTL: {token_ttl}s")
    print(f"  Bound to:  {_BIND_HOST} (localhost only)")
    print("Press Ctrl+C to stop.")

    try:
        httpd.serve_forever()
    except KeyboardInterrupt:
        pass
    finally:
        httpd.server_close()