atlas-iac/services/hermes/scripts/telegram_media_server.py

#!/usr/bin/env python3
"""Serve one Hermes tenant's generated media to the Telegram router."""

from __future__ import annotations

import hmac
import mimetypes
import os
from pathlib import Path, PurePosixPath
import posixpath
import shutil
import stat
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from urllib.parse import parse_qs, urlsplit


ALLOWED_ROOTS = (Path("/opt/data/cache/images"), Path("/opt/data/workspace"))
MAX_MEDIA_BYTES = 50 << 20


class MediaPathError(ValueError):
    """Raised when a requested path violates the tenant media boundary."""


def read_relay_key(path: Path) -> str:
    """Read the shared relay key from one runtime-only secret file."""
    value = path.read_text(encoding="utf-8").strip()
    if not value:
        raise RuntimeError("tenant relay key is unavailable")
    return value


def tenant_slot(pod_name: str) -> str:
    """Extract the StatefulSet ordinal used to bind requests to this tenant."""
    prefix, separator, ordinal = pod_name.rpartition("-")
    if not separator or not prefix or not ordinal.isdigit():
        raise RuntimeError("tenant pod name has no ordinal")
    return ordinal


def normalize_media_path(raw_path: str, allowed_roots: tuple[Path, ...] = ALLOWED_ROOTS) -> tuple[Path, tuple[str, ...]]:
    """Return an allowed root and clean relative components for an absolute path."""
    value = str(raw_path or "").strip()
    if not value or "\x00" in value or "\\" in value or not value.startswith("/"):
        raise MediaPathError("invalid media path")
    original = PurePosixPath(value)
    if ".." in original.parts:
        raise MediaPathError("parent traversal is not allowed")
    cleaned = PurePosixPath(posixpath.normpath(value))
    for root in allowed_roots:
        logical_root = PurePosixPath(root.as_posix())
        try:
            relative = cleaned.relative_to(logical_root)
        except ValueError:
            continue
        components = tuple(part for part in relative.parts if part not in {"", "."})
        if not components:
            raise MediaPathError("media path must name a file")
        return root, components
    raise MediaPathError("media path is outside allowed roots")


def open_directory_without_symlinks(path: Path) -> int:
    """Open every component of an absolute directory path with O_NOFOLLOW."""
    if not path.is_absolute() or ".." in path.parts:
        raise MediaPathError("invalid media root")
    flags = os.O_RDONLY | os.O_DIRECTORY | os.O_CLOEXEC | os.O_NOFOLLOW
    directory_fd = os.open("/", flags)
    try:
        for component in path.parts[1:]:
            next_fd = os.open(component, flags, dir_fd=directory_fd)
            os.close(directory_fd)
            directory_fd = next_fd
        return directory_fd
    except Exception:
        os.close(directory_fd)
        raise


def open_media_file(raw_path: str, allowed_roots: tuple[Path, ...] = ALLOWED_ROOTS) -> tuple[int, os.stat_result, str]:
    """Open a regular file beneath an allowed root without following symlinks."""
    root, components = normalize_media_path(raw_path, allowed_roots)
    directory_flags = os.O_RDONLY | os.O_DIRECTORY | os.O_CLOEXEC | os.O_NOFOLLOW
    file_flags = os.O_RDONLY | os.O_CLOEXEC | os.O_NOFOLLOW
    directory_fd = open_directory_without_symlinks(root)
    try:
        for component in components[:-1]:
            next_fd = os.open(component, directory_flags, dir_fd=directory_fd)
            os.close(directory_fd)
            directory_fd = next_fd
        file_fd = os.open(components[-1], file_flags, dir_fd=directory_fd)
    finally:
        os.close(directory_fd)
    metadata = os.fstat(file_fd)
    if not stat.S_ISREG(metadata.st_mode) or metadata.st_size > MAX_MEDIA_BYTES:
        os.close(file_fd)
        raise MediaPathError("media is not an uploadable regular file")
    return file_fd, metadata, components[-1]


class TenantMediaServer(ThreadingHTTPServer):
    """HTTP server carrying immutable tenant authentication and root policy."""

    daemon_threads = True

    def __init__(
        self,
        address: tuple[str, int],
        relay_key: str,
        slot: str,
        allowed_roots: tuple[Path, ...] = ALLOWED_ROOTS,
    ) -> None:
        self.relay_key = relay_key
        self.tenant_slot = slot
        self.allowed_roots = allowed_roots
        super().__init__(address, TenantMediaHandler)


class TenantMediaHandler(BaseHTTPRequestHandler):
    """Authenticate the router and stream one tenant-private file descriptor."""

    server: TenantMediaServer

    def log_message(self, _format: str, *_args: object) -> None:
        """Keep request paths out of logs because they contain internal names."""

    def _plain_status(self, status_code: int, message: bytes) -> None:
        self.send_response(status_code)
        self.send_header("Content-Type", "text/plain; charset=utf-8")
        self.send_header("Content-Length", str(len(message)))
        self.send_header("Cache-Control", "no-store")
        self.end_headers()
        self.wfile.write(message)

    def _authorized(self) -> bool:
        expected = "Bearer " + self.server.relay_key
        supplied = self.headers.get("Authorization", "")
        slot = self.headers.get("X-Hermes-Tenant-Slot", "")
        return hmac.compare_digest(supplied.encode(), expected.encode()) and hmac.compare_digest(
            slot.encode(), self.server.tenant_slot.encode()
        )

    def do_GET(self) -> None:  # noqa: N802 - BaseHTTPRequestHandler API
        parsed = urlsplit(self.path)
        if parsed.path == "/healthz":
            self._plain_status(200, b"ok\n")
            return
        if parsed.path != "/media":
            self._plain_status(404, b"not found\n")
            return
        if not self._authorized():
            self._plain_status(401, b"authentication required\n")
            return
        raw_path = parse_qs(parsed.query).get("path", [""])[0]
        try:
            file_fd, metadata, filename = open_media_file(raw_path, self.server.allowed_roots)
        except MediaPathError:
            self._plain_status(403, b"media path rejected\n")
            return
        except (FileNotFoundError, NotADirectoryError, PermissionError, OSError):
            self._plain_status(404, b"media unavailable\n")
            return
        content_type = mimetypes.guess_type(filename)[0] or "application/octet-stream"
        safe_filename = filename.replace('"', "_").replace("\r", "_").replace("\n", "_")
        self.send_response(200)
        self.send_header("Content-Type", content_type)
        self.send_header("Content-Length", str(metadata.st_size))
        self.send_header("Content-Disposition", f'attachment; filename="{safe_filename}"')
        self.send_header("Cache-Control", "no-store")
        self.send_header("X-Content-Type-Options", "nosniff")
        self.end_headers()
        try:
            with os.fdopen(file_fd, "rb") as media:
                shutil.copyfileobj(media, self.wfile, length=64 << 10)
        except (BrokenPipeError, ConnectionResetError):
            return


def main() -> None:
    """Start the tenant-local media handoff server."""
    relay_key_path = Path(
        os.getenv(
            "HERMES_MEDIA_RELAY_KEY_FILE",
            "/runtime-access/chat-relay-key",
        )
    )
    relay_key = read_relay_key(relay_key_path)
    slot = tenant_slot(os.getenv("POD_NAME", ""))
    server = TenantMediaServer(("0.0.0.0", 8788), relay_key, slot)
    server.serve_forever()


if __name__ == "__main__":
    main()