2323import random
2424import string
2525import json
26+ import time
2627from flask import Blueprint , flash
2728from flask import redirect
2829from flask import render_template
@@ -354,15 +355,98 @@ def oidc_authorise():
354355 return wrap_login_user (user , is_oidc = True )
355356
356357
358+ # MFA hardening constants. Values are deliberately conservative — a legitimate
359+ # user mistypes their token occasionally; an attacker brute-forcing 10^6 TOTP
360+ # codes should not be able to linearly grind them.
361+ _MFA_MAX_ATTEMPTS = 5
362+ _MFA_LOCKOUT_SECONDS = 15 * 60
363+
364+
365+ def _get_pre_mfa_user ():
366+ """Return the user this session just passed password auth for, or None.
367+
368+ The pre_mfa_user_id marker is set exclusively by wrap_login_user after a
369+ successful password (or LDAP) check. Any MFA handler that runs without it
370+ is being hit directly by an attacker and must be refused.
371+ """
372+ pre_mfa_user_id = session .get ('pre_mfa_user_id' )
373+ if pre_mfa_user_id is None :
374+ return None
375+ return get_user (pre_mfa_user_id , id_key = 'id' )
376+
377+
378+ def _clear_pre_mfa_state (preserve_lockout = False ):
379+ """Clear the markers that prove this session just passed password auth.
380+
381+ preserve_lockout: when True, keep the lockout timestamp and fail counter
382+ so that an attacker cannot wipe a brute-force lockout simply by hitting
383+ /login again.
384+ """
385+ session .pop ('pre_mfa_user_id' , None )
386+ session .pop ('pending_mfa_secret' , None )
387+ if not preserve_lockout :
388+ session .pop ('mfa_fail_count' , None )
389+ session .pop ('mfa_lockout_until' , None )
390+
391+
392+ def _mfa_is_locked_out ():
393+ locked_until = session .get ('mfa_lockout_until' )
394+ if locked_until and locked_until > time .time ():
395+ return True
396+ if locked_until and locked_until <= time .time ():
397+ # Lockout expired — reset the counter so the user gets a fresh window.
398+ session .pop ('mfa_lockout_until' , None )
399+ session ['mfa_fail_count' ] = 0
400+ return False
401+
402+
403+ def _register_mfa_failure (user , reason ):
404+ session ['mfa_fail_count' ] = session .get ('mfa_fail_count' , 0 ) + 1
405+ track_activity (
406+ f'Failed MFA { reason } for user { user .user } (attempt { session ["mfa_fail_count" ]} /{ _MFA_MAX_ATTEMPTS } )' ,
407+ ctx_less = True , display_in_ui = False
408+ )
409+ if session ['mfa_fail_count' ] >= _MFA_MAX_ATTEMPTS :
410+ session ['mfa_lockout_until' ] = time .time () + _MFA_LOCKOUT_SECONDS
411+ # Drop the pending-MFA marker so the attacker has to go back through
412+ # password auth before they get another burst of attempts. Preserve
413+ # the lockout timestamp so a fresh /login cannot wipe it.
414+ _clear_pre_mfa_state (preserve_lockout = True )
415+ session .pop ('username' , None )
416+
417+
357418@app .route ("/auth/mfa-setup" , methods = ["GET" , "POST" ])
358419def mfa_setup ():
359- user = retrieve_user_by_username (username = session ["username" ])
420+ user = _get_pre_mfa_user ()
421+ if user is None :
422+ return redirect (url_for ("login.login" ))
423+
424+ # mfa_setup is only for users who haven't completed MFA yet (fresh
425+ # accounts, or admin-reset accounts where mfa_setup_complete was cleared).
426+ # A user who already has MFA enrolled must go through mfa_verify — this
427+ # prevents the "I know the password, let me overwrite the enrolled secret
428+ # with one I control" bypass.
429+ if user .mfa_setup_complete and user .mfa_secrets :
430+ return redirect (url_for ("mfa_verify" ))
431+
432+ if _mfa_is_locked_out ():
433+ flash ('Too many attempts. Please try again later.' , 'danger' )
434+ return redirect (url_for ("login.login" ))
435+
360436 form = MFASetupForm ()
361437
362438 if form .submit () and form .validate ():
439+
363440 token = form .token .data
364- mfa_secret = form .mfa_secret .data
365441 user_password = form .user_password .data
442+ # The secret MUST come from the server-side session, never from the
443+ # submitted form. Trusting form.mfa_secret let an attacker enrol a
444+ # secret of their choosing and immediately log in.
445+ mfa_secret = session .get ('pending_mfa_secret' )
446+ if not mfa_secret :
447+ flash ('MFA setup expired. Please restart the login flow.' , 'danger' )
448+ return redirect (url_for ("login.login" ))
449+
366450 totp = pyotp .TOTP (mfa_secret )
367451
368452 if totp .verify (token ):
@@ -379,36 +463,36 @@ def mfa_setup():
379463 has_valid_password = True
380464
381465 if not has_valid_password :
382- track_activity (
383- f"Failed MFA setup for user { user .user } . Invalid password." ,
384- ctx_less = True ,
385- display_in_ui = False ,
386- )
466+ _register_mfa_failure (user , 'setup (invalid password)' )
387467 flash ("Invalid password. Please try again." , "danger" )
388468 return render_template ("mfa_setup.html" , form = form )
389469
390470 user .mfa_secrets = mfa_secret
391471 user .mfa_setup_complete = True
392472 db .session .commit ()
393- session ["mfa_verified" ] = False
394473 track_activity (
395474 f"MFA setup successful for user { user .user } " ,
396475 ctx_less = True ,
397476 display_in_ui = False ,
398477 )
478+ # Setup succeeded — promote this session to MFA-verified for this
479+ # user and clear the pre-MFA markers. Without this, wrap_login_user
480+ # would loop straight back to mfa_verify.
481+ session ['mfa_verified_for_user_id' ] = user .id
482+ _clear_pre_mfa_state ()
399483 return wrap_login_user (user )
400- track_activity (
401- f"Failed MFA setup for user { user .user } . Invalid token." ,
402- ctx_less = True ,
403- display_in_ui = False ,
404- )
405- flash ("Invalid token or password. Please try again." , "danger" )
484+ else :
485+ _register_mfa_failure (user , 'setup (invalid token)' )
486+ flash ("Invalid token or password. Please try again." , "danger" )
406487
488+ # Generate a fresh secret on every GET and stash it in the session. The
489+ # client only sees the QR code and the base32 key for manual entry; it
490+ # never sends the secret back.
407491 temp_otp_secret = pyotp .random_base32 ()
492+ session ['pending_mfa_secret' ] = temp_otp_secret
408493 otp_uri = pyotp .TOTP (temp_otp_secret ).provisioning_uri (
409494 user .email , issuer_name = "IRIS"
410495 )
411- form .mfa_secret .data = temp_otp_secret
412496 img = qrcode .make (otp_uri )
413497 buf = io .BytesIO ()
414498 img .save (buf , format = "PNG" )
@@ -421,11 +505,10 @@ def mfa_setup():
421505
422506@app .route ("/auth/mfa-verify" , methods = ["GET" , "POST" ])
423507def mfa_verify ():
424- if "username" not in session :
508+ user = _get_pre_mfa_user ()
509+ if user is None :
425510 return redirect (url_for ("login.login" ))
426511
427- user = retrieve_user_by_username (username = session ["username" ])
428-
429512 # Redirect user to MFA setup if MFA is not fully set up
430513 if not user .mfa_secrets or not user .mfa_setup_complete :
431514 track_activity (
@@ -435,6 +518,10 @@ def mfa_verify():
435518 )
436519 return redirect (url_for ("mfa_setup" ))
437520
521+ if _mfa_is_locked_out ():
522+ flash ('Too many attempts. Please try again later.' , 'danger' )
523+ return redirect (url_for ("login.login" ))
524+
438525 form = MFASetupForm ()
439526 form .user_password .data = "not required for verification"
440527
@@ -446,19 +533,19 @@ def mfa_verify():
446533
447534 totp = pyotp .TOTP (user .mfa_secrets )
448535 if totp .verify (token ):
449- session .pop ("username" , None )
450- session ["mfa_verified" ] = True
451536 track_activity (
452537 f"MFA verification successful for user { user .user } " ,
453538 ctx_less = True ,
454539 display_in_ui = False ,
455540 )
541+ # Bind the MFA-verified marker to this specific user id so that a
542+ # later login attempt for a different user on the same browser
543+ # session cannot reuse it.
544+ session ['mfa_verified_for_user_id' ] = user .id
545+ _clear_pre_mfa_state ()
456546 return wrap_login_user (user )
457- track_activity (
458- f"Failed MFA verification for user { user .user } . Invalid token." ,
459- ctx_less = True ,
460- display_in_ui = False ,
461- )
462- flash ("Invalid token. Please try again." , "danger" )
547+ else :
548+ _register_mfa_failure (user , 'verification (invalid token)' )
549+ flash ("Invalid token. Please try again." , "danger" )
463550
464551 return render_template ("mfa_verify.html" , form = form )
0 commit comments