enhanced error logging and REST error handing

aws_fusion/__init__.py

@@ -1 +1 @@
-__version__ = '1.3.1'
+__version__ = '1.4'

aws_fusion/app.py

@@ -4,6 +4,7 @@
 from importlib.metadata import version
 from .commands import open_browser, iam_user_credentials, generate_okta_device_auth_credentials, init
 
+
 def main():
     global_parser = argparse.ArgumentParser(add_help=False)
     global_parser.add_argument('-v', '--version', action='version', help="Display the version of this tool", version=version("aws_fusion"))
@@ -18,10 +19,10 @@ def main():
     init.setup(subparsers, global_parser)
 
     args = main_parser.parse_args()
-    
+
     if args.debug:
         logging.basicConfig(level=logging.DEBUG)
-    
+
     # Call the associated function for the selected sub-command
     if hasattr(args, 'func'):
         args.func(args)

aws_fusion/commands/generate_okta_device_auth_credentials.py

@@ -24,8 +24,8 @@ def run(args):
 
     if not assume_role_with_cache.does_valid_token_cache_exists():
         LOG.debug('Credential cache not found, invoking SAML')
-        device_code = device_auth(args.org_domain, args.oidc_client_id)
-        access_token, id_token = verification_and_token(args.org_domain, args.oidc_client_id, device_code)
+        device_code, expires_in = device_auth(args.org_domain, args.oidc_client_id)
+        access_token, id_token = verification_and_token(args.org_domain, args.oidc_client_id, device_code, expires_in)
         session_token = session_and_token(args.org_domain, args.oidc_client_id, access_token, id_token, args.aws_acct_fed_app_id)
         saml_response, roles, session_duration = saml_assertion(args.org_domain, session_token)
         assume_role_with_cache.assume_role_with_saml(saml_response, roles, session_duration)

aws_fusion/okta/api.py

@@ -1,7 +1,8 @@
+import json
+
 import requests
 import webbrowser
 import time
-import sys
 import base64
 import logging
 
@@ -10,76 +11,114 @@
 LOG = logging.getLogger(__name__)
 
 
-def device_auth(org_domain, oidc_client_id):
-    LOG.debug('Started device auth')
-    url = "https://" + org_domain + "/oauth2/v1/device/authorize"
-    payload = 'client_id=' + oidc_client_id + \
-              '&scope=openid%20okta.apps.sso%20okta.apps.read'
-    headers = {'Content-Type': 'application/x-www-form-urlencoded'}
+class OktaApiException(Exception):
+    """Exception for Okta API call"""
+    pass
 
-    request = requests.post(url, headers=headers, data=payload)
-    response = request.json()
 
-    verification_url = response['verification_uri_complete']
+def device_auth(org_domain, oidc_client_id):
+    LOG.debug('Started device auth')
+    url = f'https://{org_domain}/oauth2/v1/device/authorize'
+    payload = {
+        'client_id': oidc_client_id,
+        'scope': 'openid okta.apps.sso okta.apps.read'
+    }
+    headers = {
+        'Content-Type': 'application/x-www-form-urlencoded'
+    }
+
+    response = requests.post(url, headers=headers, data=payload)
+    response_body = response.json()
+    if response.status_code >= 300:
+        LOG.error(f'Got {response.status_code} error in getting device code: {json.dumps(response_body)}')
+        raise OktaApiException()
+
+    LOG.debug(f'Device code response: {json.dumps(response_body)}')
+
+    verification_url = response_body['verification_uri_complete']
     webbrowser.open_new_tab(verification_url)
 
-    LOG.debug(f'Got device code {response["device_code"]}')
-    return response['device_code']
+    return response_body['device_code'], response_body['expires_in']
 
 
-def verification_and_token(org_domain, oidc_client_id, device_code):
+def verification_and_token(org_domain, oidc_client_id, device_code, expires_in):
     LOG.debug('Started verification of device code')
-    url = "https://" + org_domain + "/oauth2/v1/token"
-    payload = 'client_id=' + oidc_client_id + '&device_code=' + device_code + \
-              '&grant_type=urn%3Aietf%3Aparams%3Aoauth%3Agrant-type%3Adevice_code'
-    headers = {'Content-Type': 'application/x-www-form-urlencoded'}
-
+    url = f'https://{org_domain}/oauth2/v1/token'
+    payload = {
+        'client_id': oidc_client_id,
+        'device_code': device_code,
+        'grant_type': 'urn:ietf:params:oauth:grant-type:device_code'
+    }
+    headers = {
+        'Content-Type': 'application/x-www-form-urlencoded'
+    }
+
+    time_passed = 0
+    waiting_time_each_iteration = 5
     while True:
-        request = requests.post(url, headers=headers, data=payload)
-        response = request.json()
+        response = requests.post(url, headers=headers, data=payload)
+        response_body = response.json()
 
         # Check for authorization pending
-        if request.status_code == 400 and response['error'] == 'authorization_pending':
+        if response.status_code == 400 and response_body['error'] == 'authorization_pending':
             LOG.debug('Waiting for verification')
-            time.sleep(5)
+            time.sleep(waiting_time_each_iteration)
+            time_passed += waiting_time_each_iteration
+            if time_passed >= expires_in:
+                LOG.error(f'Maximum waiting ({expires_in}s) for verification has exhausted')
+                raise OktaApiException()
             continue
 
         # Check for successful verification
-        if request.status_code == 200:
+        if response.status_code == 200:
             break
 
         # Unexpected state. Die.
-        LOG.error(response)
-        sys.exit(1)
+        LOG.error(f'Got {response.status_code} error during verification of device code: {json.dumps(response_body)}')
+        raise OktaApiException()
 
-    LOG.debug('Validated device code and got access_token & id_token')
-    return response['access_token'], response['id_token']
+    LOG.debug(f'Token response: {json.dumps(response_body)}')
+    return response_body['access_token'], response_body['id_token']
 
 
 def session_and_token(org_domain, oidc_client_id, access_token, id_token, aws_acct_fed_app_id):
     LOG.debug('Started getting of session token')
-    url = "https://" + org_domain + "/oauth2/v1/token"
-    payload = 'client_id=' + oidc_client_id + '&actor_token=' + access_token + \
-              '&actor_token_type=urn%3Aietf%3Aparams%3Aoauth%3Atoken-type%3Aaccess_token' + \
-              '&subject_token=' + id_token + \
-              '&subject_token_type=urn%3Aietf%3Aparams%3Aoauth%3Atoken-type%3Aid_token' + \
-              '&grant_type=urn%3Aietf%3Aparams%3Aoauth%3Agrant-type%3Atoken-exchange' + \
-              '&requested_token_type=urn%3Aokta%3Aoauth%3Atoken-type%3Aweb_sso_token' + \
-              '&audience=urn%3Aokta%3Aapps%3A' + aws_acct_fed_app_id
-    headers = {'Content-Type': 'application/x-www-form-urlencoded'}
-
-    request = requests.post(url, headers=headers, data=payload)
-    response = request.json()
-
-    LOG.debug('Got session token')
-    return response['access_token']
+    url = f'https://{org_domain}/oauth2/v1/token'
+    payload = {
+        'client_id': oidc_client_id,
+        'actor_token': access_token,
+        'actor_token_type': 'urn:ietf:params:oauth:token-type:access_token',
+        'subject_token': id_token,
+        'subject_token_type': 'urn:ietf:params:oauth:token-type:id_token',
+        'grant_type': 'urn:ietf:params:oauth:grant-type:token-exchange',
+        'requested_token_type': 'urn:okta:oauth:token-type:web_sso_token',
+        'audience': f'urn:okta:apps:{aws_acct_fed_app_id}'
+    }
+    headers = {
+        'Content-Type': 'application/x-www-form-urlencoded'
+    }
+
+    response = requests.post(url, headers=headers, data=payload)
+    response_body = response.json()
+    if response.status_code >= 300:
+        LOG.error(f'Got {response.status_code} error in session token: {json.dumps(response_body)}')
+        raise OktaApiException()
+
+    LOG.debug(f'Session token response: {json.dumps(response_body)}')
+    return response_body['access_token']
 
 
 def saml_assertion(org_domain, session_token):
     LOG.debug('Started SAML assertion')
     # Get SAML assertion
-    url = 'https://' + org_domain + '/login/token/sso?token=' + session_token
-    response = requests.get(url)
+    url = f'https://{org_domain}/login/token/sso'
+    query_params = {
+        'token': session_token
+    }
+    response = requests.get(url, params=query_params)
+    if response.status_code >= 300:
+        LOG.error(f'Got {response.status_code} error while getting SAML response')
+        raise OktaApiException()
 
     # Extract response value from SAML assertion call
     parser = BeautifulSoup(response.text, "html.parser")
@@ -93,8 +132,7 @@ def saml_assertion(org_domain, session_token):
         idp_and_role = role.text.split(',')
         roles[idp_and_role[1]] = idp_and_role[0]
 
-    session_duration_xml = parser.find("saml2:Attribute",
-                                       {"Name": "https://aws.amazon.com/SAML/Attributes/SessionDuration"})
+    session_duration_xml = parser.find("saml2:Attribute", {"Name": "https://aws.amazon.com/SAML/Attributes/SessionDuration"})
     session_duration = session_duration_xml.find("saml2:AttributeValue").text
 
     LOG.debug(f'Got valid SAML response: {saml_response}')