feat(cognito-idp): make TOTP MFA work (#9785)

Author: bdellegrazie

.github/workflows/tests_servermode.yml

@@ -21,7 +21,7 @@ jobs:
       run: |
         pip install build
         python -m build
-        docker run --rm -t --name motoserver -e TEST_SERVER_MODE=true -e MOTO_EC2_LOAD_DEFAULT_AMIS=false -e AWS_SECRET_ACCESS_KEY=server_secret -e AWS_ACCESS_KEY_ID=server_key -v `pwd`:/moto -p 5000:5000 -v /var/run/docker.sock:/var/run/docker.sock python:${{ matrix.python-version }}-slim /moto/scripts/ci_moto_server.sh &
+        docker run --rm -t --name motoserver -e TEST_SERVER_MODE=true -e MOTO_EC2_LOAD_DEFAULT_AMIS=false -e MOTO_COGNITO_IDP_USER_POOL_ENABLE_TOTP=true -e AWS_SECRET_ACCESS_KEY=server_secret -e AWS_ACCESS_KEY_ID=server_key -v `pwd`:/moto -p 5000:5000 -v /var/run/docker.sock:/var/run/docker.sock python:${{ matrix.python-version }}-slim /moto/scripts/ci_moto_server.sh &
         python scripts/ci_wait_for_server.py
     - name: Get pip cache dir
       id: pip-cache

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,
@@ -32,12 +35,16 @@
 from .utils import (
     PAGINATION_MODEL,
     check_secret_hash,
+    cognito_totp,
     expand_attrs,
     flatten_attrs,
     generate_id,
     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 +977,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 +1731,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 +2194,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}
@@ -2188,16 +2209,25 @@ def associate_software_token(
 
         raise NotAuthorizedError(access_token)
 
-    def verify_software_token(self, access_token: str, session: str) -> dict[str, str]:
-        """
-        The parameter UserCode has not yet been implemented
-        """
+    def verify_software_token(
+        self, access_token: str, session: str, user_code: str, friendly_device_name: str
+    ) -> dict[str, str]:
+        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)
+
+            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
 
             session = str(random.uuid4())
@@ -2209,6 +2239,14 @@ def verify_software_token(self, access_token: str, session: str) -> dict[str, st
                 _, username = user_pool.access_tokens[access_token]
                 user = self.admin_get_user(user_pool.id, username)
 
+                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
 
                 session = str(random.uuid4())
@@ -2454,9 +2492,17 @@ def associate_software_token(
         backend = self._find_backend_by_access_token_or_session(access_token, session)
         return backend.associate_software_token(access_token, session)
 
-    def verify_software_token(self, access_token: str, session: str) -> dict[str, str]:
+    def verify_software_token(
+        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)
+        return backend.verify_software_token(
+            access_token, session, user_code, friendly_device_name
+        )
 
     def set_user_mfa_preference(
         self,

moto/cognitoidp/responses.py

@@ -585,8 +585,10 @@ def associate_software_token(self) -> ActionResult:
     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
+            access_token, session, user_code, friendly_device_name
         )
         return ActionResult(result)
 

moto/cognitoidp/utils.py

@@ -5,6 +5,9 @@
 import string
 from typing import Any, Optional
 
+from cryptography.hazmat.primitives.hashes import SHA1
+from cryptography.hazmat.primitives.twofactor.totp import TOTP
+
 from moto.moto_api._internal import mock_random as random
 
 FORMATS = {
@@ -120,3 +123,19 @@ def _generate_id_hash(args: Any) -> str:
         hasher.update(str(arg).encode())
 
     return hasher.hexdigest()
+
+
+def cognito_totp(key: str) -> TOTP:
+    key_padded = key
+    # Pad the secret if required before converting it to bytes
+    padding = len(key) % 8
+    if padding != 0:
+        key_padded += "=" * (8 - padding)
+    # https://docs.aws.amazon.com/cognito/latest/developerguide/user-pool-settings-mfa-totp.html
+    return TOTP(
+        key=base64.b32decode(key_padded, casefold=True),
+        length=6,
+        algorithm=SHA1(),
+        time_step=30,
+        enforce_key_length=False,
+    )

moto/settings.py

@@ -190,6 +190,13 @@ 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("MOTO_COGNITO_IDP_USER_POOL_ENABLE_TOTP", "false").lower()
+        == "true"
+    )
+
+
 def enable_iso_regions() -> bool:
     return os.environ.get("MOTO_ENABLE_ISO_REGIONS", "false").lower() == "true"