Yasuke/moto
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

CognitoIDP: clean MFA settings capability and enable multiple MFA methods (#6853)

Browse Source
This commit is contained in:
João Pedro Schmitt 2023-09-26 10:45:40 -07:00 committed by GitHub
parent 178f2b8c03
commit 9a84423187
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 90 additions and 32 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

55
moto/cognitoidp/models.py
View File

@ -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:

67
tests/test_cognitoidp/test_cognitoidp.py
View File

@ -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():