From 9a84423187944fdec55c97d31f9c372e7c2c651a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jo=C3=A3o=20Pedro=20Schmitt?= Date: Tue, 26 Sep 2023 10:45:40 -0700 Subject: [PATCH] CognitoIDP: clean MFA settings capability and enable multiple MFA methods (#6853) --- moto/cognitoidp/models.py | 55 ++++++++++--------- tests/test_cognitoidp/test_cognitoidp.py | 67 ++++++++++++++++++++++-- 2 files changed, 90 insertions(+), 32 deletions(-) diff --git a/moto/cognitoidp/models.py b/moto/cognitoidp/models.py index 2cccdec56..1e9c82ec8 100644 --- a/moto/cognitoidp/models.py +++ b/moto/cognitoidp/models.py @@ -833,7 +833,7 @@ class CognitoIdpUser(BaseModel): user_mfa_setting_list = [] if self.software_token_mfa_enabled: user_mfa_setting_list.append("SOFTWARE_TOKEN_MFA") - elif self.sms_mfa_enabled: + if self.sms_mfa_enabled: user_mfa_setting_list.append("SMS_MFA") user_json = self._base_json() if extended: @@ -1981,26 +1981,13 @@ class CognitoIdpBackend(BaseBackend): for user_pool in self.user_pools.values(): if access_token in user_pool.access_tokens: _, username = user_pool.access_tokens[access_token] - user = self.admin_get_user(user_pool.id, username) - if software_token_mfa_settings and software_token_mfa_settings.get( - "Enabled" - ): - if user.token_verified: - user.software_token_mfa_enabled = True - else: - raise InvalidParameterException( - "User has not verified software token mfa" - ) - - if software_token_mfa_settings.get("PreferredMfa"): - user.preferred_mfa_setting = "SOFTWARE_TOKEN_MFA" - elif sms_mfa_settings and sms_mfa_settings["Enabled"]: - user.sms_mfa_enabled = True - - if sms_mfa_settings.get("PreferredMfa"): - user.preferred_mfa_setting = "SMS_MFA" - return None + return self.admin_set_user_mfa_preference( + user_pool.id, + username, + software_token_mfa_settings, + sms_mfa_settings, + ) raise NotAuthorizedError(access_token) @@ -2013,21 +2000,33 @@ class CognitoIdpBackend(BaseBackend): ) -> None: user = self.admin_get_user(user_pool_id, username) - if software_token_mfa_settings and software_token_mfa_settings.get("Enabled"): - if user.token_verified: - user.software_token_mfa_enabled = True + if software_token_mfa_settings: + if software_token_mfa_settings.get("Enabled"): + if user.token_verified: + user.software_token_mfa_enabled = True + else: + raise InvalidParameterException( + "User has not verified software token mfa" + ) else: - raise InvalidParameterException( - "User has not verified software token mfa" - ) + user.software_token_mfa_enabled = False if software_token_mfa_settings.get("PreferredMfa"): user.preferred_mfa_setting = "SOFTWARE_TOKEN_MFA" - elif sms_mfa_settings and sms_mfa_settings.get("Enabled"): - user.sms_mfa_enabled = True + elif user.preferred_mfa_setting != "SMS_MFA": + user.preferred_mfa_setting = "" + + if sms_mfa_settings: + if sms_mfa_settings.get("Enabled"): + user.sms_mfa_enabled = True + else: + user.sms_mfa_enabled = False if sms_mfa_settings.get("PreferredMfa"): user.preferred_mfa_setting = "SMS_MFA" + elif user.preferred_mfa_setting != "SOFTWARE_TOKEN_MFA": + user.preferred_mfa_setting = "" + return None def _validate_password(self, user_pool_id: str, password: str) -> None: diff --git a/tests/test_cognitoidp/test_cognitoidp.py b/tests/test_cognitoidp/test_cognitoidp.py index fa22ea810..b79464285 100644 --- a/tests/test_cognitoidp/test_cognitoidp.py +++ b/tests/test_cognitoidp/test_cognitoidp.py @@ -4235,6 +4235,8 @@ def test_setting_mfa(): for auth_flow in ["ADMIN_NO_SRP_AUTH", "ADMIN_USER_PASSWORD_AUTH"]: result = authentication_flow(conn, auth_flow) + + # Set MFA method conn.associate_software_token(AccessToken=result["access_token"]) conn.verify_software_token( AccessToken=result["access_token"], UserCode="123456" @@ -4243,12 +4245,24 @@ def test_setting_mfa(): AccessToken=result["access_token"], SoftwareTokenMfaSettings={"Enabled": True, "PreferredMfa": True}, ) - result = conn.admin_get_user( + user = conn.admin_get_user( UserPoolId=result["user_pool_id"], Username=result["username"] ) - assert len(result["UserMFASettingList"]) == 1 - assert result["PreferredMfaSetting"] == "SOFTWARE_TOKEN_MFA" + assert len(user["UserMFASettingList"]) == 1 + assert user["PreferredMfaSetting"] == "SOFTWARE_TOKEN_MFA" + + # Unset MFA method + conn.set_user_mfa_preference( + AccessToken=result["access_token"], + SoftwareTokenMfaSettings={"Enabled": False, "PreferredMfa": False}, + ) + user = conn.admin_get_user( + UserPoolId=result["user_pool_id"], Username=result["username"] + ) + + assert len(user["UserMFASettingList"]) == 0 + assert user["PreferredMfaSetting"] == "" @mock_cognitoidp @@ -4269,7 +4283,7 @@ def test_setting_mfa_when_token_not_verified(): @mock_cognitoidp -def test_admin_setting_mfa(): +def test_admin_setting_single_mfa(): conn = boto3.client("cognito-idp", "us-west-2") user_pool_id = conn.create_user_pool( @@ -4278,6 +4292,7 @@ def test_admin_setting_mfa(): username = "test@example.com" conn.admin_create_user(UserPoolId=user_pool_id, Username=username) + # Set MFA SMS method conn.admin_set_user_mfa_preference( Username=username, UserPoolId=user_pool_id, @@ -4287,6 +4302,50 @@ def test_admin_setting_mfa(): assert len(result["UserMFASettingList"]) == 1 assert result["PreferredMfaSetting"] == "SMS_MFA" + # Unset MFA SMS method + conn.admin_set_user_mfa_preference( + Username=username, + UserPoolId=user_pool_id, + SMSMfaSettings={"Enabled": False, "PreferredMfa": False}, + ) + result = conn.admin_get_user(UserPoolId=user_pool_id, Username=username) + assert len(result["UserMFASettingList"]) == 0 + assert result["PreferredMfaSetting"] == "" + + +@mock_cognitoidp +def test_admin_setting_mfa_totp_and_sms(): + 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"] + conn.associate_software_token(AccessToken=access_token) + conn.verify_software_token(AccessToken=access_token, UserCode="123456") + + # 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" + + # Unset MFA TOTP and SMS methods + conn.admin_set_user_mfa_preference( + Username=username, + UserPoolId=user_pool_id, + SoftwareTokenMfaSettings={"Enabled": False, "PreferredMfa": False}, + SMSMfaSettings={"Enabled": False, "PreferredMfa": False}, + ) + result = conn.admin_get_user(UserPoolId=user_pool_id, Username=username) + assert len(result["UserMFASettingList"]) == 0 + assert result["PreferredMfaSetting"] == "" + @mock_cognitoidp def test_admin_setting_mfa_when_token_not_verified():