#!/usr/bin/env python3
"""Serve the credential-isolated Atlas API and Git smart-HTTP boundary."""

from __future__ import annotations

import argparse
import base64
import json
import re
import socket
import tempfile
import time
import urllib.error
import urllib.parse
import urllib.request
from http.server import BaseHTTPRequestHandler
from typing import BinaryIO

import deadline_http
from gitea_api import (
    CANONICAL_BASE_URL,
    PolicyError,
    create_draft,
    read,
    read_token,
)
from gitea_api_policy import _reject_forbidden, _validate_repo, _validate_sha
from receive_pack_scan import validate_receive_pack
from scm_task_grants import TaskLedger, verify_grant
from scm_task_drafts import matches_pull, request_fields, update as update_draft
from scm_task_adoptions import seed as seed_adoptions
from scm_broker_io import RejectRedirect, response_status as _status, spool_response
from scm_broker_server import AbsoluteHeaderDeadlineMixin, BoundedThreadingHTTPServer, log_rejection, validate_ascii_headers

BROKER_PORT = 9081
GIT_USER = "hermes-automation"
MAX_CONTROL_BODY = 64 * 1024
MAX_GIT_REQUEST = 128 * 1024 * 1024
MAX_GIT_RESPONSE = 128 * 1024 * 1024
MAX_HEADERS = 32
MAX_HEADER_BYTES = 16 * 1024
SPOOL_MEMORY_LIMIT = 1024 * 1024
STREAM_CHUNK = 64 * 1024
INBOUND_HEADER_TIMEOUT = 10.0
INBOUND_BODY_TIMEOUT = 120.0
UPSTREAM_DEADLINE_SECONDS = 300.0
GIT_PATH_RE = re.compile(
    r"/git/atlas/(?P<repo>[A-Za-z0-9][A-Za-z0-9._-]{0,99})\.git/"
    r"(?P<operation>info/refs|git-upload-pack|git-receive-pack)\Z"
)
def _read_bounded(
    stream,
    maximum: int,
    length: int | None = None,
    *,
    deadline: float | None = None,
    set_timeout=None,
) -> bytes:
    if length is not None and not 0 <= length <= maximum:
        raise PolicyError("SCM request exceeds the safe size limit")
    wanted = maximum + 1 if length is None else length
    chunks: list[bytes] = []
    total = 0
    while total < wanted:
        if deadline is not None:
            remaining = deadline - time.monotonic()
            if remaining <= 0:
                raise PolicyError("SCM request body deadline exceeded")
            if set_timeout is not None:
                set_timeout(remaining)
        chunk = stream.read(min(STREAM_CHUNK, wanted - total))
        if not chunk:
            break
        chunks.append(chunk)
        total += len(chunk)
    value = b"".join(chunks)
    if len(value) > maximum or (length is not None and len(value) != length):
        raise PolicyError("SCM request exceeds the safe size limit")
    return value
def _spool_bounded(
    stream,
    maximum: int,
    length: int,
    *,
    token: str,
    context: str,
    deadline: float,
    set_timeout=None,
) -> tuple[BinaryIO, int]:
    """Copy a fixed-length exchange through a bounded-memory disk spool."""
    if not 0 <= length <= maximum:
        raise PolicyError("SCM request exceeds the safe size limit")
    spool = tempfile.SpooledTemporaryFile(  # noqa: SIM115 - caller owns returned spool
        max_size=SPOOL_MEMORY_LIMIT, dir="/tmp"
    )
    remaining = length
    forms = _credential_forms(token)
    carry = b""
    try:
        while remaining:
            remaining_time = deadline - time.monotonic()
            if remaining_time <= 0:
                raise PolicyError("SCM request body deadline exceeded")
            if set_timeout is not None:
                set_timeout(remaining_time)
            chunk = stream.read(min(STREAM_CHUNK, remaining))
            if not chunk:
                raise PolicyError("SCM request body ended early")
            candidate = carry + chunk
            if any(form in candidate for form in forms):
                raise PolicyError(f"{context} contains runtime credential material")
            width = max(len(form) for form in forms) - 1
            carry = candidate[-width:] if width else b""
            spool.write(chunk)
            remaining -= len(chunk)
        spool.seek(0)
        return spool, length
    except Exception:
        spool.close()
        raise
def _spool_response(
    response, maximum: int, token: str, deadline_seconds: float
) -> tuple[BinaryIO, int]:
    return spool_response(
        response,
        maximum,
        _credential_forms(token),
        memory_limit=SPOOL_MEMORY_LIMIT,
        chunk_size=STREAM_CHUNK,
        deadline_seconds=deadline_seconds,
    )
def _load_json(handler: BaseHTTPRequestHandler) -> dict[str, object]:
    if handler.headers.get("Transfer-Encoding"):
        raise PolicyError("chunked broker control requests are not allowed")
    if handler.headers.get_content_type() != "application/json":
        raise PolicyError("broker control request must be JSON")
    length = _content_length(handler.headers, MAX_CONTROL_BODY)
    deadline = time.monotonic() + INBOUND_BODY_TIMEOUT
    value = json.loads(
        _read_bounded(
            handler.rfile,
            MAX_CONTROL_BODY,
            length,
            deadline=deadline,
            set_timeout=handler.connection.settimeout,
        )
    )
    if not isinstance(value, dict):
        raise PolicyError("broker control request must be an object")
    return value
def _git_target(raw: str) -> tuple[str, str, str]:
    if (
        not raw.isascii()
        or any(ord(character) < 32 or ord(character) == 127 for character in raw)
        or "%" in raw
        or "\\" in raw
        or len(raw) > 512
    ):
        raise PolicyError("Git target is not canonical ASCII")
    target = urllib.parse.urlsplit(raw)
    canonical = urllib.parse.urlunsplit(("", "", target.path, target.query, ""))
    if canonical != raw:
        raise PolicyError("Git target is not in exact canonical form")
    if target.scheme or target.netloc or target.fragment:
        raise PolicyError("Git target must be relative")
    match = GIT_PATH_RE.fullmatch(target.path)
    if not match:
        raise PolicyError("Git target is outside the Atlas allowlist")
    repo = _validate_repo(match.group("repo"))
    operation = match.group("operation")
    query = target.query
    if operation == "info/refs":
        if query not in {"service=git-upload-pack", "service=git-receive-pack"}:
            raise PolicyError("Git discovery service is outside the allowlist")
        service = query.removeprefix("service=")
    elif query:
        raise PolicyError("Git RPC query parameters are not allowed")
    else:
        service = operation
    return repo, operation, service
def _content_length(headers: object, maximum: int) -> int:
    """Parse one short canonical bounded Content-Length header."""
    raw = headers.get("Content-Length", "")  # type: ignore[attr-defined]
    if (
        not isinstance(raw, str)
        or not raw.isascii()
        or len(raw) > 10
        or not raw.isdigit()
    ):
        raise PolicyError("SCM request length is invalid")
    length = int(raw)
    if raw != str(length) or not 0 <= length <= maximum:
        raise PolicyError("SCM request length is invalid")
    return length
def _credential_forms(token: str) -> tuple[bytes, ...]:
    basic = base64.b64encode(f"{GIT_USER}:{token}".encode())
    return token.encode("utf-8"), basic, b"Basic " + basic
def _validate_receive_pack(body: bytes | BinaryIO, token: str) -> None:
    """Permit only vetted new-feature-branch pushes with scanned contents."""
    validate_receive_pack(body, token, _credential_forms(token))
def _task_ledger() -> TaskLedger:
    """Open the PVC ledger lazily so read-only broker paths need no state."""
    return TaskLedger()
def _branch_head(repo: str, ref: str, token: str) -> str | None:
    """Read one exact remote branch head through the existing API boundary."""
    try:
        raw = read(f"/api/v1/repos/titan/{repo}/branches/{ref}", token=token)
        value = json.loads(raw)
        commit = value.get("commit") if isinstance(value, dict) else None
        return _validate_sha(commit.get("id") if isinstance(commit, dict) else None)
    except urllib.error.HTTPError as error:
        if error.code == 404:
            return None
        raise
def _guarded_opener(guard: deadline_http.StreamDeadline):
    """Build an opener whose connection obeys one absolute stream deadline."""
    return urllib.request.build_opener(RejectRedirect(), *guard.handlers()).open
def _upstream_git_request(
    target: str,
    *,
    method: str,
    body: bytes | BinaryIO | None,
    content_type: str | None,
    expected_type: str,
    token: str,
    body_length: int | None = None,
    stream_result: bool = False,
    opener=None,
) -> bytes | tuple[BinaryIO, int]:
    credentials = base64.b64encode(f"{GIT_USER}:{token}".encode()).decode("ascii")
    headers = {"Accept": expected_type, "Authorization": f"Basic {credentials}"}
    if content_type:
        headers["Content-Type"] = content_type
    if body is not None:
        if body_length is None:
            if not isinstance(body, bytes):
                raise PolicyError("Git upstream request length is missing")
            body_length = len(body)
        headers["Content-Length"] = str(body_length)
    request = urllib.request.Request(
        CANONICAL_BASE_URL + target,
        data=body,
        method=method,
        headers=headers,
    )
    # One wall-clock budget covers connect, send, headers, and body read; the
    # watchdog closes the upstream socket if any phase overstays it.
    guard = deadline_http.StreamDeadline(UPSTREAM_DEADLINE_SECONDS)
    try:
        open_upstream = opener if opener is not None else _guarded_opener(guard)
        with open_upstream(request, timeout=guard.remaining()) as response:
            if _status(response) != 200:
                raise PolicyError("Git upstream returned an unexpected HTTP status")
            if response.headers.get_content_type() != expected_type:
                raise PolicyError("Git upstream returned an unexpected response type")
            result, result_length = _spool_response(
                response, MAX_GIT_RESPONSE, token, guard.remaining()
            )
    finally:
        guard.cancel()
    if stream_result:
        return result, result_length
    try:
        return result.read()
    finally:
        result.close()
class BrokerHandler(AbsoluteHeaderDeadlineMixin, BaseHTTPRequestHandler):
    """Expose only bounded metadata, draft creation, and smart-HTTP Git."""

    server_version = "HermesSCMBroker/1"
    sys_version = ""
    header_deadline_seconds = INBOUND_HEADER_TIMEOUT

    def _validate_headers(self) -> None:
        validate_ascii_headers(self.headers, maximum_count=MAX_HEADERS,
                               maximum_bytes=MAX_HEADER_BYTES, error=PolicyError)

    def _stream(
        self, status: int, content_type: str, body: BinaryIO, length: int
    ) -> None:
        self.send_response(status)
        self.send_header("Content-Type", content_type)
        self.send_header("Content-Length", str(length))
        self.end_headers()
        deadline = time.monotonic() + INBOUND_BODY_TIMEOUT
        while True:
            chunk = body.read(STREAM_CHUNK)
            if not chunk:
                break
            remaining = deadline - time.monotonic()
            if remaining <= 0:
                raise PolicyError("SCM response write deadline exceeded")
            self.connection.settimeout(remaining)
            self.wfile.write(chunk)

    def log_message(self, _format: str, *_args: object) -> None:
        return

    def _json(self, status: int, body: bytes) -> None:
        self.send_response(status)
        self.send_header("Content-Type", "application/json")
        self.send_header("Content-Length", str(len(body)))
        self.send_header("Cache-Control", "no-store")
        self.end_headers()
        self.wfile.write(body)

    def _reject(self, status: int = 400, phase: str = "", category: str = "") -> None:
        self._json(status, b'{"error":"request rejected"}')

    def do_GET(self) -> None:
        try:
            self._validate_headers()
            if self.path == "/healthz":
                self._json(200, b'{"status":"ok"}')
                return
            repo, operation, service = _git_target(self.path)
            if operation != "info/refs":
                raise PolicyError("Git RPC requires POST")
            token = read_token()
            _reject_forbidden(repo, "repository", (token,))
            expected = f"application/x-{service}-advertisement"
            streamed = _upstream_git_request(
                f"/titan/{repo}.git/info/refs?service={service}",
                method="GET",
                body=None,
                content_type=None,
                expected_type=expected,
                token=token,
                stream_result=True,
            )
            body, length = streamed
            try:
                self._stream(200, expected, body, length)
            finally:
                body.close()
        except (
            OSError,
            PolicyError,
            socket.timeout,
            urllib.error.URLError,
            ValueError,
        ):
            self._reject()

    def do_POST(self) -> None:
        try:
            self._phase = "headers"
            self._validate_headers()
            self.connection.settimeout(INBOUND_BODY_TIMEOUT)
            if self.path in {"/v1/metadata", "/v1/drafts", "/v1/tasks/register", "/v1/tasks/draft-update"}:
                self._phase = "control"
                self._control()
            else:
                self._git_rpc()
        except (
            OSError,
            PolicyError,
            socket.timeout,
            urllib.error.URLError,
            ValueError,
            json.JSONDecodeError,
        ) as error:
            log_rejection(getattr(self, "_phase", "request"), error)
            self._reject()

    def _control(self) -> None:
        data = _load_json(self)
        token = read_token()
        if self.path == "/v1/metadata":
            if set(data) != {"path"} or not isinstance(data["path"], str):
                raise PolicyError("metadata request fields are invalid")
            result = read(data["path"], token=token)
        elif self.path == "/v1/drafts":
            expected = {"base", "body", "head", "head_sha", "repo", "title"}
            if set(data) != expected or not all(
                isinstance(data[key], str) for key in expected
            ):
                raise PolicyError("draft request fields are invalid")
            result = create_draft(token=token, **data)  # type: ignore[arg-type]
        elif self.path == "/v1/tasks/register":
            if set(data) != {"grant"} or not isinstance(data["grant"], str):
                raise PolicyError("task registration fields are invalid")
            claims = verify_grant(data["grant"])
            remote = _branch_head(claims["repo"], claims["ref"], token)
            _task_ledger().register(claims, remote_head=remote)
            result = b'{"registered":true}'
        else:
            grant, number, title, body = request_fields(data, token)
            claims = verify_grant(grant)
            _task_ledger().authorize_current(claims)
            if _branch_head(claims["repo"], claims["ref"], token) != claims["new_head"]:
                raise PolicyError("task branch head changed; fetch and merge before retrying")
            pull = json.loads(read(f"/api/v1/repos/titan/{claims['repo']}/pulls/{number}", token=token))
            matches_pull(pull, claims, number)
            result = update_draft(token, claims["repo"], number, title, body)
        if token.encode("utf-8") in result:
            raise PolicyError("SCM upstream reflected credential material")
        self._json(200, result)

    def _git_rpc(self) -> None:
        self._phase = "target"
        repo, operation, service = _git_target(self.path)
        if (
            operation not in {"git-upload-pack", "git-receive-pack"}
            or service != operation
        ):
            raise PolicyError("Git RPC operation is outside the allowlist")
        self._phase = "headers"
        expected_request = f"application/x-{service}-request"
        if self.headers.get_content_type() != expected_request or self.headers.get(
            "Transfer-Encoding"
        ):
            raise PolicyError("Git RPC request type is invalid")
        length = _content_length(self.headers, MAX_GIT_REQUEST)
        token = read_token()
        _reject_forbidden(repo, "repository", (token,))
        self._phase = "body"
        body, body_length = _spool_bounded(
            self.rfile,
            MAX_GIT_REQUEST,
            length,
            token=token,
            context="Git request",
            deadline=time.monotonic() + INBOUND_BODY_TIMEOUT,
            set_timeout=self.connection.settimeout,
        )
        try:
            claims = None
            if service == "git-receive-pack":
                raw_grant = self.headers.get("X-Hermes-Task-Grant", "")
                if raw_grant:
                    self._phase = "grant"
                    claims = verify_grant(raw_grant)
                    if claims["repo"] != repo:
                        raise PolicyError("task grant repository does not match Git target")
                    self._phase = "ledger"
                    _task_ledger().authorize_update(claims)
                    self._phase = "pack"
                    validate_receive_pack(
                        body, token, _credential_forms(token),
                        expected=(claims["expected_old"], claims["new_head"], claims["ref"]),
                    )
                else:
                    # Compatibility is creation-only.  Existing refs cannot be
                    # moved without an owned, signed task grant.
                    _validate_receive_pack(body, token)
            expected = f"application/x-{service}-result"
            self._phase = "upstream"
            try:
                streamed = _upstream_git_request(
                    f"/titan/{repo}.git/{service}",
                    method="POST",
                    body=body,
                    body_length=body_length,
                    content_type=expected_request,
                    expected_type=expected,
                    token=token,
                    stream_result=True,
                )
            except (OSError, socket.timeout, urllib.error.URLError):
                # The upstream may have accepted the pack but lost its response.
                # Only the exact granted new head proves that a replay is safe.
                if claims is None or _branch_head(repo, claims["ref"], token) != claims["new_head"]:
                    raise
                _task_ledger().commit(claims)
                raise PolicyError("Git upstream acknowledgement was lost; retry to reconcile")
            result, result_length = streamed
            if claims is not None:
                # Smart HTTP uses HTTP 200 for both Git success and a rejected
                # ref command. The authenticated branch read is the commit
                # point; never advance the local ledger on packet status alone.
                self._phase = "persist"
                if _branch_head(repo, claims["ref"], token) != claims["new_head"]:
                    result.close()
                    raise PolicyError("Git upstream did not advance the granted branch")
                _task_ledger().commit(claims)
            try:
                self._phase = "response"
                self._stream(200, expected, result, result_length)
            finally:
                result.close()
        finally:
            body.close()
def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--listen", default="0.0.0.0")
    parser.add_argument("--port", type=int, default=BROKER_PORT)
    args = parser.parse_args()
    seed_adoptions(_task_ledger(), read_token(), _branch_head)
    BoundedThreadingHTTPServer((args.listen, args.port), BrokerHandler).serve_forever()
    return 0


if __name__ == "__main__":
    raise SystemExit(main())