05/18 Enhanced codes and functionalities

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}>'