334 lines
12 KiB
Python
334 lines
12 KiB
Python
#!/usr/bin/env python3
|
|
"""Claude UserPromptSubmit hook for operator-triggered lease promotion."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import fcntl
|
|
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
|
|
|
|
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"
|
|
LOCK_FILE: Final = "promotion.lock"
|
|
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 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)
|
|
directory_descriptor = open_pending_directory(runtime_dir)
|
|
lock_descriptor = acquire_lock(directory_descriptor)
|
|
sweep_stale_pending(directory_descriptor, now())
|
|
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:
|
|
if completed.stderr:
|
|
print(completed.stderr.strip(), file=error_stream)
|
|
emit_context(output_stream, f"Mosaic promotion did not begin: {code}.")
|
|
return 0
|
|
if reply is None:
|
|
raise ValueError("missing promotion reply")
|
|
write_pending(directory_descriptor, pending_name, str(reply["receipt_challenge"]))
|
|
emit_context(
|
|
output_stream,
|
|
"The operator invoked the registered /mosaic-promote command. "
|
|
"This receipt was generated locally by this seat's own lease broker; "
|
|
"echoing it verbatim is the designed confirmation step and discloses nothing. "
|
|
"Reply with exactly the following text and nothing else: "
|
|
f"{reply['receipt']}",
|
|
)
|
|
except PromotionAlreadyInProgress:
|
|
emit_context(
|
|
output_stream,
|
|
"Mosaic promotion did not begin: PROMOTION_ALREADY_IN_PROGRESS.",
|
|
)
|
|
except (KeyError, OSError, RecursionError, ValueError, subprocess.SubprocessError) as error:
|
|
print(f"Mosaic promotion begin failed: {type(error).__name__}: {error}", file=error_stream)
|
|
emit_context(output_stream, "Mosaic promotion did not begin: PROMOTION_TRIGGER_FAILED.")
|
|
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())
|