Pass session to SMS auth challenge

IrinaMityayeva · Feb 22, 2018 · a465d05 · a465d05
1 parent 7293aa1
commit a465d05
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 7 deletions.
diff --git a/setup.py b/setup.py
@@ -6,7 +6,7 @@
 install_reqs = parse_requirements('requirements.txt', session=False)
 test_reqs = parse_requirements('requirements_test.txt', session=False)
 
-version = '0.6.2'
+version = '0.6.3'
 
 README="""Python class to integrate Boto3's Cognito client so it is easy to login users. With SRP support."""
 

diff --git a/warrant/__init__.py b/warrant/__init__.py
@@ -366,18 +366,22 @@ def admin_authenticate(self, password):
             AuthParameters=auth_params,
         )
 
-        msa_enabled = False
         if response['ChallengeName'] == self.SMS_MFA_CHALLENGE:
-            msa_enabled = True
-            return msa_enabled
+            mfa_resp = {
+                'mfa_enabled': True,
+                'session': response['Session']
+            }
         else:
             self.verify_token(response['AuthenticationResult']['IdToken'], 'id_token', 'id')
             self.refresh_token = response['AuthenticationResult']['RefreshToken']
             self.verify_token(response['AuthenticationResult']['AccessToken'], 'access_token', 'access')
             self.token_type = response['AuthenticationResult']['TokenType']
-            return msa_enabled
+            mfa_resp = {
+                'mfa_enabled': False,
+            }
+        return mfa_resp
 
-    def sms_mfa(self, auth_code, username=None):
+    def sms_mfa(self, auth_code, session, username=None):
         if not username:
             username = self.username
         challenge_response_params = {
@@ -387,7 +391,8 @@ def sms_mfa(self, auth_code, username=None):
         response = self.client.respond_to_auth_challenge(
             ClientId=self.client_id,
             ChallengeName='SMS_MFA',
-            ChallengeResponses=challenge_response_params
+            ChallengeResponses=challenge_response_params,
+            Session=session
         )
         self.verify_token(response['AuthenticationResult']['IdToken'], 'id_token', 'id')
         self.refresh_token = response['AuthenticationResult']['RefreshToken']