"""Digest-only, task-scoped token registry for isolated writers.""" from __future__ import annotations import argparse import hashlib import hmac import json import os import grp import pwd import secrets import sqlite3 import stat import threading from datetime import datetime, timezone from pathlib import Path from typing import Any _REGISTRY_INIT_LOCK = threading.Lock() def utc_now() -> str: return datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ") def token_digest(token: str) -> str: return hashlib.sha256(token.encode("utf-8")).hexdigest() class TokenRegistry: def __init__(self, path: Path | str): self.path = Path(path) self._prepare() def _prepare(self) -> None: self.path.parent.mkdir(parents=True, exist_ok=True) if self.path.parent.is_symlink() or self.path.is_symlink(): raise PermissionError("writer_token_registry_symlink_rejected") with _REGISTRY_INIT_LOCK: with self._connect() as db: db.executescript( """ CREATE TABLE IF NOT EXISTS token_grants ( token_id TEXT PRIMARY KEY, token_digest TEXT NOT NULL UNIQUE, grant_json TEXT NOT NULL, revoked INTEGER NOT NULL DEFAULT 0, created_at TEXT NOT NULL ); CREATE TABLE IF NOT EXISTS token_uses ( token_id TEXT NOT NULL, operation TEXT NOT NULL, use_count INTEGER NOT NULL DEFAULT 0, PRIMARY KEY(token_id, operation), FOREIGN KEY(token_id) REFERENCES token_grants(token_id) ); """ ) os.chmod(self.path, 0o600) def _connect(self) -> sqlite3.Connection: db = sqlite3.connect(self.path, timeout=30, isolation_level=None) db.row_factory = sqlite3.Row db.execute("PRAGMA journal_mode=WAL") db.execute("PRAGMA synchronous=FULL") db.execute("PRAGMA foreign_keys=ON") return db def issue(self, grant: dict[str, Any], token: str | None = None) -> tuple[str, dict[str, Any]]: value = token or secrets.token_urlsafe(32) if len(value.encode("utf-8")) < 32: raise ValueError("writer_token_entropy_insufficient") record = dict(grant) record["schema_version"] = 2 record["token_digest"] = token_digest(value) record["revoked"] = False record.setdefault("created_at", utc_now()) required = ( "token_id", "audience", "task_id", "source_cidrs", "allowed_branch", "approved_base_oid", "allowed_operations", "allowed_paths", "not_before", "expires_at", "max_uses", "issued_by", ) if any(not record.get(field) for field in required): raise ValueError("writer_token_grant_incomplete") encoded = json.dumps(record, sort_keys=True, separators=(",", ":")) with self._connect() as db: db.execute("BEGIN IMMEDIATE") try: db.execute( """ INSERT INTO token_grants(token_id, token_digest, grant_json, revoked, created_at) VALUES(?,?,?,?,?) """, ( str(record["token_id"]), str(record["token_digest"]), encoded, 0, str(record["created_at"]), ), ) db.execute("COMMIT") except Exception: db.execute("ROLLBACK") raise return value, record def find(self, token: str) -> dict[str, Any] | None: candidate = token_digest(token) with self._connect() as db: rows = db.execute( "SELECT token_digest, grant_json, revoked FROM token_grants" ).fetchall() match = None for row in rows: if hmac.compare_digest(str(row["token_digest"]), candidate): match = dict(json.loads(str(row["grant_json"]))) match["revoked"] = bool(row["revoked"]) return match def consume(self, token_id: str, operation: str, maximum: int) -> int: if maximum < 1: raise PermissionError("writer_token_operation_exhausted") with self._connect() as db: db.execute("BEGIN IMMEDIATE") try: row = db.execute( "SELECT revoked FROM token_grants WHERE token_id=?", (token_id,), ).fetchone() if row is None or bool(row["revoked"]): raise PermissionError("writer_token_revoked") used_row = db.execute( "SELECT use_count FROM token_uses WHERE token_id=? AND operation=?", (token_id, operation), ).fetchone() used = int(used_row["use_count"]) if used_row else 0 if used >= maximum: raise PermissionError("writer_token_operation_exhausted") current = used + 1 db.execute( """ INSERT INTO token_uses(token_id, operation, use_count) VALUES(?,?,?) ON CONFLICT(token_id, operation) DO UPDATE SET use_count=excluded.use_count """, (token_id, operation, current), ) db.execute("COMMIT") return current except Exception: db.execute("ROLLBACK") raise def revoke(self, token_id: str) -> bool: with self._connect() as db: db.execute("BEGIN IMMEDIATE") changed = db.execute( "UPDATE token_grants SET revoked=1 WHERE token_id=? AND revoked=0", (token_id,), ).rowcount db.execute("COMMIT") return bool(changed) def usage(self, token_id: str, operation: str) -> int: with self._connect() as db: row = db.execute( "SELECT use_count FROM token_uses WHERE token_id=? AND operation=?", (token_id, operation), ).fetchone() return int(row["use_count"]) if row else 0 def _write_secret_file( path: Path, token: str, *, owner: str | None = None, group: str | None = None, ) -> None: path.parent.mkdir(parents=True, exist_ok=True) if path.parent.is_symlink() or path.is_symlink(): raise PermissionError("writer_token_file_symlink_rejected") flags = os.O_WRONLY | os.O_CREAT | os.O_EXCL if hasattr(os, "O_NOFOLLOW"): flags |= os.O_NOFOLLOW fd = os.open(path, flags, 0o600) try: if not stat.S_ISREG(os.fstat(fd).st_mode): raise PermissionError("writer_token_file_not_regular") if owner is not None or group is not None: uid = pwd.getpwnam(owner).pw_uid if owner is not None else -1 gid = grp.getgrnam(group).gr_gid if group is not None else -1 os.fchown(fd, uid, gid) os.write(fd, token.encode("utf-8")) os.fsync(fd) finally: os.close(fd) def main(argv: list[str] | None = None) -> int: parser = argparse.ArgumentParser() parser.add_argument("--registry", required=True) sub = parser.add_subparsers(dest="command", required=True) issue = sub.add_parser("issue") issue.add_argument("--grant-file", required=True) issue.add_argument("--token-file", required=True) issue.add_argument("--owner") issue.add_argument("--group") revoke = sub.add_parser("revoke") revoke.add_argument("--token-id", required=True) args = parser.parse_args(argv) registry = TokenRegistry(args.registry) if args.command == "issue": grant = json.loads(Path(args.grant_file).read_text(encoding="utf-8")) token, record = registry.issue(grant) _write_secret_file( Path(args.token_file), token, owner=args.owner, group=args.group, ) print(json.dumps({"token_id": record["token_id"], "created": True}, sort_keys=True)) return 0 changed = registry.revoke(args.token_id) print(json.dumps({"token_id": args.token_id, "revoked": changed}, sort_keys=True)) return 0 if __name__ == "__main__": raise SystemExit(main())