Files
workdock-platform/backend/workflows/totp.py
2026-03-27 03:04:02 +01:00

67 lines
2.2 KiB
Python

from __future__ import annotations
import base64
import hashlib
import hmac
import secrets
import struct
import string
from urllib.parse import quote
def generate_totp_secret(length: int = 20) -> str:
return base64.b32encode(secrets.token_bytes(length)).decode('ascii').rstrip('=')
def normalize_totp_token(value: str | None) -> str:
return ''.join(ch for ch in (value or '').strip() if ch.isdigit())
def _secret_bytes(secret: str) -> bytes:
padded = secret.strip().replace(' ', '').upper()
padding = '=' * ((8 - len(padded) % 8) % 8)
return base64.b32decode(padded + padding, casefold=True)
def generate_totp_token(secret: str, for_time: int, *, digits: int = 6, period: int = 30) -> str:
counter = int(for_time // period)
key = _secret_bytes(secret)
msg = struct.pack('>Q', counter)
digest = hmac.new(key, msg, hashlib.sha1).digest()
offset = digest[-1] & 0x0F
code_int = struct.unpack('>I', digest[offset:offset + 4])[0] & 0x7FFFFFFF
return str(code_int % (10**digits)).zfill(digits)
def verify_totp_token(secret: str, token: str, *, for_time: int, digits: int = 6, period: int = 30, window: int = 1) -> bool:
normalized = normalize_totp_token(token)
if len(normalized) != digits:
return False
for offset in range(-window, window + 1):
candidate_time = for_time + (offset * period)
if generate_totp_token(secret, candidate_time, digits=digits, period=period) == normalized:
return True
return False
def build_otpauth_uri(secret: str, *, account_name: str, issuer: str) -> str:
label = quote(f'{issuer}:{account_name}')
issuer_q = quote(issuer)
return f'otpauth://totp/{label}?secret={secret}&issuer={issuer_q}&algorithm=SHA1&digits=6&period=30'
def normalize_recovery_code(value: str | None) -> str:
raw = (value or '').strip().upper().replace(' ', '')
return raw
def generate_recovery_codes(count: int = 8) -> list[str]:
alphabet = string.ascii_uppercase + string.digits
codes = []
for _ in range(count):
parts = []
for _part in range(2):
parts.append(''.join(secrets.choice(alphabet) for _ in range(5)))
codes.append('-'.join(parts))
return codes