aios/active-runtime/app/aios_task_identity_overlay_api.py

#!/usr/bin/env python3
"""Loopback-only HTTP API for the task identity overlay.

Endpoints:
  GET  /healthz
  POST /api/governance/owner-authorizations

The API authenticates only the short-lived overlay credential. It does not read
or modify the global bearer-token registry.
"""
from __future__ import annotations

import json
import os
import urllib.parse
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from typing import Any

from authorized_task_writer_provisioning import ProvisioningError
from owner_authorization_broker import OwnerAuthorizationError
from task_identity_overlay import (
    MAX_BODY_BYTES,
    IdentityOverlayConfig,
    IdentityOverlayError,
    TaskIdentityOverlayService,
)


_ENDPOINT = "/api/governance/owner-authorizations"


def _enabled() -> bool:
    return os.getenv(
        "AIOS_GITOPS_IDENTITY_OVERLAY_ENABLED", ""
    ).strip().lower() in {"1", "true", "yes", "on"}


def _bearer(header: str) -> str:
    prefix = "Bearer "
    if not header.startswith(prefix):
        raise IdentityOverlayError("OVERLAY_AUTHENTICATION_REQUIRED", 401)
    token = header[len(prefix):].strip()
    if not token or len(token) > 4096:
        raise IdentityOverlayError("OVERLAY_AUTHENTICATION_REQUIRED", 401)
    return token


class OverlayHandler(BaseHTTPRequestHandler):
    server_version = "aios-task-identity-overlay/1"
    sys_version = ""

    @property
    def overlay(self) -> TaskIdentityOverlayService:
        return self.server.overlay  # type: ignore[attr-defined]

    def log_message(self, format: str, *args: Any) -> None:
        # Never log Authorization or request bodies.
        return

    def _send(self, status: int, value: dict[str, Any]) -> None:
        payload = (
            json.dumps(
                value,
                ensure_ascii=False,
                sort_keys=True,
                separators=(",", ":"),
            )
            + "\n"
        ).encode("utf-8")
        self.send_response(status)
        self.send_header("Content-Type", "application/json")
        self.send_header("Cache-Control", "no-store")
        self.send_header("Content-Length", str(len(payload)))
        self.end_headers()
        self.wfile.write(payload)

    def _body(self) -> dict[str, Any]:
        try:
            length = int(self.headers.get("Content-Length", "0"))
        except ValueError:
            raise IdentityOverlayError("OVERLAY_CONTENT_LENGTH_INVALID") from None
        if length <= 0 or length > MAX_BODY_BYTES:
            raise IdentityOverlayError("OVERLAY_BODY_SIZE_REJECTED", 413)
        raw = self.rfile.read(length)
        try:
            value = json.loads(raw)
        except (UnicodeDecodeError, json.JSONDecodeError):
            raise IdentityOverlayError("OVERLAY_BODY_JSON_INVALID") from None
        if not isinstance(value, dict):
            raise IdentityOverlayError("OVERLAY_BODY_OBJECT_REQUIRED")
        return value

    def do_GET(self) -> None:
        path = urllib.parse.urlparse(self.path).path.rstrip("/") or "/"
        if path != "/healthz":
            self._send(404, {"error": "NOT_FOUND"})
            return
        self._send(
            200,
            {
                "service": "aios-task-identity-overlay",
                "status": "active" if _enabled() else "disabled",
                "global_token_registry_used": False,
                "loopback_only": True,
            },
        )

    def do_POST(self) -> None:
        path = urllib.parse.urlparse(self.path).path.rstrip("/") or "/"
        if path != _ENDPOINT:
            self._send(404, {"error": "NOT_FOUND"})
            return
        if not _enabled():
            self._send(503, {"error": "IDENTITY_OVERLAY_DISABLED"})
            return
        try:
            token = _bearer(self.headers.get("Authorization", ""))
            result = self.overlay.provision(self._body(), token)
            self._send(200, result)
        except (IdentityOverlayError, OwnerAuthorizationError, ProvisioningError) as exc:
            body: dict[str, Any] = {"error": exc.code}
            if getattr(exc, "details", None):
                body["details"] = exc.details
            self._send(exc.status, body)
        except Exception:
            self._send(500, {"error": "IDENTITY_OVERLAY_API_INTERNAL_ERROR"})


class OverlayServer(ThreadingHTTPServer):
    allow_reuse_address = True

    def __init__(
        self,
        address: tuple[str, int],
        overlay: TaskIdentityOverlayService,
    ) -> None:
        super().__init__(address, OverlayHandler)
        self.overlay = overlay


def main() -> None:
    host = os.getenv("AIOS_GITOPS_IDENTITY_OVERLAY_HOST", "127.0.0.1")
    if host not in {"127.0.0.1", "::1", "localhost"}:
        raise SystemExit("identity overlay must listen on loopback")
    try:
        port = int(os.getenv("AIOS_GITOPS_IDENTITY_OVERLAY_PORT", "19864"))
    except ValueError:
        raise SystemExit("invalid overlay port") from None
    if not 1024 <= port <= 65535:
        raise SystemExit("invalid overlay port")
    config = IdentityOverlayConfig.from_environment()
    server = OverlayServer(
        (host, port),
        TaskIdentityOverlayService(config),
    )
    server.serve_forever(poll_interval=0.25)


if __name__ == "__main__":
    main()