04/18 Enhance app (security, performance)

app/services/auth_service.py

@@ -1,9 +1,12 @@
+import base64
+import os
 import uuid
 import time
 from datetime import datetime, timedelta
 from functools import wraps
 
 import jwt
+from cryptography.hazmat.primitives.ciphers.aead import AESGCM
 from flask import current_app, request, g, jsonify
 from argon2 import PasswordHasher
 from argon2.exceptions import VerifyMismatchError, VerificationError, InvalidHashError
@@ -27,6 +30,44 @@ def verify_auth_token(auth_hash: str, stored_hash: str) -> bool:
         return False
 
 
+def _get_totp_key() -> bytes:
+    """
+    Return the 32-byte AES key used for server-side TOTP secret encryption.
+    The key is stored as a 64-char hex string in TOTP_ENCRYPTION_KEY config.
+    """
+    hex_key = current_app.config.get('TOTP_ENCRYPTION_KEY', '')
+    if not hex_key or len(hex_key) != 64:
+        raise RuntimeError(
+            'TOTP_ENCRYPTION_KEY must be set to a 64-character hex string (32 bytes). '
+            'Generate one with: python -c "import secrets; print(secrets.token_hex(32))"'
+        )
+    return bytes.fromhex(hex_key)
+
+
+def encrypt_totp_secret(plaintext_secret: str) -> tuple[str, str]:
+    """
+    Encrypt a plaintext TOTP base32 secret with AES-256-GCM.
+    Returns (ciphertext_b64, iv_b64).
+    """
+    key = _get_totp_key()
+    iv = os.urandom(12)
+    aesgcm = AESGCM(key)
+    ciphertext = aesgcm.encrypt(iv, plaintext_secret.encode(), None)
+    return base64.b64encode(ciphertext).decode(), base64.b64encode(iv).decode()
+
+
+def decrypt_totp_secret(ciphertext_b64: str, iv_b64: str) -> str:
+    """
+    Decrypt a base64-encoded AES-256-GCM TOTP secret ciphertext.
+    Returns the plaintext base32 secret string.
+    """
+    key = _get_totp_key()
+    iv = base64.b64decode(iv_b64)
+    ciphertext = base64.b64decode(ciphertext_b64)
+    aesgcm = AESGCM(key)
+    return aesgcm.decrypt(iv, ciphertext, None).decode()
+
+
 def generate_tokens(user_id: int) -> dict:
     """Return access_token and refresh_token JWTs, each with a unique jti."""
     now = datetime.utcnow()