Merge ba94a58 into 9b2d86a

flask_security/forms.py

@@ -436,7 +436,7 @@ def validate(self):
         # verify entered token with user's totp secret
         if not _security._totp_factory.verify_totp(
             token=self.code.data,
-            totp_secret=self.user.tf_totp_secret,
+            totp_secret=self.tf_totp_secret,
             user=self.user,
             window=self.window,
         ):

flask_security/twofactor.py

@@ -48,6 +48,7 @@ def tf_clean_session():
             "tf_primary_method",
             "tf_confirmed",
             "tf_remember_login",
+            "tf_totp_secret",
         ]:
             session.pop(k, None)
 
@@ -86,16 +87,23 @@ def send_security_token(user, method, totp_secret):
     )
 
 
-def complete_two_factor_process(user, primary_method, is_changing, remember_login=None):
+def complete_two_factor_process(
+    user, primary_method, totp_secret, is_changing, remember_login=None
+):
     """clean session according to process (login or changing two-factor method)
      and perform action accordingly
     """
 
-    # only update primary_method and DB if necessary
+    # update changed primary_method
     if user.tf_primary_method != primary_method:
         user.tf_primary_method = primary_method
         _datastore.put(user)
 
+    # update changed totp_secret
+    if user.tf_totp_secret != totp_secret:
+        user.tf_totp_secret = totp_secret
+        _datastore.put(user)
+
     # if we are changing two-factor method
     if is_changing:
         completion_message = "TWO_FACTOR_CHANGE_METHOD_SUCCESSFUL"

flask_security/views.py

@@ -701,7 +701,7 @@ def two_factor_setup():
         # Regenerate the TOTP secret on every call of 2FA setup unless it is
         # within the same session and method (e.g. upon entering the phone number)
         if pm != session.get("tf_primary_method", None):
-            user.tf_totp_secret = _security._totp_factory.generate_totp_secret()
+            session["tf_totp_secret"] = _security._totp_factory.generate_totp_secret()
 
         session["tf_primary_method"] = pm
         session["tf_state"] = "validating_profile"
@@ -710,7 +710,7 @@ def two_factor_setup():
         _datastore.put(user)
         after_this_request(_commit)
 
-        send_security_token(user=user, method=pm, totp_secret=user.tf_totp_secret)
+        send_security_token(user=user, method=pm, totp_secret=session["tf_totp_secret"])
         code_form = _security.two_factor_verify_code_form()
         if not _security._want_json(request):
             return _security.render_template(
@@ -783,8 +783,10 @@ def two_factor_token_validation():
 
         if session["tf_state"] == "ready":
             pm = user.tf_primary_method
+            totp_secret = user.tf_totp_secret
         else:
             pm = session["tf_primary_method"]
+            totp_secret = session["tf_totp_secret"]
     else:
         if (
             not all(
@@ -797,13 +799,15 @@ def two_factor_token_validation():
             logout_user()
             return _tf_illegal_state(form, _security.login_url)
         pm = session["tf_primary_method"]
+        totp_secret = session["tf_totp_secret"]
         form.user = current_user
 
     setattr(form, "primary_method", pm)
+    setattr(form, "tf_totp_secret", totp_secret)
     if form.validate_on_submit():
         # Success - log in user and clear all session variables
         completion_message = complete_two_factor_process(
-            form.user, pm, changing, session.pop("tf_remember_login", None)
+            form.user, pm, totp_secret, changing, session.pop("tf_remember_login", None)
         )
         after_this_request(_commit)
         if not request.is_json:
@@ -960,6 +964,8 @@ def two_factor_qrcode():
         return abort(404)
 
     totp = user.tf_totp_secret
+    if "tf_totp_secret" in session:
+        totp = session["tf_totp_secret"]
     try:
         import pyqrcode
 

tests/test_two_factor.py

@@ -265,6 +265,11 @@ def test_two_factor_flag(app, client):
     assert username == "gal@lp.com"
     assert b"svg" in qrcode_page_response.data
 
+    # Generate token from passed totp_secret and confirm setup
+    code = app.security._totp_factory.generate_totp_password(totp_secret)
+    response = client.post("/tf-validate", data=dict(code=code), follow_redirects=True)
+    assert b"You successfully changed your two-factor method" in response.data
+
     logout(client)
 
     # Test login with remember_token
@@ -617,11 +622,15 @@ def on_identity_changed(app, identity):
     response = client.post("/tf-setup", data=data, follow_redirects=True)
     assert b"To Which Phone Number Should We Send Code To" in response.data
 
-    # Retrieve the current totp secret from the user in datastore for later comparison
-    with app.app_context():
-        user = app.security.datastore.find_user(email="jill@lp.com")
-        assert "enckey" in user.tf_totp_secret
-        old_secret = user.tf_totp_secret
+    # Retrieve the currently generated totp secret for later comparison
+    session = get_session(response)
+    if "tf_totp_secret" in session:
+        generated_secret = session["tf_totp_secret"]
+    else:
+        with app.app_context():
+            user = app.security.datastore.find_user(email="jill@lp.com")
+            generated_secret = user.tf_totp_secret
+    assert "enckey" in generated_secret
 
     # Send the phone number in the second step, method remains unchanged
     data = dict(setup="sms", phone="+111111111111")
@@ -636,7 +645,7 @@ def on_identity_changed(app, identity):
     # Retrieve the final totp secret and make sure it matches the previous one
     with app.app_context():
         user = app.security.datastore.find_user(email="jill@lp.com")
-        assert old_secret == user.tf_totp_secret
+        assert generated_secret == user.tf_totp_secret
 
     # Finally opt back out and check that tf_totp_secret is None
     response = client.get("/tf-setup", data=data, follow_redirects=True)