CognitoIDP: Add SMS_MFA support to admin_initiate_auth and respond_to_auth_challenge methods (#7379)

This commit is contained in:
Ari Freyr Asgeirsson 2024-02-22 20:44:52 +00:00 committed by GitHub
parent 7c26f9f831
commit 25bc983df0
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 64 additions and 4 deletions

View File

@ -1466,6 +1466,16 @@ class CognitoIdpBackend(BaseBackend):
"Session": session, "Session": session,
} }
if user.sms_mfa_enabled and user.preferred_mfa_setting == "SMS_MFA":
session = str(random.uuid4())
self.sessions[session] = user_pool
return {
"ChallengeName": "SMS_MFA",
"ChallengeParameters": {},
"Session": session,
}
return self._log_user_in(user_pool, client, username) return self._log_user_in(user_pool, client, username)
elif auth_flow in (AuthFlow.REFRESH_TOKEN, AuthFlow.REFRESH_TOKEN_AUTH): elif auth_flow in (AuthFlow.REFRESH_TOKEN, AuthFlow.REFRESH_TOKEN_AUTH):
refresh_token: str = auth_parameters.get("REFRESH_TOKEN") # type: ignore[assignment] refresh_token: str = auth_parameters.get("REFRESH_TOKEN") # type: ignore[assignment]
@ -1578,13 +1588,13 @@ class CognitoIdpBackend(BaseBackend):
del self.sessions[session] del self.sessions[session]
return self._log_user_in(user_pool, client, username) return self._log_user_in(user_pool, client, username)
elif challenge_name == "SOFTWARE_TOKEN_MFA": elif challenge_name == "SOFTWARE_TOKEN_MFA" or challenge_name == "SMS_MFA":
username: str = challenge_responses.get("USERNAME") # type: ignore[no-redef] username: str = challenge_responses.get("USERNAME") # type: ignore[no-redef]
self.admin_get_user(user_pool.id, username) self.admin_get_user(user_pool.id, username)
software_token_mfa_code = challenge_responses.get("SOFTWARE_TOKEN_MFA_CODE") mfa_code = challenge_responses.get(f"{challenge_name}_CODE")
if not software_token_mfa_code: if not mfa_code:
raise ResourceNotFoundError(software_token_mfa_code) raise ResourceNotFoundError(mfa_code)
if client.generate_secret: if client.generate_secret:
secret_hash = challenge_responses.get("SECRET_HASH") secret_hash = challenge_responses.get("SECRET_HASH")

View File

@ -4454,6 +4454,56 @@ def test_admin_initiate_auth_when_token_totp_enabled():
assert result["AuthenticationResult"]["TokenType"] == "Bearer" assert result["AuthenticationResult"]["TokenType"] == "Bearer"
@mock_aws
def test_admin_initiate_auth_when_sms_mfa_enabled():
conn = boto3.client("cognito-idp", "us-west-2")
result = authentication_flow(conn, "ADMIN_NO_SRP_AUTH")
user_pool_id = result["user_pool_id"]
username = result["username"]
client_id = result["client_id"]
password = result["password"]
# Set MFA SMS methods
conn.admin_set_user_mfa_preference(
Username=username,
UserPoolId=user_pool_id,
SMSMfaSettings={"Enabled": True, "PreferredMfa": True},
)
result = conn.admin_get_user(UserPoolId=user_pool_id, Username=username)
assert len(result["UserMFASettingList"]) == 1
assert result["PreferredMfaSetting"] == "SMS_MFA"
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"] == "SMS_MFA"
assert result["Session"] != ""
result = conn.admin_respond_to_auth_challenge(
UserPoolId=user_pool_id,
ClientId=client_id,
ChallengeName="SMS_MFA",
Session=result["Session"],
ChallengeResponses={
"SMS_MFA_CODE": "123456",
"USERNAME": username,
},
)
assert result["AuthenticationResult"]["IdToken"] != ""
assert result["AuthenticationResult"]["AccessToken"] != ""
assert result["AuthenticationResult"]["RefreshToken"] != ""
assert result["AuthenticationResult"]["TokenType"] == "Bearer"
@mock_aws @mock_aws
def test_admin_setting_mfa_when_token_not_verified(): def test_admin_setting_mfa_when_token_not_verified():
conn = boto3.client("cognito-idp", "us-west-2") conn = boto3.client("cognito-idp", "us-west-2")