242 lines
8.6 KiB
Python
242 lines
8.6 KiB
Python
"""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())
|