Merge pull request #94 from cevoaustralia/enable-html-debug

Enable html debug

aws_google_auth/__init__.py

@@ -31,6 +31,7 @@ def parse_args(args):
     parser.add_argument('-D', '--disable-u2f', action='store_true', help='Disable U2F functionality.')
     parser.add_argument('--no-cache', dest="saml_cache", action='store_false', help='Do not cache the SAML Assertion.')
     parser.add_argument('--resolve-aliases', action='store_true', help='Resolve AWS account aliases.')
+    parser.add_argument('--save-failure-html', action='store_true', help='Write HTML failure responses to file for troubleshooting.')
 
     role_group = parser.add_mutually_exclusive_group()
     role_group.add_argument('-a', '--ask-role', action='store_true', help='Set true to always pick the role')
@@ -168,9 +169,9 @@ def process_auth(args, config):
 
         # There is no way (intentional) to pass in the password via the command
         # line nor environment variables. This prevents password leakage.
+        keyring_password = None
         if config.keyring:
-            keyring_password = keyring.get_password(
-                "aws-google-auth", config.username)
+            keyring_password = keyring.get_password("aws-google-auth", config.username)
             if keyring_password:
                 config.password = keyring_password
             else:
@@ -181,7 +182,7 @@ def process_auth(args, config):
         # Validate Options
         config.raise_if_invalid()
 
-        google_client = google.Google(config)
+        google_client = google.Google(config, args.save_failure_html)
         google_client.do_login()
         saml_xml = google_client.parse_saml()
 

aws_google_auth/google.py

@@ -1,6 +1,7 @@
 #!/usr/bin/env python
 # -*- coding: utf8 -*-
 from __future__ import print_function
+from requests import HTTPError
 from . import _version
 
 import sys
@@ -27,7 +28,7 @@ def __init__(self, *args):
 
 
 class Google:
-    def __init__(self, config):
+    def __init__(self, config, save_failure):
         """The Google object holds authentication state
         for a given session. You need to supply:
 
@@ -43,14 +44,14 @@ def __init__(self, config):
         self.version = _version.__version__
         self.config = config
         self.base_url = 'https://accounts.google.com'
+        self.save_failure = save_failure
 
     @property
     def login_url(self):
         return self.base_url + "/o/saml2/initsso?idpid={}&spid={}&forceauthn=false".format(
             self.config.idp_id, self.config.sp_id)
 
-    @staticmethod
-    def check_for_failure(sess):
+    def check_for_failure(self, sess):
 
         if isinstance(sess.reason, bytes):
             # We attempt to decode utf-8 first because some servers
@@ -68,14 +69,22 @@ def check_for_failure(sess):
             raise ExpectedGoogleException(u'{} accessing {}'.format(
                 reason, sess.url))
 
-        sess.raise_for_status()
+        try:
+            sess.raise_for_status()
+        except HTTPError as ex:
+
+            if self.save_failure:
+                print("Saving failure trace in 'failure.html'")
+                with open("failure.html", 'w') as out:
+                    out.write(sess.text)
+
+            raise ex
 
         return sess
 
     def post(self, url, data=None, json=None):
         try:
-            response = self.check_for_failure(
-                self.session.post(url, data=data, json=json))
+            response = self.check_for_failure(self.session.post(url, data=data, json=json))
         except requests.exceptions.ConnectionError as e:
             print(
                 'There was a connection error, check your network settings: {}'.
@@ -121,9 +130,7 @@ def parse_error_message(sess):
 
     def do_login(self):
         self.session = requests.Session()
-        self.session.headers[
-            'User-Agent'] = "AWS Sign-in/{} (Cevo aws-google-auth)".format(
-                self.version)
+        self.session.headers['User-Agent'] = "AWS Sign-in/{} (Cevo aws-google-auth)".format(self.version)
         sess = self.get(self.login_url)
 
         # Collect information from the page source
@@ -132,9 +139,7 @@ def do_login(self):
         self.cont = first_page.find('input', {'name': 'continue'}).get('value')
         page = first_page.find('input', {'name': 'Page'}).get('value')
         sign_in = first_page.find('input', {'name': 'signIn'}).get('value')
-        account_login_url = first_page.find('form', {
-            'id': 'gaia_loginform'
-        }).get('action')
+        account_login_url = first_page.find('form', {'id': 'gaia_loginform'}).get('action')
 
         payload = {
             'bgresponse': 'js_disabled',
@@ -261,26 +266,26 @@ def do_login(self):
 
     @staticmethod
     def check_extra_step(response):
-        extra_step = response.find(
-            text='This extra step shows that it’s really you trying to sign in'
-        )
+        extra_step = response.find(text='This extra step shows that it’s really you trying to sign in')
         if extra_step:
             if response.find(id='contactAdminMessage'):
                 raise ValueError(response.find(id='contactAdminMessage').text)
 
     def parse_saml(self):
         if self.session_state is None:
-            raise RuntimeError(
-                'You must use do_login() before calling parse_saml()')
+            raise RuntimeError('You must use do_login() before calling parse_saml()')
 
         parsed = BeautifulSoup(self.session_state.text, 'html.parser')
         try:
-            saml_element = parsed.find('input', {
-                'name': 'SAMLResponse'
-            }).get('value')
+            saml_element = parsed.find('input', {'name': 'SAMLResponse'}).get('value')
         except:
-            raise RuntimeError(
-                'Could not find SAML response, check your credentials')
+
+            if self.save_failure:
+                print("SAML lookup failed, storing failure page to 'saml.html' to assist with debugging.")
+                with open("saml.html", 'w') as out:
+                    out.write(self.session_state.text.encode('utf-8'))
+
+            raise ExpectedGoogleException('Something went wrong - Could not find SAML response, check your credentials or use --save-failure-html to debug.')
 
         return base64.b64decode(saml_element)
 
@@ -491,13 +496,22 @@ def handle_prompt(self, sess):
 
         self.check_prompt_code(response_page)
 
-        print(
-            "Open the Google App, and tap 'Yes' on the prompt to sign in ...")
+        print("Open the Google App, and tap 'Yes' on the prompt to sign in ...")
 
         self.session.headers['Referer'] = sess.url
 
-        parsed_response = json.loads(
-            self.post(await_url, json=await_body).text)
+        retry = True
+        response = None
+        while retry:
+            try:
+                response = self.post(await_url, json=await_body)
+                retry = False
+            except requests.exceptions.HTTPError as ex:
+
+                if not ex.response.status_code == 500:
+                    raise ex
+
+        parsed_response = json.loads(response.text)
 
         payload = {
             'challengeId':
@@ -750,8 +764,7 @@ def handle_selectchallenge(self, sess):
         selected_challenge = input("Enter MFA choice number ({}): ".format(
             challenge_ids[-1:][0])) or None
 
-        if selected_challenge is not None and int(
-                selected_challenge) in challenge_ids:
+        if selected_challenge is not None and int(selected_challenge) in challenge_ids:
             challenge_id = int(selected_challenge)
         else:
             # use the highest index as that will default to prompt, then sms, then totp, etc.

aws_google_auth/tests/test_args_parser.py

@@ -27,9 +27,11 @@ def test_no_arguments(self):
         self.assertEqual(parser.role_arn, None)
         self.assertEqual(parser.username, None)
 
+        self.assertFalse(parser.save_failure_html)
+
         # Assert the size of the parameter so that new parameters trigger a review of this function
         # and the appropriate defaults are added here to track backwards compatibility in the future.
-        self.assertEqual(len(vars(parser)), 12)
+        self.assertEqual(len(vars(parser)), 13)
 
     def test_username(self):
 

aws_google_auth/tests/test_init.py

@@ -53,6 +53,7 @@ def test_main_method_chaining(self, process_auth, resolve_config, exit_if_unsupp
                                          region=None,
                                          resolve_aliases=False,
                                          role_arn=None,
+                                         save_failure_html=False,
                                          saml_cache=True,
                                          sp_id=None,
                                          username=None))
@@ -68,6 +69,7 @@ def test_main_method_chaining(self, process_auth, resolve_config, exit_if_unsupp
                                          region=None,
                                          resolve_aliases=False,
                                          role_arn=None,
+                                         save_failure_html=False,
                                          saml_cache=True,
                                          sp_id=None,
                                          username=None),