53 lines
1.4 KiB
Python
53 lines
1.4 KiB
Python
import base64
|
|
import hashlib
|
|
import hmac
|
|
import os
|
|
import struct
|
|
import time
|
|
from urllib.parse import quote
|
|
|
|
STEP_SECONDS = 30
|
|
DIGITS = 6
|
|
WINDOW = 1
|
|
SECRET_BYTES = 20
|
|
ISSUER = "checkbot"
|
|
|
|
|
|
def generate_secret() -> str:
|
|
return base64.b32encode(os.urandom(SECRET_BYTES)).decode().rstrip("=")
|
|
|
|
|
|
def _decode(secret: str) -> bytes:
|
|
return base64.b32decode(secret + "=" * (-len(secret) % 8))
|
|
|
|
|
|
def _code_at(secret: str, counter: int) -> str:
|
|
digest = hmac.new(_decode(secret), struct.pack(">Q", counter), hashlib.sha1).digest()
|
|
offset = digest[-1] & 0x0F
|
|
number = struct.unpack(">I", digest[offset:offset + 4])[0] & 0x7FFFFFFF
|
|
return str(number % 10 ** DIGITS).zfill(DIGITS)
|
|
|
|
|
|
def _counter(now: float | None) -> int:
|
|
return int((time.time() if now is None else now) // STEP_SECONDS)
|
|
|
|
|
|
def current_code(secret: str, now: float | None = None) -> str:
|
|
return _code_at(secret, _counter(now))
|
|
|
|
|
|
def verify(secret: str, code: str, now: float | None = None) -> bool:
|
|
code = code.replace(" ", "")
|
|
if len(code) != DIGITS or not code.isdigit():
|
|
return False
|
|
counter = _counter(now)
|
|
return any(
|
|
hmac.compare_digest(_code_at(secret, counter + offset), code)
|
|
for offset in range(-WINDOW, WINDOW + 1)
|
|
)
|
|
|
|
|
|
def provisioning_uri(secret: str, username: str) -> str:
|
|
return (f"otpauth://totp/{quote(ISSUER)}:{quote(username)}"
|
|
f"?secret={secret}&issuer={quote(ISSUER)}&digits={DIGITS}&period={STEP_SECONDS}")
|