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

from __future__ import annotations

import argparse
import errno
import json
import os
import secrets
import signal
import socket
import stat
import struct
import sys
import time
from pathlib import Path
from typing import Final

MAX_FRAME: Final = 64 * 1024
MAX_STATE: Final = 4 * 1024 * 1024
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


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")

    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 type(token["consumed"]) is not bool:
            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.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:
        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:
                payload = (
                    json.dumps(self.value, sort_keys=True, separators=(",", ":")) + "\n"
                ).encode()
                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
            directory = os.open(self.path.parent, os.O_RDONLY | os.O_DIRECTORY)
            try:
                os.fsync(directory)
            finally:
                os.close(directory)
        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)
            self.store.commit()
        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:
        for token in self.store.tokens().values():
            if token["session_id"] == session_id:
                token["consumed"] = True

    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,
                }
                self.store.commit()
            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)
                    self.store.commit()
            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")
            token = secrets.token_hex(32)
            self.store.tokens()[token] = {
                "session_id": session_id,
                "runtime_generation": request["runtime_generation"],
                "binding": binding,
                "consumed": False,
            }
            self.store.commit()
            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")
            token["consumed"] = True
            self.store.commit()
            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 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)
    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()

    signal.signal(signal.SIGTERM, stop)
    signal.signal(signal.SIGINT, stop)
    server.listen(16)
    print("READY", flush=True)
    try:
        while not stopping:
            try:
                connection, _ = server.accept()
            except OSError:
                if stopping:
                    break
                raise
            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:
                    continue
                else:
                    try:
                        reply = broker.handle(peer, request)
                    except BrokerFailure as exc:
                        reply = {"ok": False, "code": exc.code}
                try:
                    remaining = deadline - time.monotonic()
                    if remaining <= 0:
                        continue
                    connection.settimeout(remaining)
                    connection.sendall((json.dumps(reply, separators=(",", ":")) + "\n").encode())
                except OSError:
                    continue
    finally:
        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)