79 lines
2.9 KiB
Python
79 lines
2.9 KiB
Python
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}>'
|