OTP Code Generator — Python source

"""otp-code-generator — TOTP (RFC 6238) / HOTP (RFC 4226) code generator.

Language: Python (3.9+, standard library only)
Source:   CosmoDev polyglot showcase port of the OTP Code Generator tool, ported
          from cli/otp-code-generator/otp-code-generator.go (the live Go CLI
          twin — the authoritative reference) and src/lib/otp.ts (canonical
          TypeScript, which wraps the `otpauth` dependency).
License:  display source — part of CosmoDev's polyglot tool pages.

Design goals:
  - Pure + deterministic; never raises (generate returns Optional[str], None on
    a bad secret/algorithm).
  - Functionally equivalent to the Go/TS reference: same inputs -> same outputs
    (both implement RFC 6238, so tokens agree by construction).
  - Self-contained: stdlib only (hashlib/hmac, base64, struct).

Implements RFC 4226 (HOTP: HMAC the 8-byte counter, dynamic-truncate, mod
10^digits) and RFC 6238 (TOTP: counter = floor(timestamp_ms / 1000 / period),
then HOTP). `generate` mirrors Generate() in the Go twin (defaults SHA1, 6
digits, 30-second period); `validate` accepts the current period and +/-1
adjacent periods (matching otpauth's default window=1). Secrets are base32
(RFC 4648); whitespace/case tolerated, '=' padding stripped — exactly like
secretFrom()/decodeSecret() in the TS/Go. Python's // is floor division; every
operand here is non-negative, so it matches Go's integer division exactly.
"""

from __future__ import annotations

import base64
import hashlib
import hmac
import re
import struct
from dataclasses import dataclass
from typing import Callable, Optional

__all__ = ["generate", "validate", "hotp", "decode_secret"]


@dataclass
class Options:
    """Mirrors the Go twin's Options struct + the TS TotpOptions (minus the
    injectable timestamp, which is passed to generate() explicitly)."""
    secret: str                      # Base32 (RFC 4648); spaces/case tolerated
    algorithm: str = "SHA1"          # 'SHA1' | 'SHA256' | 'SHA512'
    digits: int = 6
    period: int = 30                 # seconds


def _with_defaults(opts: Options) -> Options:
    """Apply the TS defaults (the `??` coalescing in src/lib/otp.ts config()).
    A zero digits/period and an empty algorithm fall back to the otpauth
    defaults — 6, 30, and SHA1 — exactly like the Go twin's withDefaults()."""
    return Options(
        secret=opts.secret,
        algorithm=opts.algorithm or "SHA1",
        digits=opts.digits or 6,
        period=opts.period or 30,
    )


def decode_secret(secret: str) -> bytes:
    """Normalize + base32-decode the secret. Twin of decodeSecret() in the Go:
    strip whitespace, uppercase, strip '=' padding, then base32-decode
    (NoPadding mode tolerates missing/extra padding). Raises ValueError on any
    byte outside the base32 alphabet."""
    s = re.sub(r"\s+", "", secret).upper().rstrip("=")
    try:
        return base64.b32decode(s)  # NoPadding: python ignores missing '=' by default
    except Exception as exc:  # base32 binascii error -> a single ValueError type
        raise ValueError(f"invalid base32 secret: {exc}") from exc


def _new_hash(algo: str) -> Callable[[bytes, bytes], bytes]:
    """Return an hmac(key, msg) -> digest callable for the named algorithm.
    Case-insensitive; empty falls back to SHA1 (the TS default). Mirrors
    newHasher() in the Go twin."""
    table = {
        "SHA1": hashlib.sha1,
        "SHA256": hashlib.sha256,
        "SHA512": hashlib.sha512,
    }
    h = table.get((algo or "SHA1").upper())
    if h is None:
        raise ValueError(f"unknown algorithm: {algo!r}")
    return lambda key, msg: hmac.new(key, msg, h).digest()


def _truncate(digest: bytes, digits: int) -> str:
    """RFC 4226 section 5.4 dynamic truncation + mod 10^digits, zero-padded.
    Twin of the truncation tail of Generate() in the Go twin: mask the top bit
    of the 4-byte big-endian window (equiv. to Go's &0x7f on byte[offset])."""
    offset = digest[-1] & 0x0F
    bin_code = struct.unpack(">I", digest[offset:offset + 4])[0] & 0x7FFFFFFF
    return str(bin_code % (10 ** digits)).zfill(digits)


def hotp(opts: Options, counter: int) -> str:
    """RFC 4226 HOTP for opts.Secret at the given 8-byte counter. Shared core:
    generate() builds the counter from the timestamp then calls this. Raises
    ValueError on a bad secret/algorithm. It is the Python twin of the HOTP step
    inside Generate() in the Go twin."""
    opts = _with_defaults(opts)
    key = decode_secret(opts.secret)
    digest = _new_hash(opts.algorithm)(key, struct.pack(">Q", counter))
    return _truncate(digest, opts.digits)


def generate(opts: Options, timestamp_ms: int) -> Optional[str]:
    """TOTP (RFC 6238) for opts.Secret at timestamp_ms (milliseconds since the
    Unix epoch). The Python twin of generateTotp() in src/lib/otp.ts / Generate()
    in the Go — defaults SHA1, 6 digits, 30-second period. Returns None on a bad
    secret/algorithm (mirrors the try/catch in generateTotp())."""
    try:
        opts = _with_defaults(opts)
        counter = timestamp_ms // 1000 // opts.period  # RFC 6238 section 4.2
        return hotp(opts, counter)
    except ValueError:
        return None


def validate(token: str, opts: Options, timestamp_ms: int) -> bool:
    """Check token against opts.Secret at timestamp_ms, accepting the current
    period and +/-1 adjacent periods (otpauth window=1). Never raises — mirrors
    validateTotp() in src/lib/otp.ts / Validate() in the Go twin."""
    opts = _with_defaults(opts)
    period_ms = opts.period * 1000
    for ts in (timestamp_ms, timestamp_ms - period_ms, timestamp_ms + period_ms):
        try:
            got = hotp(opts, ts // 1000 // opts.period)
        except ValueError:
            return False
        if hmac.compare_digest(got, token):
            return True
    return False


if __name__ == "__main__":
    # Showcase vectors — shared with the TS/Go/Rust/PHP/JS twins so every
    # implementation is held to one contract. RFC is the base32 of ASCII
    # "12345678901234567890" — the RFC 6238/4226 reference key.
    RFC = "GEZDGNBVGY3TQOJQGEZDGNBVGY3TQOJQ"
    assert hotp(Options(RFC), 0) == "755224"                                   # RFC 4226 c=0
    assert generate(Options(RFC, digits=8), 59_000) == "94287082"              # RFC 6238 T=59s
    assert generate(Options("JBSWY3DPEHPK3PXP"), 1_700_000_000_000) == "324550"  # Go-twin lock-step
    assert generate(Options("jbsw y3dp ehpk 3pxp"), 1_700_000_000_000) == "324550"  # spaces/lowercase
    assert generate(Options("!!!not-base32!!!"), 0) is None                    # invalid secret
    assert validate("324550", Options("JBSWY3DPEHPK3PXP"), 1_700_000_000_000)  # round-trip
    print("otp: all showcase vectors passed")