Repository: getmoto/moto (python)

Issue:
If MfaConfiguration is "ON", then software_token_mfa_enabled should be set in verify_software_token
Presently, software_token_mfa_enabled is only set to True in admin_set_user_mfa_preference.

(When MfaConfiguration is "ON", I don't know it it should be set even earlier than verify_software_token, but it should be set at least once token_verified is set to True).

Pull request:
CognitoIDP - improve MFA flows
Fixes #8029

 - Forces MFA setup when MFA_CONFIG is ON, not when it's optional
 - Respond with SOFTWARE_TOKEN_MFA when it's the user preference, or when MFA_CONFIG is ON and user already had a token verified
 - Verifies an auth flow after MFA has been configured

Patch:
diff --git a/moto/cognitoidp/models.py b/moto/cognitoidp/models.py
index 25c8450b3..81a639d8d 100644
--- a/moto/cognitoidp/models.py
+++ b/moto/cognitoidp/models.py
@@ -1389,13 +1389,14 @@ class CognitoIdpBackend(BaseBackend):
         )
 
         return {
+            "ChallengeParameters": {},
             "AuthenticationResult": {
                 "IdToken": id_token,
                 "AccessToken": access_token,
                 "RefreshToken": refresh_token,
                 "ExpiresIn": expires_in,
                 "TokenType": "Bearer",
-            }
+            },
         }
 
     def _validate_auth_flow(
@@ -1592,27 +1593,40 @@ class CognitoIdpBackend(BaseBackend):
             if not timestamp:
                 raise ResourceNotFoundError(timestamp)
 
-            if user.software_token_mfa_enabled:
+            if user.status == UserStatus.FORCE_CHANGE_PASSWORD:
                 return {
-                    "ChallengeName": "SOFTWARE_TOKEN_MFA",
+                    "ChallengeName": "NEW_PASSWORD_REQUIRED",
+                    "ChallengeParameters": {
+                        "USERNAME": username,
+                    },
                     "Session": session,
-                    "ChallengeParameters": {},
                 }
+            if user_pool.mfa_config == "ON" and not user.token_verified:
+                mfas_can_setup = []
+                if user_pool.token_mfa_config == {"Enabled": True}:
+                    mfas_can_setup.append("SOFTWARE_TOKEN_MFA")
+                if user_pool.sms_mfa_config:
+                    mfas_can_setup.append("SMS_MFA")
 
-            if user.sms_mfa_enabled:
                 return {
-                    "ChallengeName": "SMS_MFA",
+                    "ChallengeName": "MFA_SETUP",
+                    "ChallengeParameters": {"MFAS_CAN_SETUP": mfas_can_setup},
+                    "Session": session,
+                }
+            if user.software_token_mfa_enabled or (
+                user_pool.token_mfa_config == {"Enabled": True} and user.token_verified
+            ):
+                return {
+                    "ChallengeName": "SOFTWARE_TOKEN_MFA",
                     "Session": session,
                     "ChallengeParameters": {},
                 }
 
-            if user.status == UserStatus.FORCE_CHANGE_PASSWORD:
+            if user.sms_mfa_enabled:
                 return {
-                    "ChallengeName": "NEW_PASSWORD_REQUIRED",
-                    "ChallengeParameters": {
-                        "USERNAME": username,
-                    },
+                    "ChallengeName": "SMS_MFA",
                     "Session": session,
+                    "ChallengeParameters": {},
                 }
 
             del self.sessions[session]
@@ -1637,8 +1651,7 @@ class CognitoIdpBackend(BaseBackend):
 
         elif challenge_name == "MFA_SETUP":
             username, user_pool = self.sessions[session]
-            auth_params = self._log_user_in(user_pool, client, username)
-            return {"ChallengeParameters": {}, **auth_params}
+            return self._log_user_in(user_pool, client, username)
 
         else:
             return {}
@@ -2012,7 +2025,7 @@ class CognitoIdpBackend(BaseBackend):
             if (
                 user.software_token_mfa_enabled
                 and user.preferred_mfa_setting == "SOFTWARE_TOKEN_MFA"
-            ):
+            ) or (user_pool.mfa_config == "ON" and user.token_verified):
                 return {
                     "ChallengeName": "SOFTWARE_TOKEN_MFA",
                     "ChallengeParameters": {},


Test patch:
diff --git a/tests/test_cognitoidp/test_cognitoidp.py b/tests/test_cognitoidp/test_cognitoidp.py
index 26b0cec20..0e5077a78 100644
--- a/tests/test_cognitoidp/test_cognitoidp.py
+++ b/tests/test_cognitoidp/test_cognitoidp.py
@@ -3071,14 +3071,219 @@ def user_authentication_flow(
     }
 
 
+@cognitoidp_aws_verified(generate_secret=True, with_mfa="ON")
+@pytest.mark.aws_verified
+def test_user_authentication_flow_mfa_on(user_pool=None, user_pool_client=None):
+    conn = boto3.client("cognito-idp", "us-west-2")
+
+    username = str(uuid.uuid4())
+    password = "P2$Sword"
+
+    user_pool_id = user_pool["UserPool"]["Id"]
+    client_id = user_pool_client["UserPoolClient"]["ClientId"]
+    client_secret = user_pool_client["UserPoolClient"]["ClientSecret"]
+
+    secret_hash = pycognito.aws_srp.AWSSRP.get_secret_hash(
+        username=username, client_id=client_id, client_secret=client_secret
+    )
+    conn.sign_up(
+        ClientId=client_id, Username=username, Password=password, SecretHash=secret_hash
+    )
+
+    client_secret = conn.describe_user_pool_client(
+        UserPoolId=user_pool_id, ClientId=client_id
+    )["UserPoolClient"]["ClientSecret"]
+
+    conn.admin_confirm_sign_up(UserPoolId=user_pool_id, Username=username)
+
+    aws_srp = pycognito.aws_srp.AWSSRP(
+        username=username,
+        password=password,
+        pool_id=user_pool_id,
+        client_id=client_id,
+        client_secret=client_secret,
+        client=conn,
+    )
+    auth_params = aws_srp.get_auth_params()
+
+    result = conn.initiate_auth(
+        ClientId=client_id,
+        AuthFlow="USER_SRP_AUTH",
+        AuthParameters=auth_params,
+    )
+
+    challenge_response = aws_srp.process_challenge(
+        result["ChallengeParameters"], auth_params
+    )
+    result = conn.respond_to_auth_challenge(
+        ClientId=client_id,
+        ChallengeName=result["ChallengeName"],
+        ChallengeResponses=challenge_response,
+    )
+    assert result["ChallengeName"] == "MFA_SETUP"
+
+    # add mfa token
+    resp = conn.associate_software_token(Session=result["Session"])
+    secret_code = resp["SecretCode"]
+    totp = pyotp.TOTP(secret_code)
+    user_code = totp.now()
+
+    resp = conn.verify_software_token(Session=resp["Session"], UserCode=user_code)
+    assert resp["Status"] == "SUCCESS"
+
+    # authenticate user once again
+    result = conn.initiate_auth(
+        ClientId=client_id,
+        AuthFlow="USER_SRP_AUTH",
+        AuthParameters={
+            "USERNAME": username,
+            "SRP_A": pycognito.aws_srp.long_to_hex(aws_srp.large_a_value),
+            "SECRET_HASH": secret_hash,
+        },
+    )
+
+    challenge_response = aws_srp.process_challenge(
+        result["ChallengeParameters"], auth_params
+    )
+    result = conn.respond_to_auth_challenge(
+        ClientId=client_id,
+        ChallengeName=result["ChallengeName"],
+        ChallengeResponses=challenge_response,
+    )
+    assert result["ChallengeName"] == "SOFTWARE_TOKEN_MFA"
+
+    # AWS needs a new user code, otherwise they will throw an exception
+    #    ExpiredCodeException: Your software token has already been used once.
+    new_user_code = totp.now()
+    if allow_aws_request():
+        while user_code == new_user_code:
+            new_user_code = totp.now()
+            time.sleep(1)
+    result = conn.respond_to_auth_challenge(
+        ClientId=client_id,
+        Session=result["Session"],
+        ChallengeName=result["ChallengeName"],
+        ChallengeResponses={
+            "SOFTWARE_TOKEN_MFA_CODE": new_user_code,
+            "USERNAME": username,
+            "SECRET_HASH": auth_params["SECRET_HASH"],
+        },
+    )
+    assert result["AuthenticationResult"]["AccessToken"]
+    assert result["AuthenticationResult"]["IdToken"]
+    assert result["AuthenticationResult"]["RefreshToken"]
+
+
 @cognitoidp_aws_verified(generate_secret=True, with_mfa="OPTIONAL")
 @pytest.mark.aws_verified
-def test_user_authentication_flow(user_pool=None, user_pool_client=None):
+def test_user_authentication_flow_mfa_optional(user_pool=None, user_pool_client=None):
     conn = boto3.client("cognito-idp", "us-west-2")
 
-    user_authentication_flow(
-        conn, user_pool=user_pool, user_pool_client=user_pool_client
+    username = str(uuid.uuid4())
+    password = "P2$Sword"
+
+    user_pool_id = user_pool["UserPool"]["Id"]
+    client_id = user_pool_client["UserPoolClient"]["ClientId"]
+    client_secret = user_pool_client["UserPoolClient"]["ClientSecret"]
+
+    secret_hash = pycognito.aws_srp.AWSSRP.get_secret_hash(
+        username=username, client_id=client_id, client_secret=client_secret
+    )
+    conn.sign_up(
+        ClientId=client_id, Username=username, Password=password, SecretHash=secret_hash
+    )
+
+    client_secret = conn.describe_user_pool_client(
+        UserPoolId=user_pool_id, ClientId=client_id
+    )["UserPoolClient"]["ClientSecret"]
+
+    conn.admin_confirm_sign_up(UserPoolId=user_pool_id, Username=username)
+
+    aws_srp = pycognito.aws_srp.AWSSRP(
+        username=username,
+        password=password,
+        pool_id=user_pool_id,
+        client_id=client_id,
+        client_secret=client_secret,
+        client=conn,
     )
+    auth_params = aws_srp.get_auth_params()
+
+    result = conn.initiate_auth(
+        ClientId=client_id,
+        AuthFlow="USER_SRP_AUTH",
+        AuthParameters=auth_params,
+    )
+
+    challenge_response = aws_srp.process_challenge(
+        result["ChallengeParameters"], auth_params
+    )
+    result = conn.respond_to_auth_challenge(
+        ClientId=client_id,
+        ChallengeName=result["ChallengeName"],
+        ChallengeResponses=challenge_response,
+    )
+
+    # add mfa token
+    resp = conn.associate_software_token(
+        AccessToken=result["AuthenticationResult"]["AccessToken"]
+    )
+    secret_code = resp["SecretCode"]
+    totp = pyotp.TOTP(secret_code)
+    user_code = totp.now()
+
+    resp = conn.verify_software_token(
+        AccessToken=result["AuthenticationResult"]["AccessToken"],
+        UserCode=user_code,
+    )
+    assert resp["Status"] == "SUCCESS"
+
+    conn.set_user_mfa_preference(
+        AccessToken=result["AuthenticationResult"]["AccessToken"],
+        SoftwareTokenMfaSettings={"Enabled": True, "PreferredMfa": True},
+    )
+
+    # authenticate user once again
+    result = conn.initiate_auth(
+        ClientId=client_id,
+        AuthFlow="USER_SRP_AUTH",
+        AuthParameters={
+            "USERNAME": username,
+            "SRP_A": pycognito.aws_srp.long_to_hex(aws_srp.large_a_value),
+            "SECRET_HASH": secret_hash,
+        },
+    )
+
+    challenge_response = aws_srp.process_challenge(
+        result["ChallengeParameters"], auth_params
+    )
+    result = conn.respond_to_auth_challenge(
+        ClientId=client_id,
+        ChallengeName=result["ChallengeName"],
+        ChallengeResponses=challenge_response,
+    )
+    assert result["ChallengeName"] == "SOFTWARE_TOKEN_MFA"
+
+    # AWS needs a new user code, otherwise they will throw an exception
+    #    ExpiredCodeException: Your software token has already been used once.
+    new_user_code = totp.now()
+    if allow_aws_request():
+        while user_code == new_user_code:
+            new_user_code = totp.now()
+            time.sleep(1)
+    result = conn.respond_to_auth_challenge(
+        ClientId=client_id,
+        Session=result["Session"],
+        ChallengeName=result["ChallengeName"],
+        ChallengeResponses={
+            "SOFTWARE_TOKEN_MFA_CODE": new_user_code,
+            "USERNAME": username,
+            "SECRET_HASH": auth_params["SECRET_HASH"],
+        },
+    )
+    assert result["AuthenticationResult"]["AccessToken"]
+    assert result["AuthenticationResult"]["IdToken"]
+    assert result["AuthenticationResult"]["RefreshToken"]
 
 
 @mock_aws
@@ -4602,6 +4807,7 @@ def test_initiate_mfa_auth_USER_PASSWORD_AUTH_with_FORCE_CHANGE_PASSWORD_status(
     # 6. Check that the MFA challenge is received
     # 7. Return the MFA challenge
     # 8. Check that the access tokens are received
+    # 9. Login again
 
     client = boto3.client("cognito-idp", "us-west-2")
     username = str(uuid.uuid4())
@@ -4667,6 +4873,31 @@ def test_initiate_mfa_auth_USER_PASSWORD_AUTH_with_FORCE_CHANGE_PASSWORD_status(
     assert result["AuthenticationResult"]["RefreshToken"]
     assert result["AuthenticationResult"]["IdToken"]
 
+    result = client.initiate_auth(
+        ClientId=client_id,
+        AuthFlow="USER_PASSWORD_AUTH",
+        AuthParameters={"USERNAME": username, "PASSWORD": new_password},
+    )
+    assert result["ChallengeName"] == "SOFTWARE_TOKEN_MFA"
+    third_user_code = totp.now()
+    if allow_aws_request():
+        while third_user_code == new_user_code:
+            third_user_code = totp.now()
+            time.sleep(1)
+    result = client.respond_to_auth_challenge(
+        ClientId=client_id,
+        Session=result["Session"],
+        ChallengeName=result["ChallengeName"],
+        ChallengeResponses={
+            "SOFTWARE_TOKEN_MFA_CODE": third_user_code,
+            "USERNAME": username,
+        },
+    )
+    assert result["ChallengeParameters"] == {}
+    assert result["AuthenticationResult"]["AccessToken"]
+    assert result["AuthenticationResult"]["RefreshToken"]
+    assert result["AuthenticationResult"]["IdToken"]
+
 
 @mock_aws
 def test_initiate_auth_USER_PASSWORD_AUTH_user_not_found():