400 lines
16 KiB
Python
400 lines
16 KiB
Python
#!/usr/bin/env python3
|
|
"""Claude UserPromptSubmit hook for operator-triggered lease promotion."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import fcntl
|
|
import importlib.util
|
|
import json
|
|
import os
|
|
import secrets
|
|
import stat
|
|
import subprocess
|
|
import sys
|
|
import time
|
|
from collections.abc import Callable, Mapping
|
|
from pathlib import Path
|
|
from typing import Final, TextIO
|
|
|
|
_MODULE_DIRECTORY = str(Path(__file__).resolve().parent)
|
|
if _MODULE_DIRECTORY not in sys.path:
|
|
sys.path.insert(0, _MODULE_DIRECTORY)
|
|
|
|
from receipt_challenge import receipt_for # noqa: E402
|
|
|
|
_observer_spec = importlib.util.spec_from_file_location(
|
|
"mosaic_receipt_observer_client", Path(__file__).resolve().with_name("receipt-observer-client.py")
|
|
)
|
|
if _observer_spec is None or _observer_spec.loader is None:
|
|
raise RuntimeError("unable to load receipt observer client")
|
|
_observer_module = importlib.util.module_from_spec(_observer_spec)
|
|
_observer_spec.loader.exec_module(_observer_module)
|
|
observer_request = _observer_module.observer_request
|
|
|
|
MAX_FRAME: Final = 64 * 1024
|
|
PENDING_MAX_AGE_SECONDS: Final = 60 * 60
|
|
PROMOTER_TIMEOUT_SECONDS: Final = 10.0
|
|
PROMOTION_PROMPT: Final = "/mosaic-promote"
|
|
PROMOTER: Final = Path(__file__).resolve().with_name("lease_promote.py")
|
|
PENDING_DIRECTORY: Final = "mosaic-lease"
|
|
AUTHORIZATION_DIRECTORY: Final = "authorizations"
|
|
AUTHORIZATION_TTL_SECONDS: Final = 60
|
|
LEASE_TTL_SECONDS: Final = 60 * 60
|
|
LOCK_FILE: Final = "promotion.lock"
|
|
RESULT_FILE: Final = "last-result.json"
|
|
EXPECTED_BEGIN_KEYS: Final = frozenset(
|
|
{"ok", "state", "receipt_challenge", "receipt", "binding"}
|
|
)
|
|
EXPECTED_BINDING_KEYS: Final = frozenset(
|
|
{
|
|
"compaction_epoch",
|
|
"request_epoch",
|
|
"h_source",
|
|
"h_payload",
|
|
"runtime_generation",
|
|
"schema_version",
|
|
}
|
|
)
|
|
|
|
|
|
class PromotionAlreadyInProgress(RuntimeError):
|
|
pass
|
|
|
|
|
|
def reject_duplicate_json_keys(pairs: list[tuple[str, object]]) -> dict[str, object]:
|
|
value: dict[str, object] = {}
|
|
for key, item in pairs:
|
|
if key in value:
|
|
raise ValueError("duplicate promoter JSON key")
|
|
value[key] = item
|
|
return value
|
|
|
|
|
|
def read_hook_input(stream: object) -> dict[str, object]:
|
|
raw = getattr(stream, "buffer", stream).read(MAX_FRAME + 1)
|
|
if not isinstance(raw, bytes) or len(raw) > MAX_FRAME:
|
|
raise ValueError("invalid UserPromptSubmit input")
|
|
value = json.loads(raw, object_pairs_hook=reject_duplicate_json_keys)
|
|
if not isinstance(value, dict):
|
|
raise ValueError("invalid UserPromptSubmit input")
|
|
return value
|
|
|
|
|
|
def emit_context(stream: TextIO, message: str) -> None:
|
|
json.dump(
|
|
{
|
|
"hookSpecificOutput": {
|
|
"hookEventName": "UserPromptSubmit",
|
|
"additionalContext": message,
|
|
}
|
|
},
|
|
stream,
|
|
separators=(",", ":"),
|
|
)
|
|
stream.write("\n")
|
|
|
|
|
|
def session_pending_name(environ: Mapping[str, str]) -> tuple[Path, str]:
|
|
runtime_dir = Path(environ["XDG_RUNTIME_DIR"])
|
|
session_id = environ["MOSAIC_LEASE_SESSION_ID"]
|
|
if not runtime_dir.is_absolute():
|
|
raise ValueError("XDG_RUNTIME_DIR must be absolute")
|
|
if len(session_id) != 64 or any(character not in "0123456789abcdef" for character in session_id):
|
|
raise ValueError("invalid lease session id")
|
|
return runtime_dir, f"pending-{session_id}"
|
|
|
|
|
|
def open_pending_directory(runtime_dir: Path) -> int:
|
|
directory_flags = (
|
|
os.O_RDONLY
|
|
| getattr(os, "O_CLOEXEC", 0)
|
|
| getattr(os, "O_DIRECTORY", 0)
|
|
| getattr(os, "O_NOFOLLOW", 0)
|
|
)
|
|
runtime_descriptor = os.open(runtime_dir, directory_flags)
|
|
try:
|
|
runtime_metadata = os.fstat(runtime_descriptor)
|
|
if (
|
|
not stat.S_ISDIR(runtime_metadata.st_mode)
|
|
or runtime_metadata.st_uid != os.getuid()
|
|
or stat.S_IMODE(runtime_metadata.st_mode) != 0o700
|
|
):
|
|
raise ValueError("unsafe XDG runtime directory")
|
|
try:
|
|
os.mkdir(PENDING_DIRECTORY, mode=0o700, dir_fd=runtime_descriptor)
|
|
except FileExistsError:
|
|
pass
|
|
descriptor = os.open(PENDING_DIRECTORY, directory_flags, dir_fd=runtime_descriptor)
|
|
finally:
|
|
os.close(runtime_descriptor)
|
|
|
|
metadata = os.fstat(descriptor)
|
|
if (
|
|
not stat.S_ISDIR(metadata.st_mode)
|
|
or metadata.st_uid != os.getuid()
|
|
or stat.S_IMODE(metadata.st_mode) != 0o700
|
|
):
|
|
os.close(descriptor)
|
|
raise ValueError("unsafe promotion pending directory")
|
|
return descriptor
|
|
|
|
|
|
def acquire_lock(directory_descriptor: int) -> int:
|
|
flags = (
|
|
os.O_RDWR
|
|
| os.O_CREAT
|
|
| getattr(os, "O_CLOEXEC", 0)
|
|
| getattr(os, "O_NOFOLLOW", 0)
|
|
)
|
|
descriptor = os.open(LOCK_FILE, flags, 0o600, dir_fd=directory_descriptor)
|
|
metadata = os.fstat(descriptor)
|
|
if (
|
|
not stat.S_ISREG(metadata.st_mode)
|
|
or metadata.st_uid != os.getuid()
|
|
or stat.S_IMODE(metadata.st_mode) != 0o600
|
|
):
|
|
os.close(descriptor)
|
|
raise ValueError("unsafe promotion lock file")
|
|
try:
|
|
fcntl.flock(descriptor, fcntl.LOCK_EX | fcntl.LOCK_NB)
|
|
except BlockingIOError as error:
|
|
os.close(descriptor)
|
|
raise PromotionAlreadyInProgress() from error
|
|
return descriptor
|
|
|
|
|
|
def sweep_stale_pending(directory_descriptor: int, current_time: float) -> None:
|
|
cutoff = current_time - PENDING_MAX_AGE_SECONDS
|
|
removed = False
|
|
with os.scandir(directory_descriptor) as entries:
|
|
for candidate in entries:
|
|
if not (
|
|
candidate.name.startswith("pending-")
|
|
or candidate.name.startswith(".pending-")
|
|
):
|
|
continue
|
|
try:
|
|
metadata = candidate.stat(follow_symlinks=False)
|
|
if metadata.st_mtime < cutoff and not stat.S_ISDIR(metadata.st_mode):
|
|
os.unlink(candidate.name, dir_fd=directory_descriptor)
|
|
removed = True
|
|
except FileNotFoundError:
|
|
continue
|
|
if removed:
|
|
os.fsync(directory_descriptor)
|
|
|
|
|
|
def consume_authorization(directory_descriptor: int, session_id: str, wall_clock: float) -> str | None:
|
|
flags = os.O_RDONLY | getattr(os, "O_CLOEXEC", 0) | getattr(os, "O_DIRECTORY", 0) | getattr(os, "O_NOFOLLOW", 0)
|
|
try:
|
|
authorization_descriptor = os.open(AUTHORIZATION_DIRECTORY, flags, dir_fd=directory_descriptor)
|
|
except FileNotFoundError:
|
|
return None
|
|
try:
|
|
metadata = os.fstat(authorization_descriptor)
|
|
if not stat.S_ISDIR(metadata.st_mode) or metadata.st_uid != os.getuid() or stat.S_IMODE(metadata.st_mode) != 0o700:
|
|
raise ValueError("unsafe promotion authorization directory")
|
|
name = f"{session_id}.auth"
|
|
try:
|
|
descriptor = os.open(name, os.O_RDONLY | getattr(os, "O_CLOEXEC", 0) | getattr(os, "O_NOFOLLOW", 0), dir_fd=authorization_descriptor)
|
|
except FileNotFoundError:
|
|
return None
|
|
try:
|
|
token_metadata = os.fstat(descriptor)
|
|
if not stat.S_ISREG(token_metadata.st_mode) or token_metadata.st_uid != os.getuid() or stat.S_IMODE(token_metadata.st_mode) != 0o600 or token_metadata.st_size <= 0 or token_metadata.st_size > MAX_FRAME:
|
|
raise ValueError("unsafe promotion authorization")
|
|
raw = os.read(descriptor, MAX_FRAME + 1)
|
|
finally:
|
|
os.close(descriptor)
|
|
os.unlink(name, dir_fd=authorization_descriptor)
|
|
os.fsync(authorization_descriptor)
|
|
token = json.loads(raw, object_pairs_hook=reject_duplicate_json_keys)
|
|
if not isinstance(token, dict) or set(token) != {"nonce", "seat", "session_id", "expires_at", "ts"}:
|
|
return None
|
|
nonce = token.get("nonce")
|
|
expires_at = token.get("expires_at")
|
|
issued_at = token.get("ts")
|
|
if token.get("session_id") != session_id or not isinstance(token.get("seat"), str) or not isinstance(nonce, str) or len(nonce) != 64 or any(char not in "0123456789abcdef" for char in nonce) or type(expires_at) not in (int, float) or type(issued_at) not in (int, float) or expires_at <= wall_clock or expires_at > issued_at + AUTHORIZATION_TTL_SECONDS:
|
|
return None
|
|
return nonce
|
|
finally:
|
|
os.close(authorization_descriptor)
|
|
|
|
|
|
def write_result(directory_descriptor: int, attempt_id: str, verified: bool, reason: str | None, session_id: str, wall_clock: float) -> None:
|
|
temporary = f".{RESULT_FILE}.tmp-{secrets.token_hex(8)}"
|
|
descriptor = os.open(temporary, os.O_WRONLY | os.O_CREAT | os.O_EXCL | getattr(os, "O_CLOEXEC", 0) | getattr(os, "O_NOFOLLOW", 0), 0o600, dir_fd=directory_descriptor)
|
|
try:
|
|
os.fchmod(descriptor, 0o600)
|
|
with os.fdopen(descriptor, "w", encoding="utf-8", closefd=False) as stream:
|
|
json.dump({"attempt_id": attempt_id, "expires_at_wallclock": wall_clock + LEASE_TTL_SECONDS if verified else None, "reason": reason, "session_id": session_id, "ts": wall_clock, "verified": verified}, stream, separators=(",", ":"), sort_keys=True)
|
|
stream.flush(); os.fsync(stream.fileno())
|
|
os.replace(temporary, RESULT_FILE, src_dir_fd=directory_descriptor, dst_dir_fd=directory_descriptor)
|
|
os.fsync(directory_descriptor)
|
|
finally:
|
|
os.close(descriptor)
|
|
|
|
|
|
def write_pending(directory_descriptor: int, name: str, challenge: str) -> None:
|
|
temporary = f".{name}.tmp-{secrets.token_hex(8)}"
|
|
flags = (
|
|
os.O_WRONLY
|
|
| os.O_CREAT
|
|
| os.O_EXCL
|
|
| getattr(os, "O_CLOEXEC", 0)
|
|
| getattr(os, "O_NOFOLLOW", 0)
|
|
)
|
|
descriptor = os.open(temporary, flags, 0o600, dir_fd=directory_descriptor)
|
|
try:
|
|
os.fchmod(descriptor, 0o600)
|
|
with os.fdopen(descriptor, "w", encoding="utf-8", closefd=False) as stream:
|
|
stream.write(challenge)
|
|
stream.flush()
|
|
os.fsync(stream.fileno())
|
|
os.replace(
|
|
temporary,
|
|
name,
|
|
src_dir_fd=directory_descriptor,
|
|
dst_dir_fd=directory_descriptor,
|
|
)
|
|
os.fsync(directory_descriptor)
|
|
except Exception:
|
|
try:
|
|
os.unlink(temporary, dir_fd=directory_descriptor)
|
|
except FileNotFoundError:
|
|
pass
|
|
raise
|
|
finally:
|
|
os.close(descriptor)
|
|
|
|
|
|
def parse_begin_reply(
|
|
completed: subprocess.CompletedProcess[str],
|
|
) -> tuple[str, dict[str, object] | None]:
|
|
if completed.returncode != 0:
|
|
return f"PROMOTER_EXIT_{completed.returncode}", None
|
|
try:
|
|
value = json.loads(
|
|
completed.stdout,
|
|
object_pairs_hook=reject_duplicate_json_keys,
|
|
)
|
|
except (json.JSONDecodeError, RecursionError, TypeError, ValueError):
|
|
return "INVALID_PROMOTER_REPLY", None
|
|
if not isinstance(value, dict):
|
|
return "INVALID_PROMOTER_REPLY", None
|
|
if value.get("ok") is False and set(value) == {"ok", "code"}:
|
|
code = value.get("code")
|
|
return code if isinstance(code, str) and code else "PROMOTION_BEGIN_REFUSED", value
|
|
if set(value) != EXPECTED_BEGIN_KEYS or value.get("ok") is not True:
|
|
return "INVALID_PROMOTER_REPLY", None
|
|
if value.get("state") != "PENDING_VERIFICATION":
|
|
return "INVALID_PROMOTER_REPLY", None
|
|
challenge = value.get("receipt_challenge")
|
|
receipt = value.get("receipt")
|
|
binding = value.get("binding")
|
|
if (
|
|
not isinstance(challenge, str)
|
|
or len(challenge) != 64
|
|
or any(character not in "0123456789abcdef" for character in challenge)
|
|
or not isinstance(receipt, str)
|
|
or not isinstance(binding, dict)
|
|
or set(binding) != EXPECTED_BINDING_KEYS
|
|
):
|
|
return "INVALID_PROMOTER_REPLY", None
|
|
integer_fields = (
|
|
"compaction_epoch",
|
|
"request_epoch",
|
|
"runtime_generation",
|
|
"schema_version",
|
|
)
|
|
if any(type(binding.get(field)) is not int or binding[field] < 0 for field in integer_fields):
|
|
return "INVALID_PROMOTER_REPLY", None
|
|
if not all(
|
|
isinstance(binding.get(field), str)
|
|
and len(binding[field]) == 64
|
|
and all(character in "0123456789abcdef" for character in binding[field])
|
|
for field in ("h_source", "h_payload")
|
|
):
|
|
return "INVALID_PROMOTER_REPLY", None
|
|
if not secrets.compare_digest(
|
|
receipt.encode("utf-8"),
|
|
receipt_for(challenge, binding).encode("utf-8"),
|
|
):
|
|
return "INVALID_PROMOTER_REPLY", None
|
|
return "", value
|
|
|
|
|
|
def main(
|
|
*,
|
|
environ: Mapping[str, str] | None = None,
|
|
stdin: object | None = None,
|
|
stdout: TextIO | None = None,
|
|
stderr: TextIO | None = None,
|
|
run: Callable[..., subprocess.CompletedProcess[str]] = subprocess.run,
|
|
now: Callable[[], float] = time.time,
|
|
) -> int:
|
|
source_environment = os.environ if environ is None else environ
|
|
input_stream = sys.stdin if stdin is None else stdin
|
|
output_stream = sys.stdout if stdout is None else stdout
|
|
error_stream = sys.stderr if stderr is None else stderr
|
|
|
|
try:
|
|
hook_input = read_hook_input(input_stream)
|
|
except (OSError, RecursionError, ValueError, json.JSONDecodeError) as error:
|
|
print(f"Mosaic promotion trigger ignored invalid hook input: {error}", file=error_stream)
|
|
return 0
|
|
if hook_input.get("prompt") != PROMOTION_PROMPT:
|
|
return 0
|
|
|
|
directory_descriptor: int | None = None
|
|
lock_descriptor: int | None = None
|
|
try:
|
|
runtime_dir, pending_name = session_pending_name(source_environment)
|
|
session_id = source_environment["MOSAIC_LEASE_SESSION_ID"]
|
|
directory_descriptor = open_pending_directory(runtime_dir)
|
|
lock_descriptor = acquire_lock(directory_descriptor)
|
|
wall_clock = now()
|
|
nonce = consume_authorization(directory_descriptor, session_id, wall_clock)
|
|
if nonce is None:
|
|
write_result(directory_descriptor, "0" * 64, False, "NOT_AUTHORIZED", session_id, wall_clock)
|
|
print("Mosaic promotion denied: NOT_AUTHORIZED.", file=error_stream)
|
|
return 0
|
|
sweep_stale_pending(directory_descriptor, wall_clock)
|
|
completed = run([sys.executable, "-I", "-S", "-B", str(PROMOTER), "--begin"], check=False, capture_output=True, text=True, env=dict(source_environment), timeout=PROMOTER_TIMEOUT_SECONDS)
|
|
code, reply = parse_begin_reply(completed)
|
|
if code or reply is None:
|
|
write_result(directory_descriptor, nonce, False, code or "PROMOTION_BEGIN_FAILED", session_id, now())
|
|
return 0
|
|
challenge = str(reply["receipt_challenge"])
|
|
observation = observer_request(
|
|
Path(source_environment["MOSAIC_RECEIPT_OBSERVER_SOCKET"]),
|
|
{"action": "record_runtime_observation", "session_id": session_id, "runtime_generation": int(source_environment["MOSAIC_RUNTIME_GENERATION"]), "runtime": "claude", "latest_assistant_message": reply["receipt"]},
|
|
)
|
|
if set(observation) != {"ok"} or observation.get("ok") is not True:
|
|
write_result(directory_descriptor, challenge, False, "OBSERVATION_REJECTED", session_id, now())
|
|
return 0
|
|
completion = run([sys.executable, "-I", "-S", "-B", str(PROMOTER), "--complete", challenge], check=False, capture_output=True, text=True, env=dict(source_environment), timeout=PROMOTER_TIMEOUT_SECONDS)
|
|
try:
|
|
outcome = json.loads(completion.stdout, object_pairs_hook=reject_duplicate_json_keys)
|
|
except (json.JSONDecodeError, ValueError):
|
|
outcome = None
|
|
if completion.returncode == 0 and isinstance(outcome, dict) and outcome.get("stage") == "promote_lease" and outcome.get("ok") is True and outcome.get("state") == "VERIFIED":
|
|
write_result(directory_descriptor, challenge, True, None, session_id, now())
|
|
else:
|
|
reason = outcome.get("code") if isinstance(outcome, dict) and isinstance(outcome.get("code"), str) else "PROMOTION_INCOMPLETE"
|
|
write_result(directory_descriptor, challenge, False, reason, session_id, now())
|
|
except PromotionAlreadyInProgress:
|
|
print("Mosaic promotion denied: PROMOTION_ALREADY_IN_PROGRESS.", file=error_stream)
|
|
except (KeyError, OSError, RecursionError, ValueError, subprocess.SubprocessError) as error:
|
|
print(f"Mosaic promotion begin failed: {type(error).__name__}: {error}", file=error_stream)
|
|
finally:
|
|
if lock_descriptor is not None:
|
|
os.close(lock_descriptor)
|
|
if directory_descriptor is not None:
|
|
os.close(directory_descriptor)
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|