05/18 Enhanced codes and functionalities

app/routes/auth.py

@@ -19,6 +19,8 @@ from app.services.auth_service import (
     verify_recovery_proof,
     generate_backup_codes,
     verify_and_consume_backup_code,
+    is_totp_code_used,
+    mark_totp_code_used,
 )
 
 auth_bp = Blueprint('auth', __name__)
@@ -280,7 +282,11 @@ def mfa_enable():
     if not pyotp.TOTP(secret).verify(totp_code, valid_window=1):
         return jsonify({'error': 'Invalid verification code'}), 400
 
-    totp_secret_enc, totp_iv = encrypt_totp_secret(secret)
+    # Prevent replay: reject a code that was already consumed within the valid window.
+    # user.id is not yet persisted (MFA not enabled), so use g.current_user_id directly.
+    if is_totp_code_used(g.current_user_id, totp_code):
+        return jsonify({'error': 'Verification code already used. Wait for the next code.'}), 400
+    mark_totp_code_used(g.current_user_id, totp_code)
     user.totp_secret = totp_secret_enc
     user.totp_iv = totp_iv
     user.totp_enabled = True
@@ -324,7 +330,11 @@ def mfa_disable():
     verified = False
 
     if totp_code:
+        if is_totp_code_used(user.id, totp_code):
+            return jsonify({'error': 'Verification code already used. Wait for the next code.'}), 400
         verified = pyotp.TOTP(plaintext_secret).verify(totp_code, valid_window=1)
+        if verified:
+            mark_totp_code_used(user.id, totp_code)
     elif backup_code:
         stored = json.loads(user.mfa_backup_codes or '[]')
         matched, remaining = verify_and_consume_backup_code(stored, backup_code)
@@ -380,7 +390,11 @@ def mfa_verify():
     verified = False
 
     if totp_code:
+        if is_totp_code_used(user.id, totp_code):
+            return jsonify({'error': 'Verification code already used. Wait for the next code.'}), 400
         verified = pyotp.TOTP(plaintext_secret).verify(totp_code, valid_window=1)
+        if verified:
+            mark_totp_code_used(user.id, totp_code)
 
     if not verified and backup_code:
         stored = json.loads(user.mfa_backup_codes or '[]')
@@ -454,8 +468,11 @@ def mfa_backup_codes_regenerate():
 
     import pyotp
     plaintext_secret = decrypt_totp_secret(user.totp_secret, user.totp_iv)
+    if is_totp_code_used(user.id, totp_code):
+        return jsonify({'error': 'Verification code already used. Wait for the next code.'}), 400
     if not pyotp.TOTP(plaintext_secret).verify(totp_code, valid_window=1):
         return jsonify({'error': 'Invalid verification code'}), 400
+    mark_totp_code_used(user.id, totp_code)
 
     plaintext_codes, hashed_codes = generate_backup_codes()
     user.mfa_backup_codes = json.dumps(hashed_codes)
@@ -886,7 +903,6 @@ def recovery_data():
     db.session.commit()
 
     return jsonify({
-        'enc_key_salt': user.enc_key_salt,
         'recovery_enc_salt': user.recovery_enc_salt,
         'recovery_iv': user.recovery_iv,
         'nonce': nonce,
@@ -902,11 +918,12 @@ def recovery_items():
     Requires X-Recovery-Proof header containing the HMAC-SHA256 proof:
         proof = HMAC-SHA256(key=enc_key_salt_bytes, msg=nonce_from_recovery_data)
 
-    The server validates the proof against the value stored in the DB
-    during /recovery/data — enc_key_salt is never sent in plaintext.
-    The challenge row is NOT consumed here; it is consumed by the final
-    POST /recover call so the client can call this endpoint to preview
-    items before committing the full re-encryption.
+    The enc_key_salt used as the HMAC key is NOT returned by /recovery/data;
+    the client must derive it by decrypting the recovery blob with the recovery
+    code. This ensures only the holder of the recovery code can compute the proof.
+
+    The challenge row is NOT consumed here — it is consumed by the final
+    POST /recover call so that endpoint can also validate the proof.
     Items are returned as encrypted ciphertext blobs only.
     """
     from app.models.recovery_challenge import RecoveryChallenge
@@ -921,7 +938,8 @@ def recovery_items():
     if not user or not user.recovery_enc_salt:
         return jsonify({'error': 'No recovery data found'}), 404
 
-    # Validate against the DB-stored challenge — safe across all workers.
+    # Validate against the DB-stored challenge without consuming it —
+    # POST /recover will consume it atomically on commit.
     challenge = RecoveryChallenge.query.filter_by(user_id=user.id).first()
     from datetime import datetime, timezone
     if not challenge or challenge.expires_at < datetime.now(timezone.utc).replace(tzinfo=None):

app/routes/vault.py

@@ -1,6 +1,7 @@
 from flask import Blueprint, request, jsonify, g
 from app import db, limiter, client_ip
 from app.models.vault_item import VaultItem, ItemType
+from app.models.folder import Folder
 from app.models.audit_log import AuditLog
 from app.services.auth_service import require_jwt
 
@@ -9,6 +10,25 @@ vault_bp = Blueprint('vault', __name__)
 VALID_TYPES = {t.value for t in ItemType}
 
 
+def _validate_folder_id(folder_id, user_id: int):
+    """
+    Verify folder_id belongs to user_id.
+    Returns the sanitised folder_id (int or None).
+    Raises ValueError with a safe message if the folder doesn't exist or
+    belongs to another user.
+    """
+    if folder_id is None:
+        return None
+    try:
+        folder_id = int(folder_id)
+    except (TypeError, ValueError):
+        raise ValueError('folder_id must be an integer')
+    folder = Folder.query.filter_by(id=folder_id, user_id=user_id).first()
+    if not folder:
+        raise ValueError('Folder not found')
+    return folder_id
+
+
 @vault_bp.route('', methods=['GET'])
 @limiter.limit('120 per minute')
 @require_jwt
@@ -40,6 +60,11 @@ def create_item():
     if not enc_data or not iv:
         return jsonify({'error': 'enc_data and iv are required'}), 400
 
+    try:
+        folder_id = _validate_folder_id(folder_id, g.current_user_id)
+    except ValueError as e:
+        return jsonify({'error': str(e)}), 400
+
     item = VaultItem(
         user_id=g.current_user_id,
         folder_id=folder_id,
@@ -93,7 +118,10 @@ def update_item(item_id):
             return jsonify({'error': 'name cannot be empty'}), 400
         item.name = name
     if 'folder_id' in data:
-        item.folder_id = data['folder_id']
+        try:
+            item.folder_id = _validate_folder_id(data['folder_id'], g.current_user_id)
+        except ValueError as e:
+            return jsonify({'error': str(e)}), 400
     if 'enc_data' in data:
         item.enc_data = data['enc_data']
     if 'iv' in data:
@@ -204,6 +232,10 @@ def import_items():
             skipped += 1
             continue
         folder_id = row.get('folder_id')
+        try:
+            folder_id = _validate_folder_id(folder_id, g.current_user_id)
+        except ValueError:
+            folder_id = None  # invalid/foreign folder — import to root instead of skipping
         enc_name = row.get('enc_name') or None
         iv_name = row.get('iv_name') or None
         item = VaultItem(