IBeam TOTP handler

totp_handler.py

# Usage notes:
# - Set env vars:
#   IBEAM_TWO_FA_HANDLER=CUSTOM_HANDLER
#   IBEAM_CUSTOM_TWO_FA_HANDLER=totp_handler.TOTPHandler
#   IBEAM_TOTP=otpauth://totp/...:...?secret=...&issuer=...
#   IBEAM_TWO_FA_EL_ID=ID@@xyz-field-silver-response
#   IBEAM_TWO_FA_INPUT_EL_ID=ID@@xyz-field-silver-response
# - Place this file at project root and mount it to /srv/inputs/totp_handler.py.
# - Quick test: `python -c "from totp_handler import TOTPHandler; print(TOTPHandler().get_two_fa_code())"`
import base64
import hashlib
import hmac
import os
import struct
import time
from urllib.parse import parse_qs, unquote, urlparse

try:
    from ibeam.src.two_fa_handlers.two_fa_handler import TwoFaHandler
except Exception:  # pragma: no cover - fallback for local usage outside IBeam
    class TwoFaHandler:
        pass


class TOTPHandler(TwoFaHandler):
    def __init__(self, totp_value=None, *args, **kwargs):
        raw_value = totp_value if totp_value is not None else os.getenv("IBEAM_TOTP")
        if not raw_value:
            raise ValueError("IBEAM_TOTP is not set or empty.")

        secret, digits, period, algorithm = self._parse_totp_value(raw_value)
        self._digits = digits
        self._period = period
        self._hash_fn = self._hash_function(algorithm)
        self._secret_bytes = self._decode_secret(secret)

    def _parse_totp_value(self, raw_value):
        raw_value = raw_value.strip()
        if raw_value.lower().startswith("otpauth://"):
            parsed = urlparse(raw_value)
            if parsed.scheme.lower() != "otpauth" or parsed.netloc.lower() != "totp":
                raise ValueError("IBEAM_TOTP must be an otpauth://totp URL or base32 secret.")

            params = parse_qs(parsed.query)
            secret = params.get("secret", [None])[0]
            if not secret:
                raise ValueError("otpauth URL is missing a secret parameter.")

            digits = int(params.get("digits", ["6"])[0])
            period = int(params.get("period", ["30"])[0])
            algorithm = params.get("algorithm", ["SHA1"])[0].upper()
            return unquote(secret), digits, period, algorithm

        return raw_value, 6, 30, "SHA1"

    def _hash_function(self, algorithm):
        mapping = {
            "SHA1": hashlib.sha1,
            "SHA256": hashlib.sha256,
            "SHA512": hashlib.sha512,
        }
        if algorithm not in mapping:
            raise ValueError(f"Unsupported TOTP algorithm: {algorithm}")
        return mapping[algorithm]

    def _decode_secret(self, secret):
        cleaned = secret.strip().replace(" ", "").rstrip("=")
        padded = cleaned + ("=" * (-len(cleaned) % 8))
        return base64.b32decode(padded, casefold=True)

    def _generate_totp(self, timestamp=None):
        if timestamp is None:
            timestamp = time.time()
        counter = int(timestamp // self._period)
        message = struct.pack(">Q", counter)
        digest = hmac.new(self._secret_bytes, message, self._hash_fn).digest()
        offset = digest[-1] & 0x0F
        code = struct.unpack(">I", digest[offset:offset + 4])[0] & 0x7FFFFFFF
        return str(code % (10 ** self._digits)).zfill(self._digits)

    def get_two_fa_code(self, *args, **kwargs):
        return self._generate_totp()

    def __str__(self):
        return f"TOTPHandler(digits={self._digits}, period={self._period})"