feat(cognito-idp): make TOTP MFA work

* Put working TOTP MFA behind a feature flag - specifically
`MOTO_COGNITO_IDP_USER_POOL_ENABLE_TOTP`
* Support the FriendlyDeviceName field, although only for error
  reporting to end-user

Author: bdellegrazie

moto/cognitoidp/exceptions.py

@@ -59,3 +59,8 @@ def __init__(self) -> None:
             error_type="InvalidPasswordException",
             message="The provided password does not confirm to the configured password policy",
         )
+
+
+class CodeMismatchException(JsonRESTError):
+    def __init__(self, message: str):
+        super().__init__(error_type="CodeMismatchException", message=message)

moto/cognitoidp/models.py

@@ -2,8 +2,9 @@
 import re
 import time
 from collections import OrderedDict
-from typing import Any, Optional
+from typing import Any, Final, Optional
 
+from cryptography.hazmat.primitives.twofactor import InvalidToken
 from joserfc import jwk, jwt
 
 from moto.core.base_backend import BackendDict, BaseBackend
@@ -15,10 +16,12 @@
 
 from ..settings import (
     get_cognito_idp_user_pool_client_id_strategy,
+    get_cognito_idp_user_pool_enable_totp,
     get_cognito_idp_user_pool_id_strategy,
 )
 from .exceptions import (
     AliasExistsException,
+    CodeMismatchException,
     ExpiredCodeException,
     GroupExistsException,
     InvalidParameterException,
@@ -39,6 +42,9 @@
     validate_username_format,
 )
 
+# FIXME: Should be per user and stored in the user's profile
+COGNITO_TOTP_MFA_SECRET: Final[str] = "asdfasdfasdf"
+
 
 class UserStatus(str, enum.Enum):
     FORCE_CHANGE_PASSWORD = "FORCE_CHANGE_PASSWORD"
@@ -970,8 +976,12 @@ class CognitoIdpBackend(BaseBackend):
     In some cases, you need to have reproducible IDs for the user pool.
     For example, a single initialization before the start of integration tests.
 
-    This behavior can be enabled by passing the environment variable: MOTO_COGNITO_IDP_USER_POOL_ID_STRATEGY=HASH.
-    Passing MOTO_COGNITO_IDP_USER_POOL_CLIENT_ID_STRATEGY=HASH enables the same logic for user pool clients.
+    This behavior can be enabled by passing the environment variable: `MOTO_COGNITO_IDP_USER_POOL_ID_STRATEGY=HASH`.
+    Passing `MOTO_COGNITO_IDP_USER_POOL_CLIENT_ID_STRATEGY=HASH` enables the same logic for user pool clients.
+
+    Support for MFA TOTP can be enabled by setting `MOTO_COGNITO_IDP_USER_POOL_ENABLE_TOTP=true`.
+    Moto will validate the TOTP MFA provided by the user when registering MFA or subsequently authenticating.
+    At this time, Moto uses a single fixed secret across all users.
     """
 
     def __init__(self, region_name: str, account_id: str):
@@ -1720,6 +1730,16 @@ def respond_to_auth_challenge(
                 ):
                     raise NotAuthorizedError(secret_hash)
 
+            if (
+                challenge_name == "SOFTWARE_TOKEN_MFA"
+                and get_cognito_idp_user_pool_enable_totp()
+            ):
+                totp = cognito_totp(COGNITO_TOTP_MFA_SECRET)
+                try:
+                    totp.verify(mfa_code.encode("utf-8"), int(time.time()))
+                except InvalidToken:
+                    raise CodeMismatchException("MFA Code Mismatch")
+
             del self.sessions[session]
             return self._log_user_in(user_pool, client, username)
 
@@ -2173,7 +2193,7 @@ def initiate_auth(
     def associate_software_token(
         self, access_token: str, session: str
     ) -> dict[str, str]:
-        secret_code = "asdfasdfasdf"
+        secret_code = COGNITO_TOTP_MFA_SECRET
         if session:
             if session in self.sessions:
                 return {"SecretCode": secret_code, "Session": session}
@@ -2189,20 +2209,23 @@ def associate_software_token(
         raise NotAuthorizedError(access_token)
 
     def verify_software_token(
-        self, access_token: str, session: str, user_code: str
+        self, access_token: str, session: str, user_code: str, friendly_device_name: str
     ) -> dict[str, str]:
-        """
-        The parameter UserCode has not yet been implemented
-        """
-        totp = cognito_totp("asdfasdfasdf")
+        totp = cognito_totp(COGNITO_TOTP_MFA_SECRET)
         if session:
             if session not in self.sessions:
                 raise ResourceNotFoundError(session)
 
             username, user_pool = self.sessions[session]
             user = self.admin_get_user(user_pool.id, username)
 
-            totp.verify(user_code.encode("utf-8"), int(time.time()))
+            if get_cognito_idp_user_pool_enable_totp():
+                try:
+                    totp.verify(user_code.encode("utf-8"), int(time.time()))
+                except InvalidToken:
+                    raise CodeMismatchException(
+                        f"Code mismatch ({friendly_device_name})"
+                    )
 
             user.token_verified = True
 
@@ -2215,7 +2238,13 @@ def verify_software_token(
                 _, username = user_pool.access_tokens[access_token]
                 user = self.admin_get_user(user_pool.id, username)
 
-                totp.verify(user_code.encode("utf-8"), int(time.time()))
+                if get_cognito_idp_user_pool_enable_totp():
+                    try:
+                        totp.verify(user_code.encode("utf-8"), int(time.time()))
+                    except InvalidToken:
+                        raise CodeMismatchException(
+                            f"Code mismatch ({friendly_device_name})"
+                        )
 
                 user.token_verified = True
 
@@ -2463,10 +2492,16 @@ def associate_software_token(
         return backend.associate_software_token(access_token, session)
 
     def verify_software_token(
-        self, access_token: str, session: str, user_code: str
+        self,
+        access_token: str,
+        session: str,
+        user_code: str,
+        friendly_device_name: str,
     ) -> dict[str, str]:
         backend = self._find_backend_by_access_token_or_session(access_token, session)
-        return backend.verify_software_token(access_token, session, user_code)
+        return backend.verify_software_token(
+            access_token, session, user_code, friendly_device_name
+        )
 
     def set_user_mfa_preference(
         self,

moto/cognitoidp/responses.py

@@ -586,8 +586,9 @@ def verify_software_token(self) -> ActionResult:
         access_token = self._get_param("AccessToken")
         session = self._get_param("Session")
         user_code = self._get_param("UserCode")
+        friendly_device_name = self._get_param("FriendlyDeviceName")
         result = self._get_region_agnostic_backend().verify_software_token(
-            access_token, session, user_code
+            access_token, session, user_code, friendly_device_name
         )
         return ActionResult(result)
 

moto/settings.py

@@ -190,6 +190,15 @@ def get_cognito_idp_user_pool_client_id_strategy() -> Optional[str]:
     return os.environ.get("MOTO_COGNITO_IDP_USER_POOL_CLIENT_ID_STRATEGY")
 
 
+def get_cognito_idp_user_pool_enable_totp() -> bool:
+    return (
+        os.environ.get(
+            key="MOTO_COGNITO_IDP_USER_POOL_ENABLE_TOTP", default="false"
+        ).lower()
+        == "true"
+    )
+
+
 def enable_iso_regions() -> bool:
     return os.environ.get("MOTO_ENABLE_ISO_REGIONS", "false").lower() == "true"
 

tests/test_cognitoidp/test_cognitoidp.py

@@ -3294,6 +3294,7 @@ def user_authentication_flow(
 
 
 @cognitoidp_aws_verified(generate_secret=True, with_mfa="ON")
+@mock.patch.dict(os.environ, {"MOTO_COGNITO_IDP_USER_POOL_ENABLE_TOTP": "true"})
 @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")
@@ -3397,6 +3398,7 @@ def test_user_authentication_flow_mfa_on(user_pool=None, user_pool_client=None):
 
 
 @cognitoidp_aws_verified(generate_secret=True, with_mfa="OPTIONAL")
+@mock.patch.dict(os.environ, {"MOTO_COGNITO_IDP_USER_POOL_ENABLE_TOTP": "true"})
 @pytest.mark.aws_verified
 def test_user_authentication_flow_mfa_optional(user_pool=None, user_pool_client=None):
     conn = boto3.client("cognito-idp", "us-west-2")
@@ -5019,6 +5021,7 @@ def test_initiate_auth_USER_PASSWORD_AUTH_with_FORCE_CHANGE_PASSWORD_status():
 
 
 @cognitoidp_aws_verified(explicit_auth_flows=["USER_PASSWORD_AUTH"], with_mfa="ON")
+@mock.patch.dict(os.environ, {"MOTO_COGNITO_IDP_USER_POOL_ENABLE_TOTP": "true"})
 @pytest.mark.aws_verified
 def test_initiate_mfa_auth_USER_PASSWORD_AUTH_with_FORCE_CHANGE_PASSWORD_status(
     user_pool=None, user_pool_client=None
@@ -5241,6 +5244,7 @@ def test_initiate_auth_with_invalid_secret_hash():
 
 
 @mock_aws
+@mock.patch.dict(os.environ, {"MOTO_COGNITO_IDP_USER_POOL_ENABLE_TOTP": "true"})
 def test_setting_mfa():
     conn = boto3.client("cognito-idp", "us-west-2")
 
@@ -5329,6 +5333,7 @@ def test_admin_setting_single_mfa():
 
 
 @mock_aws
+@mock.patch.dict(os.environ, {"MOTO_COGNITO_IDP_USER_POOL_ENABLE_TOTP": "true"})
 def test_admin_setting_mfa_totp_and_sms():
     conn = boto3.client("cognito-idp", "us-west-2")
 
@@ -5367,7 +5372,7 @@ def test_admin_setting_mfa_totp_and_sms():
 
 
 @mock_aws
-def test_admin_initiate_auth_when_token_totp_enabled():
+def test_admin_initiate_auth_when_token_totp_masked():
     conn = boto3.client("cognito-idp", "us-west-2")
 
     result = authentication_flow(conn, "ADMIN_NO_SRP_AUTH")
@@ -5425,6 +5430,123 @@ def test_admin_initiate_auth_when_token_totp_enabled():
     assert result["AuthenticationResult"]["TokenType"] == "Bearer"
 
 
+@mock_aws
+@mock.patch.dict(os.environ, {"MOTO_COGNITO_IDP_USER_POOL_ENABLE_TOTP": "true"})
+def test_admin_initiate_auth_when_token_totp_enabled():
+    conn = boto3.client("cognito-idp", "us-west-2")
+
+    result = authentication_flow(conn, "ADMIN_NO_SRP_AUTH")
+    access_token = result["access_token"]
+    user_pool_id = result["user_pool_id"]
+    username = result["username"]
+    client_id = result["client_id"]
+    password = result["password"]
+    resp = conn.associate_software_token(AccessToken=access_token)
+    secret_code = resp["SecretCode"]
+    totp = pyotp.TOTP(secret_code)
+    user_code = totp.now()
+    conn.verify_software_token(AccessToken=access_token, UserCode=user_code)
+
+    # Set MFA TOTP and SMS methods
+    conn.admin_set_user_mfa_preference(
+        Username=username,
+        UserPoolId=user_pool_id,
+        SoftwareTokenMfaSettings={"Enabled": True, "PreferredMfa": True},
+        SMSMfaSettings={"Enabled": True, "PreferredMfa": False},
+    )
+    result = conn.admin_get_user(UserPoolId=user_pool_id, Username=username)
+    assert len(result["UserMFASettingList"]) == 2
+    assert result["PreferredMfaSetting"] == "SOFTWARE_TOKEN_MFA"
+
+    # Initiate auth with TOTP
+    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"] == "SOFTWARE_TOKEN_MFA"
+    assert result["Session"] != ""
+
+    # Respond to challenge with TOTP
+    result = conn.admin_respond_to_auth_challenge(
+        UserPoolId=user_pool_id,
+        ClientId=client_id,
+        ChallengeName="SOFTWARE_TOKEN_MFA",
+        Session=result["Session"],
+        ChallengeResponses={
+            "SOFTWARE_TOKEN_MFA_CODE": totp.now(),
+            "USERNAME": username,
+        },
+    )
+
+    assert result["AuthenticationResult"]["IdToken"] != ""
+    assert result["AuthenticationResult"]["AccessToken"] != ""
+    assert result["AuthenticationResult"]["RefreshToken"] != ""
+    assert result["AuthenticationResult"]["TokenType"] == "Bearer"
+
+
+@mock_aws
+@mock.patch.dict(os.environ, {"MOTO_COGNITO_IDP_USER_POOL_ENABLE_TOTP": "true"})
+def test_admin_initiate_auth_when_token_totp_enabled_invalid():
+    conn = boto3.client("cognito-idp", "us-west-2")
+
+    result = authentication_flow(conn, "ADMIN_NO_SRP_AUTH")
+    access_token = result["access_token"]
+    user_pool_id = result["user_pool_id"]
+    username = result["username"]
+    client_id = result["client_id"]
+    password = result["password"]
+    resp = conn.associate_software_token(AccessToken=access_token)
+    secret_code = resp["SecretCode"]
+    totp = pyotp.TOTP(secret_code)
+    user_code = totp.now()
+    conn.verify_software_token(AccessToken=access_token, UserCode=user_code)
+
+    # Set MFA TOTP and SMS methods
+    conn.admin_set_user_mfa_preference(
+        Username=username,
+        UserPoolId=user_pool_id,
+        SoftwareTokenMfaSettings={"Enabled": True, "PreferredMfa": True},
+        SMSMfaSettings={"Enabled": True, "PreferredMfa": False},
+    )
+    result = conn.admin_get_user(UserPoolId=user_pool_id, Username=username)
+    assert len(result["UserMFASettingList"]) == 2
+    assert result["PreferredMfaSetting"] == "SOFTWARE_TOKEN_MFA"
+
+    # Initiate auth with TOTP
+    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"] == "SOFTWARE_TOKEN_MFA"
+    assert result["Session"] != ""
+
+    with pytest.raises(ClientError) as exc:
+        result = conn.admin_respond_to_auth_challenge(
+            UserPoolId=user_pool_id,
+            ClientId=client_id,
+            ChallengeName="SOFTWARE_TOKEN_MFA",
+            Session=result["Session"],
+            ChallengeResponses={
+                "SOFTWARE_TOKEN_MFA_CODE": "123456",
+                "USERNAME": username,
+            },
+        )
+        err = exc.value.response["Error"]
+        assert err["Code"] == "CodeMismatch"
+
+
 @mock_aws
 def test_admin_initiate_auth_when_sms_mfa_enabled():
     conn = boto3.client("cognito-idp", "us-west-2")

tests/test_cognitoidp/test_cognitoidp_utils.py

@@ -2,13 +2,13 @@
 
 import pyotp
 
+from moto.cognitoidp.models import COGNITO_TOTP_MFA_SECRET
 from moto.cognitoidp.utils import cognito_totp
 
 
 def test_cognito_totp():
-    key = "asdfasdfasdf"
-    client_totp = pyotp.TOTP(s=key)
-    internal_totp = cognito_totp(key)
+    client_totp = pyotp.TOTP(s=COGNITO_TOTP_MFA_SECRET)
+    internal_totp = cognito_totp(COGNITO_TOTP_MFA_SECRET)
 
     client_code = client_totp.now()
     internal_code = internal_totp.generate(int(time.time())).decode("utf-8")