CIS490/receiver/app.py

from __future__ import annotations

import logging
import secrets
import time
from pathlib import Path
from typing import Awaitable, Callable

from starlette.applications import Starlette
from starlette.requests import Request
from starlette.responses import JSONResponse, Response
from starlette.routing import Route

from .store import EpisodeStore, is_valid_id


log = logging.getLogger("cis490.receiver")


SUFFIX = ".tar.zst"
SCHEMA_VERSION = 1


def _bearer_check(request: Request, expected: str | None) -> Response | None:
    if expected is None:
        return None
    auth = request.headers.get("authorization", "")
    prefix = "Bearer "
    if not auth.startswith(prefix):
        return JSONResponse({"error": "missing bearer token"}, status_code=401)
    presented = auth[len(prefix):]
    if not secrets.compare_digest(presented, expected):
        return JSONResponse({"error": "bad bearer token"}, status_code=401)
    return None


def make_app(
    store: EpisodeStore,
    max_episode_bytes: int,
    bearer_token: str | None = None,
) -> Starlette:
    async def health(request: Request) -> JSONResponse:
        return JSONResponse({"status": "ok"})

    async def ping(request: Request) -> JSONResponse:
        """Smoke-test endpoint. Verifies that the auth layer and the
        WG/Caddy/receiver pipe are alive end-to-end without persisting
        anything — index.jsonl is untouched. Used by ``cis490-shipper
        --ping`` during initial bring-up of a new lab host."""
        guard = _bearer_check(request, bearer_token)
        if guard is not None:
            return guard
        return JSONResponse(
            {
                "ok": True,
                "host_id": request.headers.get("x-lab-host"),
                "t_wall_ns": time.time_ns(),
                "schema_version": SCHEMA_VERSION,
            }
        )

    async def put_episode(request: Request) -> JSONResponse:
        guard = _bearer_check(request, bearer_token)
        if guard is not None:
            return guard

        host_id: str = request.path_params["host_id"]
        filename: str = request.path_params["filename"]

        if not filename.endswith(SUFFIX):
            return JSONResponse(
                {"error": f"expected {SUFFIX} suffix"}, status_code=400
            )
        episode_id = filename[: -len(SUFFIX)]

        if not is_valid_id(host_id) or not is_valid_id(episode_id):
            return JSONResponse({"error": "bad host_id or episode_id"}, status_code=400)

        expected_sha = request.headers.get("x-content-sha256")
        if not expected_sha or len(expected_sha) != 64:
            return JSONResponse(
                {"error": "X-Content-SHA256 header (64 hex chars) required"},
                status_code=400,
            )
        expected_sha = expected_sha.lower()

        try:
            schema_version = int(request.headers.get("x-schema-version", "1"))
        except ValueError:
            return JSONResponse({"error": "bad X-Schema-Version"}, status_code=400)

        cl = request.headers.get("content-length")
        if cl is not None:
            try:
                if int(cl) > max_episode_bytes:
                    return JSONResponse(
                        {"error": "episode exceeds max size"}, status_code=413
                    )
            except ValueError:
                return JSONResponse({"error": "bad Content-Length"}, status_code=400)

        result = await store.ingest_stream(
            host_id=host_id,
            episode_id=episode_id,
            expected_sha256=expected_sha,
            schema_version=schema_version,
            body=request.stream(),
            max_bytes=max_episode_bytes,
        )

        if result.status == "stored":
            log.info(
                "stored episode host=%s id=%s sha=%s size=%d",
                host_id, episode_id, result.sha256, result.size_bytes,
            )
            return JSONResponse(
                {"status": "stored", "sha256": result.sha256, "size_bytes": result.size_bytes},
                status_code=201,
            )
        if result.status == "already-present":
            return JSONResponse(
                {"status": "already-present", "sha256": result.sha256},
                status_code=200,
            )
        if result.status == "conflict":
            log.warning(
                "conflict host=%s id=%s existing=%s",
                host_id, episode_id, result.existing_sha256,
            )
            return JSONResponse(
                {"status": "conflict", "existing_sha256": result.existing_sha256},
                status_code=409,
            )
        if result.status == "sha-mismatch":
            return JSONResponse(
                {"status": "sha-mismatch", "actual_sha256": result.sha256},
                status_code=400,
            )
        if result.status == "too-large":
            return JSONResponse({"error": "episode exceeds max size"}, status_code=413)

        return JSONResponse({"error": "unknown ingest result"}, status_code=500)

    routes = [
        Route("/v1/health", health, methods=["GET"]),
        Route("/v1/ping", ping, methods=["POST"]),
        Route(
            "/v1/episodes/{host_id}/{filename}",
            put_episode,
            methods=["PUT"],
        ),
    ]
    return Starlette(routes=routes)