atlas-iac/dockerfiles/hermes-webui-hux-bff-patch.py

#!/usr/bin/env python3
"""Install the pinned, same-origin HUX BFF into Hermes WebUI."""

from __future__ import annotations

import hashlib
import hmac
import http.client
import json
import os
from pathlib import Path
import re
import socket
import stat
import time
from urllib.parse import parse_qsl, urlencode


PINNED_UPSTREAM_COMMIT = "7a94e34a6d639576576baa9131acf6765f6d2b98"
UPSTREAM_HOST = "127.0.0.1"
UPSTREAM_PORT = 8790
RELAY_KEY_FILE = Path("/run/hermes-webui-hux/relay-key")
CONTEXT_KEY_FILE = Path("/run/hermes-hux-context/context-key")
MAX_REQUEST_BYTES = 1024 * 1024
MAX_RESPONSE_BYTES = 4 * 1024 * 1024
MAX_ARTIFACT_REQUEST_BYTES = 34 * 1024 * 1024
MAX_ARTIFACT_RESPONSE_BYTES = 52 * 1024 * 1024
MAX_STREAM_BYTES = 16 * 1024 * 1024
UPSTREAM_TIMEOUT = 10.0
STREAM_DEADLINE = 65.0
CHUNK_BYTES = 64 * 1024
SAFE_PATH = re.compile(r"^/hux/v1(?:/[A-Za-z0-9._:-]+)*$")
SAFE_QUERY_KEY = re.compile(r"^[A-Za-z0-9._:@+-]{1,120}$")
SAFE_QUERY_VALUE = re.compile(r"^[A-Za-z0-9 ._:@+,-]{0,240}$")
SAFE_SLOT = re.compile(r"^slot-[0-9]{1,3}$")
SAFE_SUBJECT = re.compile(r"^usr_[0-9a-f]{64}$")
SAFE_KEY = re.compile(r"^[A-Za-z0-9._:-]{32,256}$")
SAFE_IDEMPOTENCY = re.compile(r"^[A-Za-z0-9._:-]{8,120}$")
SAFE_REVISION = re.compile(r"^[0-9]{1,10}$")
SAFE_EVENT_ID = re.compile(r"^[0-9]{1,20}$")
MUTATIONS = frozenset({"POST", "PUT", "PATCH", "DELETE"})
CONTEXT_KEY_BYTES = 32


class BffError(Exception):
    """A bounded, user-safe proxy rejection."""

    def __init__(self, status: int, message: str):
        super().__init__(message)
        self.status = status


def _header_values(handler, name: str) -> list[str]:
    headers = getattr(handler, "headers", None)
    if headers is None:
        return []
    if hasattr(headers, "get_all"):
        return [str(value) for value in (headers.get_all(name) or [])]
    value = headers.get(name)
    return [] if value is None else [str(value)]


def _one_header(handler, name: str) -> str:
    values = _header_values(handler, name)
    if len(values) > 1:
        raise BffError(400, f"Duplicate {name} header")
    return values[0].strip() if values else ""


def _trusted_identity(handler) -> tuple[str, str, str]:
    from api.auth import ensure_trusted_auth_session, parse_cookie, verify_session

    cookie = parse_cookie(handler)
    if not cookie or not verify_session(cookie):
        raise BffError(401, "Trusted WebUI session required")
    info = ensure_trusted_auth_session(handler)
    slot = _one_header(handler, "X-Hermes-Tenant-Identity")
    supplied_subject = _one_header(handler, "X-Hux-Subject")
    username = str((info or {}).get("username") or "")
    if (not info or info.get("auth_type") != "trusted" or not SAFE_SLOT.fullmatch(slot)
            or username != slot):
        raise BffError(401, "Trusted WebUI session required")
    subject = _derived_subject(slot)
    if supplied_subject and supplied_subject != subject:
        raise BffError(401, "Trusted WebUI session required")
    return slot, subject, "chat"


def _context_key_path() -> Path:
    raw = os.environ.get("HUX_CONTEXT_KEY_FILE", "").strip()
    path = Path(raw) if raw else CONTEXT_KEY_FILE
    if not path.is_absolute() or ".." in path.parts:
        raise BffError(503, "HUX identity is unavailable")
    return path


def _secure_context_key() -> bytes:
    path = _context_key_path()
    try:
        parent = path.parent.stat(follow_symlinks=False)
        if not stat.S_ISDIR(parent.st_mode) or parent.st_uid != os.geteuid():
            raise BffError(503, "HUX identity is unavailable")
        descriptor = os.open(path, os.O_RDONLY | getattr(os, "O_CLOEXEC", 0) |
                             getattr(os, "O_NOFOLLOW", 0))
        try:
            info = os.fstat(descriptor)
            if (not stat.S_ISREG(info.st_mode) or info.st_uid != os.geteuid() or
                    stat.S_IMODE(info.st_mode) != 0o600 or info.st_nlink != 1 or
                    info.st_size != CONTEXT_KEY_BYTES):
                raise BffError(503, "HUX identity is unavailable")
            key = os.read(descriptor, CONTEXT_KEY_BYTES + 1)
        finally:
            os.close(descriptor)
    except BffError:
        raise
    except OSError as exc:
        raise BffError(503, "HUX identity is unavailable") from exc
    if len(key) != CONTEXT_KEY_BYTES:
        raise BffError(503, "HUX identity is unavailable")
    return key


def _derived_subject(slot: str) -> str:
    digest = hmac.new(_secure_context_key(), b"hux.subject.id.v1\0" + slot.encode("ascii"),
                      hashlib.sha256).hexdigest()
    subject = "usr_" + digest
    if not SAFE_SUBJECT.fullmatch(subject):  # defensive parity with the public identity contract
        raise BffError(503, "HUX identity is unavailable")
    return subject


def _require_csrf(handler, cookie: str) -> None:
    from api.auth import CSRF_HEADER_NAME, verify_csrf_token

    supplied = _one_header(handler, CSRF_HEADER_NAME)
    if not supplied or not verify_csrf_token(cookie, supplied):
        raise BffError(403, "Valid WebUI CSRF token required")


def _secure_relay_key() -> str:
    parent = RELAY_KEY_FILE.parent
    try:
        parent_stat = parent.stat(follow_symlinks=False)
    except OSError as exc:
        raise BffError(503, "HUX relay is unavailable") from exc
    if not stat.S_ISDIR(parent_stat.st_mode) or parent_stat.st_uid != os.geteuid() or parent_stat.st_mode & 0o077:
        raise BffError(503, "HUX relay is unavailable")
    flags = os.O_RDONLY | getattr(os, "O_CLOEXEC", 0) | getattr(os, "O_NOFOLLOW", 0)
    try:
        descriptor = os.open(RELAY_KEY_FILE, flags)
        try:
            file_stat = os.fstat(descriptor)
            if (not stat.S_ISREG(file_stat.st_mode) or file_stat.st_uid != os.geteuid() or
                    stat.S_IMODE(file_stat.st_mode) != 0o400 or file_stat.st_nlink != 1):
                raise BffError(503, "HUX relay is unavailable")
            raw = os.read(descriptor, 257)
        finally:
            os.close(descriptor)
    except BffError:
        raise
    except OSError as exc:
        raise BffError(503, "HUX relay is unavailable") from exc
    try:
        key = raw.decode("ascii").strip()
    except UnicodeDecodeError as exc:
        raise BffError(503, "HUX relay is unavailable") from exc
    if len(raw) > 256 or not SAFE_KEY.fullmatch(key):
        raise BffError(503, "HUX relay is unavailable")
    return key


def _safe_target(parsed) -> str:
    path = str(getattr(parsed, "path", "") or "")
    segments = path.split("/")[3:]
    if not SAFE_PATH.fullmatch(path) or any(segment in {".", ".."} for segment in segments):
        raise BffError(400, "Invalid HUX path")
    raw_query = str(getattr(parsed, "query", "") or "")
    if len(raw_query) > 4096:
        raise BffError(400, "Invalid HUX query")
    try:
        pairs = parse_qsl(raw_query, keep_blank_values=True, strict_parsing=True, max_num_fields=64)
    except ValueError as exc:
        raise BffError(400, "Invalid HUX query") from exc
    if len({key for key, _ in pairs}) != len(pairs) or any(
        not SAFE_QUERY_KEY.fullmatch(key) or not SAFE_QUERY_VALUE.fullmatch(value)
        for key, value in pairs
    ):
        raise BffError(400, "Invalid HUX query")
    query = urlencode(pairs)
    return path + ("?" + query if query else "")


def _request_limit(target: str, method: str) -> int:
    path = target.split("?", 1)[0]
    artifact_write = re.fullmatch(
        r"/hux/v1/artifacts(?:/[A-Za-z0-9._:-]+/versions)?", path
    )
    return MAX_ARTIFACT_REQUEST_BYTES if method == "POST" and artifact_write else MAX_REQUEST_BYTES


def _response_limit(target: str) -> int:
    path = target.split("?", 1)[0]
    artifact_read = re.fullmatch(
        r"/hux/v1/artifacts/[A-Za-z0-9._:-]+/versions/[0-9]+(?:/diff)?", path
    )
    return MAX_ARTIFACT_RESPONSE_BYTES if artifact_read else MAX_RESPONSE_BYTES


def _read_request(handler, method: str, limit: int = MAX_REQUEST_BYTES) -> bytes | None:
    if _one_header(handler, "Transfer-Encoding"):
        raise BffError(400, "Chunked HUX requests are not supported")
    raw_length = _one_header(handler, "Content-Length")
    if not raw_length:
        return None
    if not raw_length.isdigit():
        raise BffError(400, "Invalid Content-Length")
    length = int(raw_length)
    if length > limit:
        raise BffError(413, "HUX request body is too large")
    if method == "GET" and length:
        raise BffError(400, "GET request body is not allowed")
    if length and _one_header(handler, "Content-Type").lower() != "application/json":
        raise BffError(415, "HUX mutations require application/json")
    body = bytearray()
    while len(body) < length:
        chunk = handler.rfile.read(length - len(body))
        if not chunk:
            raise BffError(400, "Incomplete HUX request body")
        body.extend(chunk)
    return bytes(body)


def _validated_passthrough_headers(handler, method: str, body: bytes | None) -> dict[str, str]:
    headers = {"Accept": "text/event-stream" if str(handler.path).split("?", 1)[0].endswith("/stream")
               else "application/vnd.hermes.hux+json; version=1"}
    validators = (("If-Match", SAFE_REVISION), ("Idempotency-Key", SAFE_IDEMPOTENCY),
                  ("Last-Event-ID", SAFE_EVENT_ID))
    for name, pattern in validators:
        value = _one_header(handler, name)
        if value:
            if not pattern.fullmatch(value) or (name == "Last-Event-ID" and method != "GET"):
                raise BffError(400, f"Invalid {name} header")
            headers[name] = value
    if body is not None:
        headers["Content-Type"] = "application/json"
    return headers


def _send_headers(handler, status: int, content_type: str, length: int | None = None) -> None:
    handler.send_response(status)
    handler.send_header("Content-Type", content_type)
    if length is not None:
        handler.send_header("Content-Length", str(length))
    handler.send_header("Cache-Control", "no-store")
    handler.send_header("X-Content-Type-Options", "nosniff")
    handler.send_header("Referrer-Policy", "no-referrer")


def _send_error(handler, error: BffError) -> bool:
    body = json.dumps({"error": str(error)}, separators=(",", ":")).encode("utf-8")
    _send_headers(handler, error.status, "application/json", len(body))
    handler.end_headers()
    handler.wfile.write(body)
    return True


def _response_type(response) -> str:
    content_type = str(response.getheader("Content-Type") or "application/json").split(";", 1)[0].strip().lower()
    if content_type not in {"application/json", "application/vnd.hermes.hux+json", "text/event-stream"}:
        raise BffError(502, "HUX returned an unsupported response")
    return content_type


def _relay_stream(handler, response, content_type: str) -> bool:
    _send_headers(handler, response.status, content_type)
    handler.send_header("Connection", "close")
    handler.end_headers()
    handler.close_connection = True
    total = 0
    deadline = time.monotonic() + STREAM_DEADLINE
    while time.monotonic() < deadline:
        chunk = response.read(min(CHUNK_BYTES, MAX_STREAM_BYTES - total + 1))
        if not chunk:
            break
        total += len(chunk)
        if total > MAX_STREAM_BYTES:
            break
        handler.wfile.write(chunk)
        handler.wfile.flush()
    return True


def _relay_response(handler, response, limit: int = MAX_RESPONSE_BYTES) -> bool:
    if 300 <= response.status < 400:
        raise BffError(502, "HUX redirects are refused")
    content_type = _response_type(response)
    if content_type == "text/event-stream":
        return _relay_stream(handler, response, content_type)
    body = response.read(limit + 1)
    if len(body) > limit:
        raise BffError(502, "HUX response is too large")
    _send_headers(handler, response.status, content_type, len(body))
    etag = response.getheader("ETag")
    if etag and re.fullmatch(r'"?[0-9]{1,10}"?', etag):
        handler.send_header("ETag", etag)
    if str(response.getheader("HUX-Replayed") or "").lower() == "true":
        handler.send_header("HUX-Replayed", "true")
    audit_stale = str(response.getheader("HUX-Audit-Stale") or "").lower()
    if audit_stale in {"true", "false"}:
        handler.send_header("HUX-Audit-Stale", audit_stale)
    if str(response.getheader("Content-Disposition") or "").lower() == "attachment":
        handler.send_header("Content-Disposition", "attachment")
    handler.end_headers()
    handler.wfile.write(body)
    return True


def proxy_hux(handler, parsed, method: str) -> bool:
    """Proxy one authenticated WebUI request to the fixed loopback HUX service."""
    connection = None
    try:
        method = str(method or "").upper()
        if method != "GET" and method not in MUTATIONS:
            raise BffError(405, "HUX method is not allowed")
        slot, subject, surface = _trusted_identity(handler)
        from api.auth import parse_cookie
        if method in MUTATIONS:
            _require_csrf(handler, parse_cookie(handler) or "")
        target = _safe_target(parsed)
        body = _read_request(handler, method, _request_limit(target, method))
        headers = _validated_passthrough_headers(handler, method, body)
        headers.update({"X-Hermes-Tenant-Identity": slot, "X-Hux-Subject": subject,
                        "X-Hux-Surface": surface, "X-Hux-Trust": "relay",
                        "X-Hux-Relay-Key": _secure_relay_key()})
        connection = http.client.HTTPConnection(UPSTREAM_HOST, UPSTREAM_PORT, timeout=UPSTREAM_TIMEOUT)
        connection.request(method, target, body=body, headers=headers)
        return _relay_response(handler, connection.getresponse(), _response_limit(target))
    except BffError as error:
        return _send_error(handler, error)
    except (TimeoutError, socket.timeout, ConnectionError, OSError, http.client.HTTPException):
        return _send_error(handler, BffError(502, "HUX service is unavailable"))
    finally:
        if connection is not None:
            connection.close()


def _patched_routes(source: str) -> str:
    if "from api.hux_bff import proxy_hux" in source:
        raise SystemExit("Hermes HUX BFF patch is already installed")
    replacements = []
    for method in ("GET", "POST", "PUT", "PATCH", "DELETE"):
        before = (f"def handle_{method.lower()}(handler, parsed) -> bool:\n"
                  f"    \"\"\"Handle all {method} routes. Returns True if handled, False for 404.\"\"\"\n")
        after = before + ("    if parsed.path == \"/hux/v1\" or parsed.path.startswith(\"/hux/v1/\"):\n"
                          "        from api.hux_bff import proxy_hux\n"
                          f"        return proxy_hux(handler, parsed, \"{method}\")\n")
        if source.count(before) != 1:
            raise SystemExit(f"Hermes HUX BFF patch context changed: {before[:80]!r}")
        replacements.append((before, after))
    for before, after in replacements:
        source = source.replace(before, after, 1)
    return source


def apply(root: Path | None = None) -> None:
    """Install this module and patch all five pinned WebUI route dispatchers."""
    root = root or Path(os.environ.get("HERMES_WEBUI_PATCH_ROOT", "/opt/hermes-webui"))
    api_dir = root / "api"
    routes = api_dir / "routes.py"
    installed = api_dir / "hux_bff.py"
    patched = _patched_routes(routes.read_text(encoding="utf-8"))
    installed.write_text(Path(__file__).read_text(encoding="utf-8"), encoding="utf-8")
    routes.write_text(patched, encoding="utf-8")


if __name__ == "__main__":
    apply()