Files
shd.build/tools/game-data-capture/src/game_data_capture/ingest/clipboard.py
T

511 lines
17 KiB
Python

"""Bounded, shell-free supervision for Wayland PNG clipboard intake."""
from __future__ import annotations
import hashlib
import json
import math
import os
import re
import signal
import stat
import subprocess
import sys
import threading
import time
from collections import deque
from collections.abc import Iterator, Mapping
from contextlib import suppress
from dataclasses import dataclass
from pathlib import Path
from typing import Final, cast
from game_data_capture.ingest.spool_helper import (
PNG_SIGNATURE,
acknowledge_receipt,
durable_receipts,
reconcile_spool,
)
_SAFE_ENV: Final = frozenset(
{"PATH", "HOME", "LANG", "LC_ALL", "WAYLAND_DISPLAY", "XDG_RUNTIME_DIR"}
)
_DENIED_ENV = re.compile(
r"(?:openai|authorization|token|secret|python(?:home|path|startup|userbase))", re.I
)
_SHA256 = re.compile(r"^[0-9a-f]{64}$")
_ID = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._:-]{0,127}$")
_MAX_METADATA = 16 * 1024
_MAX_RECEIPT = 1024
def _integer(value: object, minimum: int, maximum: int) -> bool:
return isinstance(value, int) and not isinstance(value, bool) and minimum <= value <= maximum
def _finite(value: object, minimum: float, maximum: float) -> bool:
return (
isinstance(value, int | float)
and not isinstance(value, bool)
and math.isfinite(value)
and minimum < value <= maximum
)
@dataclass(frozen=True, slots=True)
class WatchConfig:
spool: Path
max_image_bytes: int
max_queue_items: int
policy_version: str
max_restarts: int = 3
restart_backoff_seconds: float = 0.25
shutdown_timeout_seconds: float = 2.0
def __post_init__(self) -> None:
if not self.spool.is_absolute() or len(str(self.spool)) > 240:
raise ValueError("spool must be a bounded absolute path")
if not _integer(self.max_image_bytes, 8, 64 * 1024 * 1024):
raise ValueError("invalid image limit")
if not _integer(self.max_queue_items, 1, 256) or not _integer(self.max_restarts, 0, 10):
raise ValueError("invalid queue or restart limit")
if not _finite(self.restart_backoff_seconds, 0, 60) or not _finite(
self.shutdown_timeout_seconds, 0, 30
):
raise ValueError("invalid timeout")
if not re.fullmatch(r"[A-Za-z0-9._:-]{1,128}", self.policy_version):
raise ValueError("invalid policy version")
@dataclass(frozen=True, slots=True)
class FocusedWindowMetadata:
app_id: str | None
title: str | None
window_id: int | None
@dataclass(frozen=True, slots=True)
class IntakeReceipt:
occurrence_id: str
receipt_token: str
sha256: str
byte_size: int
media_type: str
spool_path: Path
policy_version: str
processing_version: str
focused_window: FocusedWindowMetadata
class VerifiedSpoolReader:
"""Descriptor-verified stream for handing one receipt to GDC-003 safely."""
def __init__(self, receipt: IntakeReceipt) -> None:
self.receipt = receipt
self._fd: int | None = None
def __enter__(self) -> VerifiedSpoolReader:
self._fd = _verify_spool_file(self.receipt)
return self
def __exit__(self, *_: object) -> None:
self.close()
def __iter__(self) -> Iterator[bytes]:
if self._fd is None:
raise RuntimeError("verified reader is not open")
while chunk := os.read(self._fd, 64 * 1024):
yield chunk
def close(self) -> None:
if self._fd is not None:
os.close(self._fd)
self._fd = None
def restricted_environment(environment: Mapping[str, str] | None = None) -> dict[str, str]:
source = os.environ if environment is None else environment
return {
key: value[:4096]
for key, value in source.items()
if key in _SAFE_ENV
and not _DENIED_ENV.search(key)
and isinstance(value, str)
and "\x00" not in value
}
def _bounded_text(value: object) -> str | None:
return " ".join(value.split())[:256] or None if isinstance(value, str) else None
def _bounded_command_output(
arguments: list[str], environment: Mapping[str, str], timeout: float
) -> bytes | None:
try:
process = subprocess.Popen(
arguments,
shell=False,
env=environment,
stdin=subprocess.DEVNULL,
stdout=subprocess.PIPE,
stderr=subprocess.DEVNULL,
start_new_session=True,
)
except OSError:
return None
output = bytearray()
done = threading.Event()
def collect() -> None:
if process.stdout is not None:
while len(output) <= _MAX_METADATA:
chunk = process.stdout.read(min(4096, _MAX_METADATA + 1 - len(output)))
if not chunk:
break
output.extend(chunk)
done.set()
reader = threading.Thread(target=collect, daemon=True)
reader.start()
if not done.wait(timeout):
with suppress(OSError):
os.killpg(process.pid, signal.SIGKILL)
with suppress(subprocess.TimeoutExpired):
process.wait(timeout=timeout)
if process.poll() is None:
with suppress(OSError):
os.killpg(process.pid, signal.SIGKILL)
process.wait()
if process.stdout is not None:
process.stdout.close()
reader.join(timeout)
return (
bytes(output)
if done.is_set() and process.returncode == 0 and len(output) <= _MAX_METADATA
else None
)
def focused_window(
*, environment: Mapping[str, str] | None = None, timeout_seconds: float = 0.5
) -> FocusedWindowMetadata:
if not _finite(timeout_seconds, 0, 30):
return FocusedWindowMetadata(None, None, None)
raw = _bounded_command_output(
["niri", "msg", "-j", "focused-window"],
restricted_environment(environment),
float(timeout_seconds),
)
if raw is None:
return FocusedWindowMetadata(None, None, None)
try:
payload = json.loads(raw)
except (json.JSONDecodeError, UnicodeDecodeError):
return FocusedWindowMetadata(None, None, None)
if not isinstance(payload, dict):
return FocusedWindowMetadata(None, None, None)
identifier = payload.get("id")
return FocusedWindowMetadata(
_bounded_text(payload.get("app_id")),
_bounded_text(payload.get("title")),
identifier if _integer(identifier, 0, 2**63 - 1) else None,
)
def _verify_spool_file(receipt: IntakeReceipt) -> int:
spool = receipt.spool_path.parent
if spool.is_symlink() or not spool.is_dir() or (spool.stat().st_mode & 0o777) != 0o700:
raise ValueError("unsafe spool")
fd = os.open(receipt.spool_path, os.O_RDONLY | os.O_NOFOLLOW)
try:
info = os.fstat(fd)
if (
not stat.S_ISREG(info.st_mode)
or (info.st_mode & 0o777) != 0o600
or info.st_size != receipt.byte_size
):
raise ValueError("invalid spool file")
first = os.read(fd, len(PNG_SIGNATURE))
if first != PNG_SIGNATURE:
raise ValueError("invalid spool png")
hasher = hashlib.sha256(first)
while chunk := os.read(fd, 64 * 1024):
hasher.update(chunk)
if hasher.hexdigest() != receipt.sha256:
raise ValueError("spool digest mismatch")
os.lseek(fd, 0, os.SEEK_SET)
return fd
except Exception:
os.close(fd)
raise
class ClipboardIntake:
def __init__(
self, config: WatchConfig, *, environment: Mapping[str, str] | None = None
) -> None:
self.config, self.environment = config, restricted_environment(environment)
self.armed = self.paused = False
self._watcher: subprocess.Popen[bytes] | None = None
self._reader: threading.Thread | None = None
self._queue: deque[IntakeReceipt] = deque()
self._helper_payloads: deque[dict[str, object]] = deque()
self._seen_tokens: set[str] = set()
self._seen_order: deque[str] = deque()
self._recovered_tokens: set[str] = set()
self._restart_count = 0
self._restart_after = 0.0
self._restarts_exhausted = False
@property
def queue_depth(self) -> int:
return len(self._queue)
@property
def helper_queue_depth(self) -> int:
return len(self._helper_payloads)
def arm(self) -> None:
self.armed, self.paused = True, False
self._start_if_due(time.monotonic())
def pause(self) -> None:
self.paused = True
self._stop_watcher()
def resume(self) -> None:
if self.armed:
self.paused = False
self._start_if_due(time.monotonic())
def disarm(self) -> None:
self.armed = self.paused = False
self._stop_watcher()
def cancel(self) -> None:
self.shutdown()
def shutdown(self) -> None:
self.disarm()
self._queue.clear()
self._helper_payloads.clear()
self._seen_tokens.clear()
self._seen_order.clear()
self._recovered_tokens.clear()
def tick(self, now: float | None = None) -> None:
current = time.monotonic() if now is None else now
if self._watcher is not None and self._watcher.poll() is not None:
self._stop_watcher()
self._crashed(current)
self._start_if_due(current)
def take(self) -> IntakeReceipt | None:
return self._queue.popleft() if self._queue else None
def take_helper_receipt(self) -> dict[str, object] | None:
return self._helper_payloads.popleft() if self._helper_payloads else None
def unreferenced_spool_candidates(self) -> list[Path]:
queued = {item.spool_path.name for item in self._queue}
return [
path
for path in self.config.spool.glob("*.png")
if not path.is_symlink() and path.name not in queued
]
def open_verified(self, receipt: IntakeReceipt) -> VerifiedSpoolReader:
return VerifiedSpoolReader(receipt)
def accept_helper_receipt(
self,
payload: Mapping[str, object],
*,
occurrence_id: str,
processing_version: str,
metadata: FocusedWindowMetadata | None = None,
) -> bool:
token, digest, size, name = (
payload.get("receipt_token"),
payload.get("sha256"),
payload.get("byte_size"),
payload.get("spool_name"),
)
if (
not self.armed
or self.paused
or len(self._queue) >= self.config.max_queue_items
or payload.get("receipt_version") != 1
or not isinstance(token, str)
or not _ID.fullmatch(token)
or token in self._seen_tokens
or not isinstance(digest, str)
or not _SHA256.fullmatch(digest)
or not _integer(size, 8, self.config.max_image_bytes)
or payload.get("media_type") != "image/png"
or payload.get("policy_version") != self.config.policy_version
or not isinstance(name, str)
or name != f"{digest}.png"
or not isinstance(occurrence_id, str)
or not _ID.fullmatch(occurrence_id)
or not isinstance(processing_version, str)
or not _ID.fullmatch(processing_version)
):
return False
recovered = token in self._recovered_tokens
receipt = IntakeReceipt(
occurrence_id,
token,
digest,
cast(int, size),
"image/png",
self.config.spool / name,
self.config.policy_version,
processing_version,
FocusedWindowMetadata(None, None, None)
if recovered
else focused_window()
if metadata is None
else metadata,
)
try:
verified_fd = _verify_spool_file(receipt)
except (OSError, ValueError):
return False
os.close(verified_fd)
self._seen_tokens.add(token)
self._seen_order.append(token)
if len(self._seen_order) > self.config.max_queue_items:
self._seen_tokens.remove(self._seen_order.popleft())
self._queue.append(receipt)
return True
def reconcile(self) -> list[str]:
removed = reconcile_spool(self.config.spool)
capacity = self.config.max_queue_items - len(self._helper_payloads)
if capacity <= 0:
return removed
for payload in durable_receipts(self.config.spool, max_items=256):
if len(self._helper_payloads) >= self.config.max_queue_items:
break
token = payload["receipt_token"]
if isinstance(token, str) and token not in self._seen_tokens:
self._helper_payloads.append(payload)
self._recovered_tokens.add(token)
return removed
def acknowledge(self, receipt: IntakeReceipt) -> bool:
"""Delete this occurrence's durable receipt after repository confirmation only."""
if (
receipt.spool_path.parent != self.config.spool
or receipt.policy_version != self.config.policy_version
):
return False
payload: dict[str, object] = {
"receipt_version": 1,
"receipt_token": receipt.receipt_token,
"sha256": receipt.sha256,
"byte_size": receipt.byte_size,
"media_type": receipt.media_type,
"spool_name": receipt.spool_path.name,
"policy_version": receipt.policy_version,
}
acknowledged = acknowledge_receipt(self.config.spool, payload)
if acknowledged:
self._recovered_tokens.discard(receipt.receipt_token)
return acknowledged
def _crashed(self, now: float) -> None:
if self.armed and not self.paused and self._restart_count < self.config.max_restarts:
self._restart_count += 1
self._restart_after = now + float(self.config.restart_backoff_seconds) * 2 ** (
self._restart_count - 1
)
else:
self._restarts_exhausted = True
def _start_if_due(self, now: float) -> None:
if (
not self.armed
or self.paused
or self._watcher is not None
or self._restarts_exhausted
or now < self._restart_after
):
return
try:
self._watcher = subprocess.Popen(
[
"wl-paste",
"--type",
"image/png",
"--watch",
sys.executable,
"-m",
"game_data_capture.ingest.spool_helper",
"--spool",
str(self.config.spool),
"--max-bytes",
str(self.config.max_image_bytes),
"--policy-version",
self.config.policy_version,
"--media-type",
"image/png",
],
shell=False,
env=self.environment,
stdin=subprocess.DEVNULL,
stdout=subprocess.PIPE,
stderr=subprocess.DEVNULL,
start_new_session=True,
)
except OSError:
self._crashed(now)
return
self._reader = threading.Thread(
target=self._read_receipts, args=(self._watcher,), daemon=True
)
self._reader.start()
def _read_receipts(self, process: subprocess.Popen[bytes]) -> None:
if process.stdout is None:
return
while True:
line = process.stdout.readline(_MAX_RECEIPT + 1)
if not line:
return
if len(line) > _MAX_RECEIPT and not line.endswith(b"\n"):
while line and not line.endswith(b"\n"):
line = process.stdout.readline(_MAX_RECEIPT + 1)
continue
try:
payload = json.loads(line)
except (json.JSONDecodeError, UnicodeDecodeError):
continue
if (
isinstance(payload, dict)
and all(isinstance(key, str) for key in payload)
and len(self._helper_payloads) < self.config.max_queue_items
):
self._helper_payloads.append(payload)
def _stop_watcher(self) -> None:
process, self._watcher = self._watcher, None
if process is None:
return
if process.poll() is None:
with suppress(OSError):
os.killpg(process.pid, signal.SIGTERM)
try:
process.wait(timeout=float(self.config.shutdown_timeout_seconds))
except subprocess.TimeoutExpired:
with suppress(OSError):
os.killpg(process.pid, signal.SIGKILL)
process.wait()
if process.stdout is not None:
process.stdout.close()
if self._reader is not None:
self._reader.join(timeout=float(self.config.shutdown_timeout_seconds))
self._reader = None