Secure Token Generator — Python source

#!/usr/bin/env python3
# ─────────────────────────────────────────────────────────────────────────────
# Secure token generator — Python polyglot showcase port.
# Language: Python 3.10+ (standard library only — no pip dependencies).
#
# CosmoDev polyglot showcase port of `token-generator`, ported from the
# canonical TypeScript logic in src/lib/token-generator.ts.
#
# This is display source — part of CosmoDev's polyglot tool pages, where each
# tool's pure logic is shown side-by-side in many languages.
#
# Design:
#   - Pure logic with an *injectable* RNG, so generation is unit-testable
#     without touching the secure RNG. Callers pass None for the CSPRNG
#     default; tests pass a seeded generator for exact, reproducible output.
#   - Each output symbol is selected without modulo bias via rejection
#     sampling (see `constant_time_select`), so even non-power-of-two
#     alphabets like base62 are unbiased.
#
# CSPRNG: Python's `secrets` module is the documented primitive for
# cryptographically secure randomness (it wraps os.urandom / getrandom).
# ─────────────────────────────────────────────────────────────────────────────

from __future__ import annotations

import math
import secrets
from dataclasses import dataclass
from typing import Callable, Literal, Optional

# Strength labels as a Literal type — mirrors the TS string-literal union.
StrengthLabel = Literal["weak", "fair", "strong", "very strong"]

#: The fixed alphabet strings. "custom" is intentionally absent — its symbols
#: are caller-supplied via ``GenerateOptions.custom_alphabet``.
ALPHABETS: dict[str, str] = {
    "hex":               "0123456789abcdef",
    "hex-upper":         "0123456789ABCDEF",
    "base32":            "ABCDEFGHIJKLMNOPQRSTUVWXYZ234567",        # RFC 4648
    "base32-crockford":  "0123456789ABCDEFGHJKMNPQRSTVWXYZ",      # Crockford (no I/L/O/U)
    "base64":            "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/",
    "base64url":         "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789-_",
    "base62":            "0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz",
    "alphanumeric":      "0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ",
}

#: Convenience encoding aliases mapping 1:1 onto a fixed alphabet.
_ENCODING_TO_ALPHABET: dict[str, str] = {
    "hex": "hex",
    "base32": "base32",
    "base64": "base64",
    "base64url": "base64url",
    "base62": "base62",
    "alphanumeric": "alphanumeric",
}


@dataclass
class GenerateOptions:
    """Controls token generation. All fields optional.

    ``alphabet`` wins over ``encoding`` when both are given. ``custom_alphabet``
    is used only when ``alphabet == "custom"``. ``rng`` overrides the default
    CSPRNG (useful for deterministic tests).
    """

    alphabet: Optional[str] = None
    custom_alphabet: Optional[str] = None
    encoding: Optional[str] = None
    rng: Optional[Callable[[], float]] = None


def _default_rng() -> float:
    """Draw a uniform float in [0, 1) using the kernel CSPRNG.

    ``secrets.randbelow(2**32)`` returns an unbiased integer in [0, 2^32)
    backed by the OS cryptographically secure RNG (getrandom/BCryptGenRandom).
    Dividing by 2^32 maps every outcome uniformly into the unit interval.
    """
    return secrets.randbelow(1 << 32) / (1 << 32)


def resolve_alphabet(
    opts: Optional[GenerateOptions] = None,
    encoding: Optional[str] = None,
) -> str:
    """Resolve the effective alphabet string from options and an encoding.

    Precedence: explicit ``opts.alphabet`` → ``encoding`` → ``"hex"``. A
    ``"custom"`` alphabet with no/empty ``custom_alphabet`` resolves to ``""``
    (an invalid, empty set).
    """
    if opts is not None and opts.alphabet == "custom":
        return opts.custom_alphabet or ""
    if opts is not None and opts.alphabet:
        return ALPHABETS.get(opts.alphabet, "")
    enc = encoding if encoding is not None else (opts.encoding if opts is not None else None)
    if enc and enc in _ENCODING_TO_ALPHABET:
        return ALPHABETS.get(_ENCODING_TO_ALPHABET[enc], "")
    return ALPHABETS["hex"]


def constant_time_select(
    alphabet: str,
    n: int,
    rng: Optional[Callable[[], float]] = None,
) -> str:
    """Build an ``n``-character string from ``alphabet`` WITHOUT modulo bias.

    Naive ``draw % size`` is biased whenever ``size`` does not divide the draw
    range: for base62 the trailing symbols would be slightly over-represented.
    Instead we reject any 32-bit draw landing in the uneven remainder
    (``>= limit``) and redraw, keeping every symbol exactly equally likely.
    The ``guard`` cap stops a pathological/constant RNG from looping forever.
    """
    r = rng or _default_rng
    size = len(alphabet)
    if size < 1 or n is None or n < 1:
        return ""
    # Largest multiple of ``size`` that fits in [0, 2^32-1]. Draws at or above
    # this boundary map unevenly under ``% size``, so we redraw.
    limit = (0xFFFFFFFF // size) * size
    out: list[str] = []
    for _ in range(n):
        x = r() * 0x100000000  # [0, 2^32)
        guard = 0
        while x >= limit and guard < 64:
            x = r() * 0x100000000
            guard += 1
        out.append(alphabet[math.floor(x) % size])
    return "".join(out)


def output_length(bytes_: float, alphabet_size: int) -> int:
    """Characters needed to carry ``bytes_`` bytes of entropy through an
    alphabet of ``alphabet_size`` symbols."""
    if not math.isfinite(bytes_) or bytes_ < 1 or alphabet_size < 2:
        return 0
    return math.ceil((bytes_ * 8) / math.log2(alphabet_size))


def generate_token(
    bytes_: float,
    opts: Optional[GenerateOptions] = None,
    encoding: Optional[str] = None,
) -> str:
    """Generate a token carrying ``bytes_`` bytes of underlying entropy.

    The token is rendered through ``opts.alphabet`` (or ``encoding``). Each
    character is sampled uniformly without modulo bias, so output is unbiased
    even for base62. Returns ``""`` for invalid input (non-positive or
    non-finite bytes, an alphabet under 2 symbols).

    Example: ``generate_token(16, GenerateOptions(alphabet="hex"))`` → 32 hex
    chars (128 bits).
    """
    if not math.isfinite(bytes_) or bytes_ < 1:
        return ""
    alphabet = resolve_alphabet(opts, encoding)
    if len(alphabet) < 2:
        return ""
    n = output_length(bytes_, len(alphabet))
    rng = opts.rng if opts is not None else None
    return constant_time_select(alphabet, n, rng)


def estimate_entropy(bytes_: float, alphabet_size: int) -> float:
    """Entropy (in bits) of a token of ``bytes_`` entropy in a ``size``-symbol
    alphabet.

    Equals ``output_length * log2(size)``, which is ``>= bytes * 8`` because
    the character count is rounded up to the next whole symbol.
    """
    if not math.isfinite(bytes_) or bytes_ < 1 or alphabet_size < 2:
        return 0.0
    return output_length(bytes_, alphabet_size) * math.log2(alphabet_size)


def strength_label(entropy_bits: float) -> StrengthLabel:
    """Bucket an entropy estimate (bits) into a human strength label.

    Tiers: weak <64 · fair 64–127 · strong 128–255 · very strong ≥256.
    """
    if not math.isfinite(entropy_bits) or entropy_bits < 64:
        return "weak"
    if entropy_bits < 128:
        return "fair"
    if entropy_bits < 256:
        return "strong"
    return "very strong"