stack/packages/mosaic/framework/tools/lease-broker/mutator-gate.py

#!/usr/bin/env python3
"""Runtime-neutral whole mutator-class gate backed by the Mosaic lease broker."""

from __future__ import annotations

import argparse
import json
import os
import socket
import sys
import re
from collections.abc import Callable, Mapping, Sequence
from pathlib import Path
from typing import BinaryIO, Final

# Isolated (`python -I`) adapter invocations must still import co-located
# framework modules; never depend on the caller's PYTHONPATH.
_MODULE_DIRECTORY = str(Path(__file__).resolve().parent)
if _MODULE_DIRECTORY not in sys.path:
    sys.path.insert(0, _MODULE_DIRECTORY)
from lease_generation import read_runtime_generation

MAX_FRAME: Final = 64 * 1024
BROKER_TIMEOUT_SECONDS: Final = 1.5
RECOVERY_TOOL: Final = "mosaic_context_recover"
_LITERAL_ABSOLUTE_PATH: Final = re.compile(r"/[A-Za-z0-9._/-]+\Z")
_SHELL_ACTIVE: Final = frozenset("$`~*?[]{}<>;|&" + '"' + "'" + "\\" + "\n\r\t")


def deny(code: str) -> int:
    print(f"BLOCKED: Mosaic mutator gate denied this tool ({code}).", file=sys.stderr)
    return 2


def read_tool_request(stream: BinaryIO | None = None) -> dict[str, object]:
    source = sys.stdin.buffer if stream is None else stream
    raw = source.read(MAX_FRAME + 1)
    if len(raw) > MAX_FRAME:
        raise ValueError("INVALID_GATE_INPUT")
    value = json.loads(raw)
    if not isinstance(value, dict):
        raise ValueError("INVALID_GATE_INPUT")
    tool_name = value.get("tool_name")
    if not isinstance(tool_name, str) or not tool_name or len(tool_name) > 256:
        raise ValueError("INVALID_GATE_INPUT")
    return value


def read_tool_name(stream: BinaryIO | None = None) -> str:
    """Backward-compatible strict extraction for callers that need only the name."""

    return str(read_tool_request(stream)["tool_name"])


def recovery_invocation_name(request: dict[str, object], recovery_command: Path | None) -> str:
    """Map only a byte-literal Claude recovery argv to ``RECOVERY_TOOL``.

    Claude's Bash tool evaluates its raw command with a real shell. Therefore
    the gate never attempts a second shell parser: any quote, expansion,
    redirection, operator, glob, newline, or non-space whitespace is refused
    before tokenizing. The remaining plain-space split is an exact argv proof,
    not a best-effort interpretation of shell syntax.
    """

    tool_name = request["tool_name"]
    if tool_name != "Bash" or recovery_command is None:
        return str(tool_name)
    tool_input = request.get("tool_input")
    if not isinstance(tool_input, dict) or set(tool_input) != {"command"}:
        return str(tool_name)
    command = tool_input.get("command")
    if (
        not isinstance(command, str)
        or not command
        or any(character in _SHELL_ACTIVE for character in command)
        or command.startswith(" ")
        or command.endswith(" ")
        or "  " in command
    ):
        return str(tool_name)
    argv = command.split(" ")
    if " ".join(argv) != command or len(argv) < 3 or argv[0] != "python3":
        return str(tool_name)
    if argv[1] != str(recovery_command):
        return str(tool_name)
    phase = argv[2]
    if phase == "complete" and len(argv) == 3:
        return RECOVERY_TOOL
    if phase != "begin" or len(argv) != 9:
        return str(tool_name)
    if argv[3::2] != ["--construction", "--compaction-epoch", "--request-epoch"]:
        return str(tool_name)
    construction, compaction_epoch, request_epoch = argv[4::2]
    if (
        _LITERAL_ABSOLUTE_PATH.fullmatch(construction) is None
        or not compaction_epoch.isdecimal()
        or not request_epoch.isdecimal()
    ):
        return str(tool_name)
    return RECOVERY_TOOL


def broker_request(socket_path: Path, request: dict[str, object]) -> dict[str, object]:
    payload = (json.dumps(request, separators=(",", ":")) + "\n").encode()
    if len(payload) > MAX_FRAME:
        raise ValueError("INVALID_GATE_INPUT")
    response = bytearray()
    with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as connection:
        connection.settimeout(BROKER_TIMEOUT_SECONDS)
        connection.connect(str(socket_path))
        connection.sendall(payload)
        connection.shutdown(socket.SHUT_WR)
        while len(response) <= MAX_FRAME:
            chunk = connection.recv(min(4096, MAX_FRAME + 1 - len(response)))
            if not chunk:
                break
            response.extend(chunk)
    if len(response) > MAX_FRAME or not response.endswith(b"\n"):
        raise ValueError("INVALID_BROKER_REPLY")
    value = json.loads(response)
    if not isinstance(value, dict):
        raise ValueError("INVALID_BROKER_REPLY")
    return value


def main(
    argv: Sequence[str] | None = None,
    *,
    environ: Mapping[str, str] | None = None,
    stream: BinaryIO | None = None,
    request: Callable[[Path, dict[str, object]], dict[str, object]] = broker_request,
    resolve_generation: Callable[[Mapping[str, str]], int] = read_runtime_generation,
) -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--runtime", required=True, choices=("claude", "pi"))
    parser.add_argument("--recovery-command", type=Path)
    arguments = parser.parse_args(argv)
    source_environment = os.environ if environ is None else environ

    try:
        request_input = read_tool_request(stream)
        tool_name = recovery_invocation_name(request_input, arguments.recovery_command)
        socket_value = source_environment["MOSAIC_LEASE_BROKER_SOCKET"]
        session_id = source_environment["MOSAIC_LEASE_SESSION_ID"]
        generation = resolve_generation(source_environment)
        reply = request(
            Path(socket_value),
            {
                "action": "authorize_tool",
                "session_id": session_id,
                "runtime_generation": generation,
                "runtime": arguments.runtime,
                "tool_name": tool_name,
            },
        )
    except (KeyError, ValueError, OSError, json.JSONDecodeError):
        return deny("GATE_UNAVAILABLE")

    if reply.get("ok") is True and reply.get("decision") == "allow":
        return 0
    code = reply.get("code")
    return deny(code if isinstance(code, str) else "MUTATOR_UNVERIFIED")


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