101 lines
3.0 KiB
Python
101 lines
3.0 KiB
Python
"""One-time verification codes backed by Redis.
|
|
|
|
Reusable primitive for register/verify-email/reset-password (and any future 2FA or
|
|
phone-number verification). Codes are bcrypt-hashed before storage so a Redis dump
|
|
does not leak active codes.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import secrets
|
|
from typing import Literal
|
|
|
|
import bcrypt
|
|
|
|
from app.core.security import redis_client
|
|
|
|
OtpPurpose = Literal["register", "verify_email", "reset_password"]
|
|
|
|
CODE_TTL_SECONDS = 600 # 10 minutes
|
|
RESEND_COOLDOWN_SECONDS = 60
|
|
MAX_ATTEMPTS = 5
|
|
CODE_LENGTH = 6
|
|
|
|
|
|
class OtpError(Exception):
|
|
code: str = "OTP_ERROR"
|
|
|
|
|
|
class OtpResendRateLimited(OtpError):
|
|
code = "OTP_RESEND_RATE_LIMITED"
|
|
|
|
def __init__(self, retry_after_seconds: int) -> None:
|
|
super().__init__(f"Resend allowed in {retry_after_seconds}s")
|
|
self.retry_after_seconds = retry_after_seconds
|
|
|
|
|
|
class OtpInvalid(OtpError):
|
|
code = "OTP_INVALID"
|
|
|
|
|
|
class OtpExpired(OtpError):
|
|
code = "OTP_EXPIRED"
|
|
|
|
|
|
class OtpAttemptsExceeded(OtpError):
|
|
code = "OTP_ATTEMPTS_EXCEEDED"
|
|
|
|
|
|
def _code_key(email: str, purpose: OtpPurpose) -> str:
|
|
return f"otp:{purpose}:{email.lower()}"
|
|
|
|
|
|
def _rate_key(email: str, purpose: OtpPurpose) -> str:
|
|
return f"otp_rate:{purpose}:{email.lower()}"
|
|
|
|
|
|
def _generate_code() -> str:
|
|
# secrets.randbelow gives uniform 0..10**CODE_LENGTH-1 without modulo bias
|
|
return f"{secrets.randbelow(10 ** CODE_LENGTH):0{CODE_LENGTH}d}"
|
|
|
|
|
|
def check_resend_allowed(email: str, purpose: OtpPurpose) -> None:
|
|
ttl = redis_client.ttl(_rate_key(email, purpose))
|
|
if ttl and ttl > 0:
|
|
raise OtpResendRateLimited(ttl)
|
|
|
|
|
|
def issue_code(email: str, purpose: OtpPurpose) -> str:
|
|
"""Generate a new code, persist its hash, and start the resend cooldown.
|
|
|
|
Caller is responsible for delivering the returned plaintext (e.g. via email).
|
|
Any pre-existing code for the same (purpose, email) is overwritten.
|
|
"""
|
|
check_resend_allowed(email, purpose)
|
|
code = _generate_code()
|
|
hashed = bcrypt.hashpw(code.encode(), bcrypt.gensalt()).decode()
|
|
payload = json.dumps({"hash": hashed, "attempts": 0})
|
|
redis_client.set(_code_key(email, purpose), payload, ex=CODE_TTL_SECONDS)
|
|
redis_client.set(_rate_key(email, purpose), "1", ex=RESEND_COOLDOWN_SECONDS)
|
|
return code
|
|
|
|
|
|
def verify_code(email: str, purpose: OtpPurpose, code: str) -> None:
|
|
"""Validate and consume a code. Raises subclasses of OtpError on failure."""
|
|
key = _code_key(email, purpose)
|
|
raw = redis_client.get(key)
|
|
if raw is None:
|
|
raise OtpExpired("Code expired or never issued")
|
|
record = json.loads(raw)
|
|
attempts = int(record.get("attempts", 0))
|
|
if attempts >= MAX_ATTEMPTS:
|
|
redis_client.delete(key)
|
|
raise OtpAttemptsExceeded("Too many invalid attempts")
|
|
if not bcrypt.checkpw(code.encode(), record["hash"].encode()):
|
|
record["attempts"] = attempts + 1
|
|
ttl = redis_client.ttl(key)
|
|
redis_client.set(key, json.dumps(record), ex=max(ttl, 1))
|
|
raise OtpInvalid("Incorrect code")
|
|
redis_client.delete(key)
|