From 25bc983df0a3006d5ce354596190dd173decb03b Mon Sep 17 00:00:00 2001
From: Ari Freyr Asgeirsson <haugurjr@gmail.com>
Date: Thu, 22 Feb 2024 20:44:52 +0000
Subject: [PATCH] CognitoIDP: Add `SMS_MFA` support to admin_initiate_auth and
 respond_to_auth_challenge methods (#7379)

---
 moto/cognitoidp/models.py                | 18 +++++++--
 tests/test_cognitoidp/test_cognitoidp.py | 50 ++++++++++++++++++++++++
 2 files changed, 64 insertions(+), 4 deletions(-)

diff --git a/moto/cognitoidp/models.py b/moto/cognitoidp/models.py
index 15416d4e7..0e960764e 100644
--- a/moto/cognitoidp/models.py
+++ b/moto/cognitoidp/models.py
@@ -1466,6 +1466,16 @@ class CognitoIdpBackend(BaseBackend):
                     "Session": session,
                 }
 
+            if user.sms_mfa_enabled and user.preferred_mfa_setting == "SMS_MFA":
+                session = str(random.uuid4())
+                self.sessions[session] = user_pool
+
+                return {
+                    "ChallengeName": "SMS_MFA",
+                    "ChallengeParameters": {},
+                    "Session": session,
+                }
+
             return self._log_user_in(user_pool, client, username)
         elif auth_flow in (AuthFlow.REFRESH_TOKEN, AuthFlow.REFRESH_TOKEN_AUTH):
             refresh_token: str = auth_parameters.get("REFRESH_TOKEN")  # type: ignore[assignment]
@@ -1578,13 +1588,13 @@ class CognitoIdpBackend(BaseBackend):
 
             del self.sessions[session]
             return self._log_user_in(user_pool, client, username)
-        elif challenge_name == "SOFTWARE_TOKEN_MFA":
+        elif challenge_name == "SOFTWARE_TOKEN_MFA" or challenge_name == "SMS_MFA":
             username: str = challenge_responses.get("USERNAME")  # type: ignore[no-redef]
             self.admin_get_user(user_pool.id, username)
 
-            software_token_mfa_code = challenge_responses.get("SOFTWARE_TOKEN_MFA_CODE")
-            if not software_token_mfa_code:
-                raise ResourceNotFoundError(software_token_mfa_code)
+            mfa_code = challenge_responses.get(f"{challenge_name}_CODE")
+            if not mfa_code:
+                raise ResourceNotFoundError(mfa_code)
 
             if client.generate_secret:
                 secret_hash = challenge_responses.get("SECRET_HASH")
diff --git a/tests/test_cognitoidp/test_cognitoidp.py b/tests/test_cognitoidp/test_cognitoidp.py
index 8867282e4..1ff151bfb 100644
--- a/tests/test_cognitoidp/test_cognitoidp.py
+++ b/tests/test_cognitoidp/test_cognitoidp.py
@@ -4454,6 +4454,56 @@ def test_admin_initiate_auth_when_token_totp_enabled():
     assert result["AuthenticationResult"]["TokenType"] == "Bearer"
 
 
+@mock_aws
+def test_admin_initiate_auth_when_sms_mfa_enabled():
+    conn = boto3.client("cognito-idp", "us-west-2")
+
+    result = authentication_flow(conn, "ADMIN_NO_SRP_AUTH")
+    user_pool_id = result["user_pool_id"]
+    username = result["username"]
+    client_id = result["client_id"]
+    password = result["password"]
+
+    # Set MFA SMS methods
+    conn.admin_set_user_mfa_preference(
+        Username=username,
+        UserPoolId=user_pool_id,
+        SMSMfaSettings={"Enabled": True, "PreferredMfa": True},
+    )
+    result = conn.admin_get_user(UserPoolId=user_pool_id, Username=username)
+    assert len(result["UserMFASettingList"]) == 1
+    assert result["PreferredMfaSetting"] == "SMS_MFA"
+
+    result = conn.admin_initiate_auth(
+        UserPoolId=user_pool_id,
+        ClientId=client_id,
+        AuthFlow="ADMIN_NO_SRP_AUTH",
+        AuthParameters={
+            "USERNAME": username,
+            "PASSWORD": password,
+        },
+    )
+
+    assert result["ChallengeName"] == "SMS_MFA"
+    assert result["Session"] != ""
+
+    result = conn.admin_respond_to_auth_challenge(
+        UserPoolId=user_pool_id,
+        ClientId=client_id,
+        ChallengeName="SMS_MFA",
+        Session=result["Session"],
+        ChallengeResponses={
+            "SMS_MFA_CODE": "123456",
+            "USERNAME": username,
+        },
+    )
+
+    assert result["AuthenticationResult"]["IdToken"] != ""
+    assert result["AuthenticationResult"]["AccessToken"] != ""
+    assert result["AuthenticationResult"]["RefreshToken"] != ""
+    assert result["AuthenticationResult"]["TokenType"] == "Bearer"
+
+
 @mock_aws
 def test_admin_setting_mfa_when_token_not_verified():
     conn = boto3.client("cognito-idp", "us-west-2")