comfy-discord-web/token_store.py

"""
token_store.py
==============

Invite-token CRUD for the web UI.

Tokens are stored as SHA-256 hashes so the plaintext is never at rest.
The plaintext token is only returned once (at creation time).

File format (invite_tokens.json)::

    [
        {
            "id": "uuid",
            "label": "alice",
            "hash": "<sha256-hex>",
            "admin": false,
            "created_at": "2024-01-01T00:00:00"
        },
        ...
    ]

Usage::

    # Create a token (CLI)
    python -c "from token_store import create_token; print(create_token('alice'))"

    # Verify a token (used by auth.py)
    from token_store import verify_token
    record = verify_token(plaintext, token_file)
"""

from __future__ import annotations

import hashlib
import json
import logging
import secrets
import uuid
from datetime import datetime, timezone
from pathlib import Path
from typing import Optional

logger = logging.getLogger(__name__)


def _hash(token: str) -> str:
    return hashlib.sha256(token.encode()).hexdigest()


def _load(token_file: str) -> list[dict]:
    path = Path(token_file)
    if not path.exists():
        return []
    try:
        with open(path, "r", encoding="utf-8") as f:
            return json.load(f)
    except Exception as exc:
        logger.warning("Failed to load token file %s: %s", token_file, exc)
        return []


def _save(token_file: str, records: list[dict]) -> None:
    with open(token_file, "w", encoding="utf-8") as f:
        json.dump(records, f, indent=2)


def create_token(
    label: str,
    token_file: str = "invite_tokens.json",
    *,
    admin: bool = False,
) -> str:
    """
    Create a new invite token, persist its hash, and return the plaintext.

    Parameters
    ----------
    label : str
        Human-readable name for this token (e.g. ``"alice"``).
    token_file : str
        Path to the JSON store.
    admin : bool
        If True, grants admin privileges.

    Returns
    -------
    str
        The plaintext token (shown once; not stored).
    """
    plaintext = secrets.token_urlsafe(32)
    record = {
        "id": str(uuid.uuid4()),
        "label": label,
        "hash": _hash(plaintext),
        "admin": admin,
        "created_at": datetime.now(timezone.utc).isoformat(),
    }
    records = _load(token_file)
    records.append(record)
    _save(token_file, records)
    logger.info("Created token for '%s' (admin=%s)", label, admin)
    return plaintext


def verify_token(
    plaintext: str,
    token_file: str = "invite_tokens.json",
) -> Optional[dict]:
    """
    Verify a plaintext token against the stored hashes.

    Parameters
    ----------
    plaintext : str
        The token string provided by the user.
    token_file : str
        Path to the JSON store.

    Returns
    -------
    Optional[dict]
        The matching record dict (with ``label`` and ``admin`` fields),
        or None if no match.
    """
    h = _hash(plaintext)
    for record in _load(token_file):
        if record.get("hash") == h:
            return record
    return None


def list_tokens(token_file: str = "invite_tokens.json") -> list[dict]:
    """Return all token records (hashes included, labels safe to show)."""
    return _load(token_file)


def revoke_token(token_id: str, token_file: str = "invite_tokens.json") -> bool:
    """
    Delete a token by its UUID.

    Returns
    -------
    bool
        True if found and deleted, False if not found.
    """
    records = _load(token_file)
    new_records = [r for r in records if r.get("id") != token_id]
    if len(new_records) == len(records):
        return False
    _save(token_file, new_records)
    logger.info("Revoked token %s", token_id)
    return True


if __name__ == "__main__":
    import sys

    label = sys.argv[1] if len(sys.argv) > 1 else "default"
    is_admin = "--admin" in sys.argv
    tok = create_token(label, admin=is_admin)
    print(f"Token for '{label}': {tok}")