feat(cognito-idp): make TOTP MFA work

Implement TOTP MFA using the existing cryptography API.
This allows clients to behave correctly and use proper TOTP validation
compared with previously where the UserCode was not handled.

Warning: this implementation still uses the fixed constant "secret" and
therefore should not be considered for anything other than testing.

A solution implementing "real" TOTP was considered but would need
storing the secret temporarily in the session during the auth process
but the session object is unfortunately a tuple at this time.

Author: bdellegrazie

moto/cognitoidp/models.py

@@ -32,6 +32,7 @@
 from .utils import (
     PAGINATION_MODEL,
     check_secret_hash,
+    cognito_totp,
     expand_attrs,
     flatten_attrs,
     generate_id,
@@ -2187,16 +2188,22 @@ def associate_software_token(
 
         raise NotAuthorizedError(access_token)
 
-    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
+    ) -> dict[str, str]:
         """
         The parameter UserCode has not yet been implemented
         """
+        totp = cognito_totp("asdfasdfasdf")
         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()))
+
             user.token_verified = True
 
             session = str(random.uuid4())
@@ -2208,6 +2215,8 @@ 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)
 
+                totp.verify(user_code.encode("utf-8"), int(time.time()))
+
                 user.token_verified = True
 
                 session = str(random.uuid4())
@@ -2453,9 +2462,11 @@ 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
+    ) -> 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)
 
     def set_user_mfa_preference(
         self,

moto/cognitoidp/responses.py

@@ -585,8 +585,9 @@ 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")
         result = self._get_region_agnostic_backend().verify_software_token(
-            access_token, session
+            access_token, session, user_code
         )
         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,
+    )

tests/test_cognitoidp/test_cognitoidp.py

@@ -5248,9 +5248,13 @@ def test_setting_mfa():
         result = authentication_flow(conn, auth_flow)
 
         # Set MFA method
-        conn.associate_software_token(AccessToken=result["access_token"])
+        resp = conn.associate_software_token(AccessToken=result["access_token"])
+        secret_code = resp["SecretCode"]
+        totp = pyotp.TOTP(secret_code)
+        user_code = totp.now()
+
         conn.verify_software_token(
-            AccessToken=result["access_token"], UserCode="123456"
+            AccessToken=result["access_token"], UserCode=user_code
         )
         conn.set_user_mfa_preference(
             AccessToken=result["access_token"],
@@ -5332,8 +5336,12 @@ def test_admin_setting_mfa_totp_and_sms():
     access_token = result["access_token"]
     user_pool_id = result["user_pool_id"]
     username = result["username"]
-    conn.associate_software_token(AccessToken=access_token)
-    conn.verify_software_token(AccessToken=access_token, UserCode="123456")
+    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(
@@ -5368,8 +5376,11 @@ def test_admin_initiate_auth_when_token_totp_enabled():
     username = result["username"]
     client_id = result["client_id"]
     password = result["password"]
-    conn.associate_software_token(AccessToken=access_token)
-    conn.verify_software_token(AccessToken=access_token, UserCode="123456")
+    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(

tests/test_cognitoidp/test_cognitoidp_utils.py

@@ -0,0 +1,16 @@
+import time
+
+import pyotp
+
+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_code = client_totp.now()
+    internal_code = internal_totp.generate(int(time.time())).decode("utf-8")
+
+    assert internal_code == client_code