72 lines
2.3 KiB
Python
72 lines
2.3 KiB
Python
import base64
|
|
import io
|
|
import json
|
|
|
|
import pyotp
|
|
import qrcode
|
|
|
|
from infrasynth.shared.crypto import decrypt, encrypt
|
|
from infrasynth.shared.settings_utils import get_setting
|
|
|
|
|
|
class TOTPService:
|
|
def __init__(self):
|
|
issuer = get_setting("INFRASYNTH_SECURITY", "TWO_FACTOR_ISSUER_NAME", "InfraSynth")
|
|
self.issuer = issuer
|
|
|
|
def generate_secret(self) -> str:
|
|
return pyotp.random_base32()
|
|
|
|
def get_provisioning_uri(self, secret: str, email: str) -> str:
|
|
return pyotp.totp.TOTP(secret).provisioning_uri(name=email, issuer_name=self.issuer)
|
|
|
|
def generate_qr_base64(self, secret: str, email: str) -> str:
|
|
uri = self.get_provisioning_uri(secret, email)
|
|
qr = qrcode.make(uri)
|
|
buf = io.BytesIO()
|
|
qr.save(buf, format="PNG")
|
|
return base64.b64encode(buf.getvalue()).decode()
|
|
|
|
def verify(self, secret: str, code: str) -> bool:
|
|
window = get_setting("INFRASYNTH_SECURITY", "TWO_FACTOR_TOTP_VALIDITY_WINDOW", 1)
|
|
totp = pyotp.TOTP(secret)
|
|
return totp.verify(code, valid_window=window)
|
|
|
|
def encrypt_secret(self, secret: str) -> str:
|
|
return encrypt(secret)
|
|
|
|
def decrypt_secret(self, encrypted: str) -> str:
|
|
return decrypt(encrypted)
|
|
|
|
|
|
class RecoveryCodeService:
|
|
def __init__(self):
|
|
self.count = get_setting("INFRASYNTH_SECURITY", "TWO_FACTOR_RECOVERY_CODES_COUNT", 8)
|
|
|
|
def generate_codes(self) -> list[str]:
|
|
import secrets
|
|
|
|
return [f"RC-{secrets.token_hex(4).upper()}-{secrets.token_hex(4).upper()}" for _ in range(self.count)]
|
|
|
|
def encrypt_codes(self, codes: list[str]) -> str:
|
|
return encrypt(json.dumps(codes))
|
|
|
|
def decrypt_codes(self, encrypted: str) -> list[str]:
|
|
return json.loads(decrypt(encrypted))
|
|
|
|
def verify_code(self, code: str, stored_encrypted: str) -> bool:
|
|
try:
|
|
codes = self.decrypt_codes(stored_encrypted)
|
|
return code in codes
|
|
except Exception:
|
|
return False
|
|
|
|
def remove_used_code(self, code: str, stored_encrypted: str) -> str | None:
|
|
try:
|
|
codes = self.decrypt_codes(stored_encrypted)
|
|
if code in codes:
|
|
codes.remove(code)
|
|
return self.encrypt_codes(codes)
|
|
return None
|
|
except Exception:
|
|
return None
|