101 lines
3.7 KiB
Python
101 lines
3.7 KiB
Python
#!/usr/bin/env python3
|
|
"""Revoke a runtime lease from a compaction or lifecycle observer."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
import os
|
|
import socket
|
|
import sys
|
|
from collections.abc import Callable, Mapping, Sequence
|
|
from pathlib import Path
|
|
from typing import Final
|
|
|
|
from lease_generation import bump_runtime_generation, read_runtime_generation
|
|
|
|
MAX_FRAME: Final = 64 * 1024
|
|
BROKER_TIMEOUT_SECONDS: Final = 1.5
|
|
|
|
|
|
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 revoke request")
|
|
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,
|
|
request: Callable[[Path, dict[str, object]], dict[str, object]] = broker_request,
|
|
) -> int:
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--runtime", required=True, choices=("claude", "pi"))
|
|
parser.add_argument("--reason", required=True)
|
|
parser.add_argument("--bump-generation", action="store_true")
|
|
arguments = parser.parse_args(argv)
|
|
source_environment = os.environ if environ is None else environ
|
|
|
|
try:
|
|
if not arguments.reason or len(arguments.reason) > 128:
|
|
raise ValueError("invalid revoke reason")
|
|
socket_path = Path(source_environment["MOSAIC_LEASE_BROKER_SOCKET"])
|
|
session_id = source_environment["MOSAIC_LEASE_SESSION_ID"]
|
|
if (
|
|
len(session_id) != 64
|
|
or any(character not in "0123456789abcdef" for character in session_id)
|
|
):
|
|
raise ValueError("invalid broker session")
|
|
generation = (
|
|
bump_runtime_generation(source_environment)
|
|
if arguments.bump_generation
|
|
else read_runtime_generation(source_environment)
|
|
)
|
|
reply = request(
|
|
socket_path,
|
|
{
|
|
"action": "revoke_lease",
|
|
"session_id": session_id,
|
|
"runtime_generation": generation,
|
|
"reason": arguments.reason,
|
|
"runtime": arguments.runtime,
|
|
},
|
|
)
|
|
if reply.get("ok") is not True or reply.get("state") != "UNVERIFIED":
|
|
raise ValueError("revocation refused")
|
|
except (KeyError, ValueError, OSError, json.JSONDecodeError):
|
|
# A fired observer must remain fail-closed even if the broker transport
|
|
# is unavailable. Advancing the private local generation fences every
|
|
# later tool check; the broker revokes the old lease when it next sees
|
|
# that higher generation. Explicit rollover already advanced it above.
|
|
if not arguments.bump_generation:
|
|
try:
|
|
bump_runtime_generation(source_environment)
|
|
except (KeyError, ValueError, OSError):
|
|
pass
|
|
print("Mosaic lease revocation failed; lifecycle transition denied.", file=sys.stderr)
|
|
return 2
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|