#!/usr/bin/env python3
"""Mosaic external lease broker for Linux SO_PEERCRED authenticated clients."""

from __future__ import annotations

import argparse
import copy
import errno
from concurrent.futures import ThreadPoolExecutor
import json
import os
import secrets
import signal
import socket
import stat
import struct
import sys
import threading
import time
from pathlib import Path
from typing import Final

MAX_FRAME: Final = 64 * 1024
MAX_STATE: Final = 4 * 1024 * 1024
MAX_PENDING_TOKENS: Final = 256
MAX_IN_FLIGHT_CONNECTIONS: Final = 16
STATE_VERSION: Final = 1
CONNECTION_DEADLINE_SECONDS: Final = 1.0
HEX_256_LENGTH: Final = 64


class BrokerFailure(Exception):
    def __init__(self, code: str) -> None:
        super().__init__(code)
        self.code = code


class StateCommitUncertain(RuntimeError):
    def __init__(self) -> None:
        super().__init__("STATE_COMMIT_UNCERTAIN")


def is_non_negative_integer(value: object) -> bool:
    return type(value) is int and value >= 0


def is_hex_256(value: object) -> bool:
    return (
        isinstance(value, str)
        and len(value) == HEX_256_LENGTH
        and all(character in "0123456789abcdef" for character in value)
    )


def is_positive_decimal(value: object) -> bool:
    return (
        isinstance(value, str)
        and len(value) > 0
        and value[0] in "123456789"
        and all(character in "0123456789" for character in value)
    )


def valid_binding(binding: object) -> bool:
    if not isinstance(binding, dict):
        return False
    required = {"compaction_epoch", "request_epoch", "h_source", "h_payload", "schema_version"}
    if set(binding) != required:
        return False
    if not all(
        is_non_negative_integer(binding[field])
        for field in ("compaction_epoch", "request_epoch", "schema_version")
    ):
        return False
    return all(
        is_hex_256(binding[field])
        for field in ("h_source", "h_payload")
    )


def validate_state(value: object) -> dict[str, object]:
    if not isinstance(value, dict) or set(value) != {"version", "sessions", "tokens"}:
        raise BrokerFailure("STATE_INTEGRITY")
    if type(value["version"]) is not int or value["version"] != STATE_VERSION:
        raise BrokerFailure("STATE_INTEGRITY")
    sessions = value["sessions"]
    tokens = value["tokens"]
    if not isinstance(sessions, dict) or not isinstance(tokens, dict):
        raise BrokerFailure("STATE_INTEGRITY")
    if len(tokens) > MAX_PENDING_TOKENS:
        raise BrokerFailure("STATE_INTEGRITY")

    anchors: set[tuple[int, str]] = set()
    for session_id, session in sessions.items():
        if not is_hex_256(session_id) or not isinstance(session, dict):
            raise BrokerFailure("STATE_INTEGRITY")
        if set(session) != {"anchor_pid", "anchor_starttime", "runtime_generation"}:
            raise BrokerFailure("STATE_INTEGRITY")
        anchor_pid = session["anchor_pid"]
        anchor_starttime = session["anchor_starttime"]
        if type(anchor_pid) is not int or anchor_pid <= 0:
            raise BrokerFailure("STATE_INTEGRITY")
        if not is_positive_decimal(anchor_starttime):
            raise BrokerFailure("STATE_INTEGRITY")
        if not is_non_negative_integer(session["runtime_generation"]):
            raise BrokerFailure("STATE_INTEGRITY")
        anchor = (anchor_pid, anchor_starttime)
        if anchor in anchors:
            raise BrokerFailure("STATE_INTEGRITY")
        anchors.add(anchor)

    for token_value, token in tokens.items():
        if not is_hex_256(token_value) or not isinstance(token, dict):
            raise BrokerFailure("STATE_INTEGRITY")
        if set(token) != {"session_id", "runtime_generation", "binding", "consumed"}:
            raise BrokerFailure("STATE_INTEGRITY")
        session_id = token["session_id"]
        generation = token["runtime_generation"]
        session = sessions.get(session_id) if isinstance(session_id, str) else None
        if not is_hex_256(session_id) or not isinstance(session, dict):
            raise BrokerFailure("STATE_INTEGRITY")
        if not is_non_negative_integer(generation):
            raise BrokerFailure("STATE_INTEGRITY")
        if generation != session["runtime_generation"]:
            raise BrokerFailure("STATE_INTEGRITY")
        if not valid_binding(token["binding"]) or token["consumed"] is not False:
            raise BrokerFailure("STATE_INTEGRITY")
    return value


def proc_node(pid: int) -> dict[str, int | str]:
    try:
        text = Path(f"/proc/{pid}/stat").read_text(encoding="utf-8")
    except (FileNotFoundError, PermissionError, ProcessLookupError) as exc:
        raise BrokerFailure("PID_UNAVAILABLE") from exc
    close = text.rfind(")")
    fields = text[close + 2 :].split()
    if close < 0 or len(fields) < 20:
        raise BrokerFailure("PROC_STAT_INVALID")
    return {"pid": pid, "ppid": int(fields[1]), "starttime": fields[19]}


def verified_ancestry(peer_pid: int, anchor_pid: int, anchor_starttime: str) -> bool:
    chain: list[dict[str, int | str]] = []
    seen: set[int] = set()
    current = peer_pid
    while current > 0 and current not in seen:
        seen.add(current)
        node = proc_node(current)
        chain.append(node)
        if current == anchor_pid:
            if node["starttime"] != anchor_starttime:
                return False
            break
        current = int(node["ppid"])
    else:
        return False
    if int(chain[-1]["pid"]) != anchor_pid:
        return False
    for original in chain:
        repeated = proc_node(int(original["pid"]))
        if repeated["starttime"] != original["starttime"]:
            raise BrokerFailure("PID_STARTTIME_RACE")
    return True


def secure_parent(path: Path) -> None:
    parent = path.parent
    if not parent.is_dir() or stat.S_IMODE(parent.stat().st_mode) != 0o700:
        raise BrokerFailure("INSECURE_PARENT_MODE")


class StateStore:
    def __init__(self, path: Path) -> None:
        self.path = path
        secure_parent(path)
        self.poisoned = False
        self.value: dict[str, object] = {"version": STATE_VERSION, "sessions": {}, "tokens": {}}
        flags = os.O_RDONLY | getattr(os, "O_CLOEXEC", 0) | getattr(os, "O_NOFOLLOW", 0)
        try:
            descriptor = os.open(path, flags)
        except FileNotFoundError:
            return
        except OSError as exc:
            raise BrokerFailure("STATE_INTEGRITY") from exc
        try:
            metadata = os.fstat(descriptor)
            if not stat.S_ISREG(metadata.st_mode):
                raise BrokerFailure("STATE_INTEGRITY")
            if stat.S_IMODE(metadata.st_mode) != 0o600:
                raise BrokerFailure("INSECURE_STATE_MODE")
            if metadata.st_size > MAX_STATE:
                raise BrokerFailure("STATE_INTEGRITY")
            chunks = bytearray()
            while len(chunks) <= MAX_STATE:
                chunk = os.read(descriptor, min(64 * 1024, MAX_STATE + 1 - len(chunks)))
                if not chunk:
                    break
                chunks.extend(chunk)
            if len(chunks) > MAX_STATE:
                raise BrokerFailure("STATE_INTEGRITY")
            try:
                loaded = json.loads(chunks)
            except (json.JSONDecodeError, UnicodeDecodeError) as exc:
                raise BrokerFailure("STATE_INTEGRITY") from exc
            self.value = validate_state(loaded)
        except OSError as exc:
            raise BrokerFailure("STATE_INTEGRITY") from exc
        finally:
            os.close(descriptor)

    def sessions(self) -> dict[str, dict[str, object]]:
        sessions = self.value.get("sessions")
        if not isinstance(sessions, dict):
            raise BrokerFailure("STATE_INTEGRITY")
        return sessions

    def tokens(self) -> dict[str, dict[str, object]]:
        tokens = self.value.get("tokens")
        if not isinstance(tokens, dict):
            raise BrokerFailure("STATE_INTEGRITY")
        return tokens

    def commit(self) -> None:
        if self.poisoned:
            raise StateCommitUncertain()
        payload = (
            json.dumps(self.value, sort_keys=True, separators=(",", ":")) + "\n"
        ).encode()
        if len(payload) > MAX_STATE:
            raise BrokerFailure("STATE_TOO_LARGE")
        temporary = self.path.with_name(f".{self.path.name}.{os.getpid()}.tmp")
        descriptor = os.open(temporary, os.O_WRONLY | os.O_CREAT | os.O_EXCL, 0o600)
        replaced = False
        try:
            try:
                remaining = memoryview(payload)
                while remaining:
                    written = os.write(descriptor, remaining)
                    if written == 0:
                        raise OSError(errno.EIO, "state write made no progress")
                    remaining = remaining[written:]
                os.fsync(descriptor)
            finally:
                os.close(descriptor)
            os.replace(temporary, self.path)
            replaced = True
            try:
                directory = os.open(self.path.parent, os.O_RDONLY | os.O_DIRECTORY)
                try:
                    os.fsync(directory)
                finally:
                    os.close(directory)
            except OSError as exc:
                self.poisoned = True
                raise StateCommitUncertain() from exc
        finally:
            if not replaced:
                try:
                    temporary.unlink()
                except FileNotFoundError:
                    pass


class Broker:
    def __init__(self, store: StateStore) -> None:
        self.store = store

    def authenticate(self, peer_pid: int, request: dict[str, object]) -> tuple[str, dict[str, object]]:
        session_id = request.get("session_id")
        generation = request.get("runtime_generation")
        if not isinstance(session_id, str) or not is_non_negative_integer(generation):
            raise BrokerFailure("INVALID_IDENTITY")
        session = self.store.sessions().get(session_id)
        if not isinstance(session, dict):
            raise BrokerFailure("UNKNOWN_SESSION")
        anchor_pid = session.get("anchor_pid")
        anchor_starttime = session.get("anchor_starttime")
        current_generation = session.get("runtime_generation")
        if (
            type(anchor_pid) is not int
            or not isinstance(anchor_starttime, str)
            or not is_non_negative_integer(current_generation)
        ):
            raise BrokerFailure("STATE_INTEGRITY")
        if not verified_ancestry(peer_pid, anchor_pid, anchor_starttime):
            raise BrokerFailure("ANCESTRY_MISMATCH")
        if generation < current_generation:
            raise BrokerFailure("STALE_GENERATION")
        if generation > current_generation:
            session["runtime_generation"] = generation
            self.revoke_session_tokens(session_id)
        return session_id, session

    def session_for_anchor(
        self, anchor_pid: int, anchor_starttime: str
    ) -> tuple[str, dict[str, object]] | None:
        for session_id, session in self.store.sessions().items():
            if (
                session["anchor_pid"] == anchor_pid
                and session["anchor_starttime"] == anchor_starttime
            ):
                return session_id, session
        return None

    def revoke_session_tokens(self, session_id: str) -> None:
        tokens = self.store.tokens()
        for token_value in [
            value for value, token in tokens.items() if token["session_id"] == session_id
        ]:
            del tokens[token_value]

    def handle(self, peer: tuple[int, int, int], request: dict[str, object]) -> dict[str, object]:
        if self.store.poisoned:
            raise StateCommitUncertain()
        previous = copy.deepcopy(self.store.value)
        try:
            response = self._handle(peer, request)
            if self.store.value != previous:
                self.store.commit()
            return response
        except StateCommitUncertain:
            raise
        except Exception:
            self.store.value = previous
            raise

    def _handle(self, peer: tuple[int, int, int], request: dict[str, object]) -> dict[str, object]:
        peer_pid, peer_uid, peer_gid = peer
        action = request.get("action")
        if action == "register_anchor":
            if "session_id" in request:
                raise BrokerFailure("CALLER_SESSION_ID_REFUSED")
            generation = request.get("runtime_generation")
            if not is_non_negative_integer(generation):
                raise BrokerFailure("INVALID_GENERATION")
            anchor = proc_node(peer_pid)
            anchor_starttime = str(anchor["starttime"])
            existing = self.session_for_anchor(peer_pid, anchor_starttime)
            if existing is None:
                session_id = secrets.token_hex(32)
                self.store.sessions()[session_id] = {
                    "anchor_pid": peer_pid,
                    "anchor_starttime": anchor_starttime,
                    "runtime_generation": generation,
                }
            else:
                session_id, session = existing
                current_generation = session["runtime_generation"]
                if generation < current_generation:
                    raise BrokerFailure("STALE_GENERATION")
                if generation > current_generation:
                    session["runtime_generation"] = generation
                    self.revoke_session_tokens(session_id)
            return {"ok": True, "session_id": session_id, "peer": {"pid": peer_pid, "uid": peer_uid, "gid": peer_gid, "starttime": anchor["starttime"]}}
        if action == "authenticate":
            self.authenticate(peer_pid, request)
            return {"ok": True}
        if action == "mint_token":
            session_id, _ = self.authenticate(peer_pid, request)
            binding = request.get("binding")
            if not valid_binding(binding):
                raise BrokerFailure("INVALID_BINDING")
            if len(self.store.tokens()) >= MAX_PENDING_TOKENS:
                raise BrokerFailure("TOKEN_CAPACITY")
            token = secrets.token_hex(32)
            self.store.tokens()[token] = {
                "session_id": session_id,
                "runtime_generation": request["runtime_generation"],
                "binding": binding,
                "consumed": False,
            }
            return {"ok": True, "token": token}
        if action == "consume_token":
            session_id, _ = self.authenticate(peer_pid, request)
            token_value = request.get("token")
            token = self.store.tokens().get(token_value) if isinstance(token_value, str) else None
            if not isinstance(token, dict) or token.get("session_id") != session_id or token.get("runtime_generation") != request.get("runtime_generation") or token.get("consumed") is not False:
                raise BrokerFailure("TOKEN_REPLAY")
            del self.store.tokens()[token_value]
            return {"ok": True}
        raise BrokerFailure("UNKNOWN_ACTION")


def read_frame(connection: socket.socket, deadline: float) -> dict[str, object]:
    data = bytearray()
    while len(data) <= MAX_FRAME:
        remaining = deadline - time.monotonic()
        if remaining <= 0:
            raise BrokerFailure("MALFORMED_REQUEST")
        connection.settimeout(remaining)
        chunk = connection.recv(min(4096, MAX_FRAME + 1 - len(data)))
        if not chunk:
            break
        data.extend(chunk)
    if len(data) > MAX_FRAME:
        while True:
            remaining = deadline - time.monotonic()
            if remaining <= 0:
                raise BrokerFailure("MALFORMED_REQUEST")
            connection.settimeout(remaining)
            if not connection.recv(4096):
                break
        raise BrokerFailure("MALFORMED_REQUEST")
    if not data.endswith(b"\n") or data.count(b"\n") != 1:
        raise BrokerFailure("MALFORMED_REQUEST")
    try:
        value = json.loads(data)
    except (json.JSONDecodeError, UnicodeDecodeError) as exc:
        raise BrokerFailure("MALFORMED_REQUEST") from exc
    if not isinstance(value, dict):
        raise BrokerFailure("MALFORMED_REQUEST")
    return value


def handle_connection(
    connection: socket.socket,
    broker: Broker,
    broker_lock: threading.Lock,
) -> None:
    with connection:
        deadline = time.monotonic() + CONNECTION_DEADLINE_SECONDS
        try:
            raw = connection.getsockopt(socket.SOL_SOCKET, socket.SO_PEERCRED, 12)
            peer = struct.unpack("3i", raw)
            request = read_frame(connection, deadline)
        except BrokerFailure as exc:
            reply = {"ok": False, "code": exc.code}
        except OSError:
            return
        else:
            try:
                with broker_lock:
                    reply = broker.handle(peer, request)
            except BrokerFailure as exc:
                reply = {"ok": False, "code": exc.code}
        try:
            remaining = deadline - time.monotonic()
            if remaining <= 0:
                return
            connection.settimeout(remaining)
            connection.sendall((json.dumps(reply, separators=(",", ":")) + "\n").encode())
        except OSError:
            return


def serve(socket_path: Path, state_path: Path) -> None:
    secure_parent(socket_path)
    if socket_path.exists() or socket_path.is_symlink():
        raise BrokerFailure("SOCKET_ALREADY_EXISTS")
    store = StateStore(state_path)
    broker = Broker(store)
    broker_lock = threading.Lock()
    slots = threading.BoundedSemaphore(MAX_IN_FLIGHT_CONNECTIONS)
    fatal_lock = threading.Lock()
    fatal_errors: list[Exception] = []
    executor = ThreadPoolExecutor(
        max_workers=MAX_IN_FLIGHT_CONNECTIONS,
        thread_name_prefix="mosaic-lease-broker",
    )
    server = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
    server.bind(str(socket_path))
    os.chmod(socket_path, 0o600)
    owned = (socket_path.stat().st_dev, socket_path.stat().st_ino)
    stopping = False

    def stop(_signum: int, _frame: object) -> None:
        nonlocal stopping
        stopping = True
        server.close()

    def process_connection(connection: socket.socket) -> None:
        try:
            handle_connection(connection, broker, broker_lock)
        except Exception as exc:
            with fatal_lock:
                if not fatal_errors:
                    fatal_errors.append(exc)
        finally:
            slots.release()

    def fatal_error() -> Exception | None:
        with fatal_lock:
            return fatal_errors[0] if fatal_errors else None

    signal.signal(signal.SIGTERM, stop)
    signal.signal(signal.SIGINT, stop)
    server.listen(MAX_IN_FLIGHT_CONNECTIONS)
    server.settimeout(0.1)
    print("READY", flush=True)
    try:
        while not stopping:
            failure = fatal_error()
            if failure is not None:
                raise failure
            if not slots.acquire(timeout=0.1):
                continue
            try:
                connection, _ = server.accept()
            except socket.timeout:
                slots.release()
                continue
            except OSError:
                slots.release()
                failure = fatal_error()
                if failure is not None:
                    raise failure
                if stopping:
                    break
                raise
            try:
                executor.submit(process_connection, connection)
            except Exception:
                slots.release()
                connection.close()
                raise
    finally:
        server.close()
        executor.shutdown(wait=True)
        try:
            current = socket_path.stat()
            if (current.st_dev, current.st_ino) == owned:
                socket_path.unlink()
        except FileNotFoundError:
            pass


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--socket", required=True, type=Path)
    parser.add_argument("--state", required=True, type=Path)
    arguments = parser.parse_args()
    serve(arguments.socket, arguments.state)


if __name__ == "__main__":
    try:
        main()
    except BrokerFailure as failure:
        print(failure.code, file=sys.stderr)
        raise SystemExit(1)
    except StateCommitUncertain as failure:
        print(str(failure), file=sys.stderr)
        raise SystemExit(1)