spb/company-atlas

1"""`catlas api-key create|list|revoke` — API keys for rate-limit tiers (the raw key is printed once; only its sha256 hash is stored)."""2from __future__ import annotations34import hashlib5import secrets6from datetime import UTC, datetime7from typing import Annotated89import typer10from rich.table import Table1112from companyatlas.cli import out, run_async13from companyatlas.db import execute, fetch_all, fetch_val, transaction14from companyatlas.ids import new_id1516TIERS = ("authenticated", "paid", "internal")17keys_app = typer.Typer(name="api-key", help="API keys (X-CA-API-Key) and their rate-limit tiers.", no_args_is_help=True)181920def _new_raw_key(tier: str) -> str:21    return f"ca_{tier[:4]}_{secrets.token_urlsafe(32)}"222324@keys_app.command("create")25def create(name: Annotated[str, typer.Argument(help="Human label (customer, service…)")],26           tier: Annotated[str, typer.Option("--tier", "-t", help="authenticated | paid | internal")] = "authenticated") -> None:27    """Create a key and print it once."""28    if tier not in TIERS:29        raise typer.BadParameter(f"tier must be one of {', '.join(TIERS)}")30    raw = _new_raw_key(tier)31    key_hash = hashlib.sha256(raw.encode("utf-8")).hexdigest()32    kid = new_id("api_key")3334    async def _insert() -> None:35        async with transaction() as conn:36            await execute(conn, "insert into api_keys (id, key_hash, prefix, name, tier) values (:id, :h, :p, :n, :t)", id=kid, h=key_hash, p=raw[:12], n=name, t=tier)3738    run_async(_insert())39    out.print(f"[green]created[/] {kid} tier={tier} name={name!r}")40    out.print("[bold]API key (shown once, store it now):[/]")41    out.print(raw)424344@keys_app.command("list")45def list_keys(include_revoked: bool = False) -> None:46    async def _rows():  # type: ignore[no-untyped-def]47        async with transaction() as conn:48            extra = "" if include_revoked else " where revoked_at is null"49            return await fetch_all(conn, f"select id, prefix, name, tier, created_at, last_used_at, request_count, revoked_at from api_keys{extra} order by created_at desc")5051    rows = run_async(_rows())52    t = Table("id", "prefix", "name", "tier", "created", "last used", "requests", "revoked")53    for r in rows:54        t.add_row(r["id"], r["prefix"] + "…", r["name"], r["tier"], str(r["created_at"])[:19], str(r["last_used_at"] or "")[:19], str(r["request_count"]),55                  str(r["revoked_at"] or "")[:19])56    out.print(t)575859@keys_app.command("revoke")60def revoke(key_id: Annotated[str, typer.Argument(help="key id (key_…)")]) -> None:61    async def _revoke() -> int:62        async with transaction() as conn:63            n = await fetch_val(conn, "with u as (update api_keys set revoked_at = :t where id = :id and revoked_at is null returning 1) select count(*) from u",64                                t=datetime.now(UTC), id=key_id)65            return int(n or 0)6667    n = run_async(_revoke())68    out.print("[green]revoked[/]" if n else "[yellow]no active key with that id[/]")697071def register(app: typer.Typer) -> None:72    app.add_typer(keys_app, name="api-key")737475__all__ = ["keys_app", "register"]76