92 lines
3.0 KiB
Python
92 lines
3.0 KiB
Python
#!/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
|
|
from pathlib import Path
|
|
from typing import Final
|
|
|
|
MAX_FRAME: Final = 64 * 1024
|
|
BROKER_TIMEOUT_SECONDS: Final = 1.5
|
|
|
|
|
|
def deny(code: str) -> int:
|
|
print(f"BLOCKED: Mosaic mutator gate denied this tool ({code}).", file=sys.stderr)
|
|
return 2
|
|
|
|
|
|
def read_tool_name() -> str:
|
|
raw = sys.stdin.buffer.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 tool_name
|
|
|
|
|
|
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() -> int:
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--runtime", required=True, choices=("claude", "pi"))
|
|
arguments = parser.parse_args()
|
|
|
|
try:
|
|
tool_name = read_tool_name()
|
|
socket_value = os.environ["MOSAIC_LEASE_BROKER_SOCKET"]
|
|
session_id = os.environ["MOSAIC_LEASE_SESSION_ID"]
|
|
generation = int(os.environ["MOSAIC_RUNTIME_GENERATION"])
|
|
if generation < 0:
|
|
raise ValueError("INVALID_GENERATION")
|
|
reply = broker_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())
|