Merge pull request #55 from ewhitfield-godaddy/egw/retry-okta-session

Retry okta session
godaddy · Oct 25, 2022 · c9d680e · c9d680e
2 parents a792d17 + 262ba5f
commit c9d680e
Show file tree

Hide file tree

Showing 5 changed files with 59 additions and 10 deletions.
diff --git a/aws_okta_processor/core/fetcher.py b/aws_okta_processor/core/fetcher.py
@@ -1,5 +1,6 @@
 import boto3
 import json
+import sys
 
 import aws_okta_processor.core.saml as saml
 import aws_okta_processor.core.prompt as prompt
@@ -46,14 +47,17 @@ def fetch_credentials(self):
 
     def _get_app_roles(self):
         user = self._configuration["AWS_OKTA_USER"]
+        user_pass = self._authenticate.get_pass()
         organization = self._configuration["AWS_OKTA_ORGANIZATION"]
+        no_okta_cache = self._configuration["AWS_OKTA_NO_OKTA_CACHE"]
+
         okta = Okta(
             user_name=user,
-            user_pass=self._authenticate.get_pass(),
+            user_pass=user_pass,
             organization=organization,
             factor=self._configuration["AWS_OKTA_FACTOR"],
             silent=self._configuration["AWS_OKTA_SILENT"],
-            no_okta_cache=self._configuration["AWS_OKTA_NO_OKTA_CACHE"]
+            no_okta_cache=no_okta_cache
         )
 
         self._configuration["AWS_OKTA_USER"] = ''
@@ -73,11 +77,32 @@ def _get_app_roles(self):
         saml_response = okta.get_saml_response(
             application_url=application_url
         )
-
         saml_assertion = saml.get_saml_assertion(
             saml_response=saml_response
         )
 
+        if not saml_assertion and not no_okta_cache:
+            # Try again, but without using the cached Okta session
+            print_tty("Creating new Okta session.")
+            okta = Okta(
+                user_name=user,
+                user_pass=user_pass,
+                organization=organization,
+                factor=self._configuration["AWS_OKTA_FACTOR"],
+                silent=self._configuration["AWS_OKTA_SILENT"],
+                no_okta_cache=True
+            )
+            saml_response = okta.get_saml_response(
+                application_url=application_url
+            )
+            saml_assertion = saml.get_saml_assertion(
+                saml_response=saml_response
+            )
+
+        if not saml_assertion:
+            print_tty("ERROR: SAMLResponse tag was not found!")
+            sys.exit(1)
+
         aws_roles = saml.get_aws_roles(
             saml_assertion=saml_assertion,
             accounts_filter=self._configuration.get(

diff --git a/aws_okta_processor/core/okta.py b/aws_okta_processor/core/okta.py
@@ -60,9 +60,11 @@ def __init__(
         okta_session = None
 
         if not no_okta_cache:
-            okta_session = self.get_okta_session()
+            # Get session from cache
+            okta_session = self.get_okta_session_from_cache_file()
 
         if okta_session:
+            # Refresh the session ID of the cached session
             self.read_aop_from_okta_session(okta_session)
 
             self.refresh_okta_session_id(
@@ -94,7 +96,8 @@ def __init__(
                 user_pass=user_pass
             )
 
-            self.get_okta_session_id()
+            # This call sets self.okta_session_id
+            self.create_and_store_okta_session()
 
     def read_aop_from_okta_session(self, okta_session):
         if "aws-okta-processor" in okta_session:
@@ -105,6 +108,9 @@ def read_aop_from_okta_session(self, okta_session):
             del okta_session["aws-okta-processor"]
 
     def get_cache_file_path(self):
+        """ Returns the file path for the session cache file:
+        ~/.aws-okta-processor/cache/<username>-<organization>-session.json
+        """
         home_directory = os.path.expanduser('~')
         cache_directory = os.path.join(
             home_directory,
@@ -125,6 +131,7 @@ def get_cache_file_path(self):
         return cache_file_path
 
     def set_okta_session(self, okta_session=None):
+        """ Saves the given Okta session in our cache file. """
         session_data = dict(okta_session, **{
             "aws-okta-processor": {
                 "user_name": self.user_name,
@@ -136,7 +143,7 @@ def set_okta_session(self, okta_session=None):
 
         os.chmod(self.cache_file_path, 0o600)
 
-    def get_okta_session(self):
+    def get_okta_session_from_cache_file(self):
         session = {}
 
         if os.path.isfile(self.cache_file_path):
@@ -229,7 +236,10 @@ def verify_factor(self, factor=None, state_token=None):
 
         send_error(response=response)
 
-    def get_okta_session_id(self):
+    def create_and_store_okta_session(self):
+        """ Creates a new Okta session and caches it in our cache file for future use.
+        https://developer.okta.com/docs/reference/api/sessions/#get-started
+        """
         headers = {
             "Accept": "application/json",
             "Content-Type": "application/json"

diff --git a/aws_okta_processor/core/saml.py b/aws_okta_processor/core/saml.py
@@ -18,12 +18,25 @@
 
 
 def get_saml_assertion(saml_response=None):
+    """ Extracts the SAML assertion from the saml_response HTML by finding the appropriate HTML element. """
     soup = BeautifulSoup(saml_response, "html.parser")
 
     for input_tag in soup.find_all('input'):
         if input_tag.get('name') == 'SAMLResponse':
             return input_tag.get('value')
 
+    if soup.find('div', {"id": "okta-sign-in"}):
+        # Supplied Okta session not sufficient to get SAML assertion.
+        # This condition may be missed if Okta significantly changes the app-level MFA page
+        print_tty("SAMLResponse tag not found due to MFA challenge.")
+        return None
+
+    if soup.find('div', {"id": "password-verification-challenge"}):
+        # Supplied Okta session not sufficient to get SAML assertion.
+        # This condition may be missed if Okta significantly changes the app-level re-auth page
+        print_tty("SAMLResponse tag not found due to password verification challenge.")
+        return None
+
     print_tty("ERROR: SAMLResponse tag was not found!")
     sys.exit(1)
 

diff --git a/setup.py b/setup.py
@@ -27,7 +27,8 @@
     'boto3>=1.9.134',
     'bs4>=0.0.1',
     'contextlib2>=0.5.5',
-    'six>=1.12.0'
+    'six>=1.12.0',
+    'defusedxml>=0.7.1'
 ]
 
 TEST_REQUIREMENTS = [

diff --git a/tests/core/test_okta.py b/tests/core/test_okta.py
@@ -297,7 +297,7 @@ def test_okta_mfa_totp_challenge(
         self.assertEqual(okta.okta_session_id, "session_token")
 
     @patch('aws_okta_processor.core.okta.Okta.get_okta_single_use_token')
-    @patch('aws_okta_processor.core.okta.Okta.get_okta_session_id')
+    @patch('aws_okta_processor.core.okta.Okta.create_and_store_okta_session')
     @patch('aws_okta_processor.core.okta.input')
     def test_read_aop_from_okta_session_should_read_aop_options(
         self,
@@ -325,7 +325,7 @@ def test_read_aop_from_okta_session_should_read_aop_options(
     @patch('aws_okta_processor.core.okta.os.chmod')
     @patch('aws_okta_processor.core.okta.Okta.get_cache_file_path', return_value='/tmp/test.json')
     @patch('aws_okta_processor.core.okta.Okta.get_okta_single_use_token')
-    @patch('aws_okta_processor.core.okta.Okta.get_okta_session_id')
+    @patch('aws_okta_processor.core.okta.Okta.create_and_store_okta_session')
     @patch('aws_okta_processor.core.okta.input')
     @patch('builtins.open', new_callable=mock_open)
     def test_set_okta_session_should_write_session_data(