diff --git a/aws_google_auth/__init__.py b/aws_google_auth/__init__.py index c332bb3..d6d3df4 100644 --- a/aws_google_auth/__init__.py +++ b/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() diff --git a/aws_google_auth/google.py b/aws_google_auth/google.py index 2583ea5..16509f0 100644 --- a/aws_google_auth/google.py +++ b/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. diff --git a/aws_google_auth/tests/test_args_parser.py b/aws_google_auth/tests/test_args_parser.py index 292c61e..b04e286 100644 --- a/aws_google_auth/tests/test_args_parser.py +++ b/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): diff --git a/aws_google_auth/tests/test_init.py b/aws_google_auth/tests/test_init.py index 9c4901f..bd6387e 100644 --- a/aws_google_auth/tests/test_init.py +++ b/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),