05/18 Enhanced codes and functionalities

app/__init__.py

@@ -105,6 +105,7 @@ def create_app(config_name: str = 'development') -> Flask:
     from .models.emergency_access import EmergencyAccess
     from .models.audit_log import AuditLog
     from .models.recovery_challenge import RecoveryChallenge
+    from .models.totp_used_code import TotpUsedCode
 
     @login_manager.user_loader
     def load_user(user_id):
@@ -190,8 +191,10 @@ def create_app(config_name: str = 'development') -> Flask:
                 try:
                     from app.models.token_blacklist import TokenBlacklist
                     from app.models.recovery_challenge import RecoveryChallenge
+                    from app.models.totp_used_code import TotpUsedCode
                     TokenBlacklist.cleanup_expired()
                     RecoveryChallenge.cleanup_expired()
+                    TotpUsedCode.cleanup_expired()
                     import logging
                     logging.getLogger(__name__).debug(
                         '[PassKeeper] token_blacklist + recovery_challenges cleanup completed'

app/models/totp_used_code.py

@@ -0,0 +1,78 @@
+from datetime import datetime, timezone, timedelta
+
+from sqlalchemy.dialects.mysql import INTEGER
+
+from app import db
+
+# A TOTP code is valid for at most one 30-second window on each side
+# of the current window (valid_window=1), giving a 90-second replay
+# window. We keep used-code records for 120 seconds to be safe.
+TOTP_CODE_TTL_SECONDS = 120
+
+
+class TotpUsedCode(db.Model):
+    """
+    One-time consumption record for TOTP codes.
+
+    Prevents replay attacks within the valid_window: a code that has
+    already been accepted for a given user cannot be reused within the
+    TOTP_CODE_TTL_SECONDS window, even though the code is still
+    mathematically valid according to pyotp.
+
+    Rows older than TOTP_CODE_TTL_SECONDS are cleaned up by the same
+    APScheduler job that handles token_blacklist and recovery_challenges.
+    """
+    __tablename__ = 'totp_used_codes'
+
+    id = db.Column(INTEGER(unsigned=True), autoincrement=True, primary_key=True)
+    user_id = db.Column(
+        INTEGER(unsigned=True),
+        db.ForeignKey('users.id', ondelete='CASCADE'),
+        nullable=False,
+        index=True,
+    )
+    # The 6-digit code as a zero-padded string, e.g. "042857"
+    code = db.Column(db.String(6), nullable=False)
+    expires_at = db.Column(db.DateTime, nullable=False)
+
+    __table_args__ = (
+        db.UniqueConstraint('user_id', 'code', name='uq_totp_used_user_code'),
+    )
+
+    @classmethod
+    def is_used(cls, user_id: int, code: str) -> bool:
+        """Return True if this code has already been consumed for this user."""
+        entry = cls.query.filter_by(user_id=user_id, code=code).first()
+        if not entry:
+            return False
+        return entry.expires_at > datetime.now(timezone.utc).replace(tzinfo=None)
+
+    @classmethod
+    def mark_used(cls, user_id: int, code: str) -> None:
+        """
+        Record that this code was consumed.
+        Caller must commit the session (the surrounding request handler does this).
+        Silently ignores IntegrityError (duplicate insert) — already marked used.
+        """
+        from sqlalchemy.exc import IntegrityError
+        expires_at = (
+            datetime.now(timezone.utc).replace(tzinfo=None)
+            + timedelta(seconds=TOTP_CODE_TTL_SECONDS)
+        )
+        entry = cls(user_id=user_id, code=code, expires_at=expires_at)
+        db.session.add(entry)
+        try:
+            db.session.flush()
+        except IntegrityError:
+            db.session.rollback()  # already marked — safe to ignore
+
+    @classmethod
+    def cleanup_expired(cls) -> None:
+        """Delete expired records — called by the APScheduler cleanup job."""
+        cls.query.filter(
+            cls.expires_at <= datetime.now(timezone.utc).replace(tzinfo=None)
+        ).delete()
+        db.session.commit()
+
+    def __repr__(self):
+        return f'<TotpUsedCode user_id={self.user_id} code={self.code} expires={self.expires_at}>'

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(

app/services/auth_service.py

@@ -177,7 +177,23 @@ def require_jwt(f):
     return decorated
 
 
-# ── Recovery proof helpers (HMAC-nonce) ──────────────────────────────────────
+# ── TOTP replay prevention ────────────────────────────────────────────────────
+
+def is_totp_code_used(user_id: int, code: str) -> bool:
+    """Return True if this TOTP code was already consumed for this user."""
+    from app.models.totp_used_code import TotpUsedCode
+    return TotpUsedCode.is_used(user_id, code)
+
+
+def mark_totp_code_used(user_id: int, code: str) -> None:
+    """
+    Record that this TOTP code was consumed so it cannot be replayed.
+    Caller must commit the session — the surrounding request handler does this.
+    """
+    from app.models.totp_used_code import TotpUsedCode
+    TotpUsedCode.mark_used(user_id, code)
+
+
 
 def generate_recovery_nonce() -> str:
     """

app/static/js/recover.js

@@ -2,9 +2,11 @@
  * recover.js — Account recovery flow
  *
  * Step 1: User provides email + recovery code.
- *   - Fetch recovery data (enc_key_salt, recovery_enc_salt, recovery_iv) from server.
+ *   - Fetch recovery data (recovery_enc_salt, recovery_iv, nonce) from server.
+ *     NOTE: enc_key_salt is NOT returned here — the client must derive it by
+ *     decrypting the recovery blob with the recovery code.
  *   - Derive recovery key from the recovery code (PBKDF2).
- *   - Decrypt enc_key_salt using the recovery key.
+ *   - Decrypt recovery_enc_salt using the recovery key → decrypted enc_key_salt.
  *   - If decryption succeeds, store decrypted enc_key_salt in module state → show step 2.
  *
  * Step 2: User provides new master password.