call /authorize as soon as we're logged in, still no idx cookies at s…

…aml2 call...

tokendito/okta.py

@@ -308,24 +308,24 @@ def get_session_token(config, primary_auth, headers):
     return session_token
 
 
-def get_access_token(authz_code_flow_data, authorize_code):
+def get_access_token(oauth2_config, oauth2_session_data, authorize_code):
     """Get OAuth token from Okta by calling /token endpoint.
 
     :param url: URL of the Okta OAuth token endpoint
     :return: OAuth token
     """
     payload = {
         "code": authorize_code,
-        "state": authz_code_flow_data["state"],
-        "grant_type": authz_code_flow_data["grant_type"],
-        "redirect_uri": authz_code_flow_data["redirect_uri"],
-        "client_id": authz_code_flow_data["client_id"],
-        "code_verifier": authz_code_flow_data["code_verifier"],
+        "state": oauth2_session_data["state"],
+        "grant_type": oauth2_session_data["grant_type"],
+        "redirect_uri": oauth2_session_data["redirect_uri"],
+        "client_id": oauth2_config["client_id"],
+        "code_verifier": oauth2_session_data["code_verifier"],
     }
     headers = {"accept": "application/json"}
     # Using the http_client to make the POST request
     response = HTTP_client.post(
-        authz_code_flow_data["token_endpoint_url"], data=payload, headers=headers, return_json=True
+        oauth2_config["token_endpoint"], data=payload, headers=headers, return_json=True
     )
     return response
 
@@ -344,14 +344,14 @@ def get_client_id(config):
     return config.okta["client_id"]
 
 
-def get_redirect_uri(config):
+def get_redirect_uri(oauth2_url):
     """
     Get the redirect uri needed by the Authorization Code Flow.
 
     Return url
     """
-    url = f"{config.okta['org']}/enduser/callback"
-    return url
+    uri = f"{oauth2_url}/enduser/callback"
+    return uri
 
 
 def get_response_type():
@@ -447,24 +447,40 @@ def get_authorize_code(response, payload):
         return authorize_code.group()
 
 
-def authorize_request(config, oauth2_config, oauth2_session_data):
+def authorization_code_enabled(oauth2_config):
+    """
+    Determine if authorization code grant is enabled.
+
+    Returns True if the dict key is in authorization server info, and False otherwise,
+    """
+    try:
+        if "authorization_code" not in oauth2_config["grant_types_supported"]:
+            return False
+    except (KeyError, ValueError) as e:
+        logger.error(f"No grant types supported on {oauth2_config['org']}:{str(e)}")
+        sys.exit(1)
+
+    return True
+
+
+def authorize_request(oauth2_config, oauth2_session_data):
     """
     Call /authorize endpoint.
     :param
     :return: authorization code, needed for /token call
     """
-    logger.debug(f"oauth_code_request({config}, {oauth2_config}, {oauth2_session_data})")
+    logger.debug(f"oauth_code_request({oauth2_config}, {oauth2_session_data})")
     headers = {"accept": "application/json", "content-type": "application/json"}
 
     payload = {
-        "client_id": config.okta["client_id"],
+        "client_id": oauth2_config["client_id"],
         "redirect_uri": oauth2_session_data["redirect_uri"],
         "response_type": oauth2_session_data["response_type"],
         "scope": oauth2_session_data["scope"],
         "state": oauth2_session_data["state"],
         "code_challenge": oauth2_session_data["code_challenge"],
         "code_challenge_method": oauth2_session_data["code_challenge_method"],
-        #        "prompt": "none",  # dont authenticate
+        "prompt": "none",  # dont authenticate
     }
 
     response = HTTP_client.get(
@@ -477,13 +493,14 @@ def authorize_request(config, oauth2_config, oauth2_session_data):
     return authorize_code
 
 
-def generate_oauth2_session_data(config):
+def generate_oauth2_session_data(url):
     "Generate some oauth2 session data, to have the same in oath2 cookies and /authorize call"
     authz_session_data = {
         "response_type": get_response_type(),
         "scope": get_authorize_scope(),
         "state": get_oauth2_state(),
-        "redirect_uri": get_redirect_uri(config),
+        "redirect_uri": get_redirect_uri(url),
+        "grant_type": "authorization_code",
     }
     if pkce_enabled():
         code_verifier = get_pkce_code_verifier()
@@ -494,34 +511,20 @@ def generate_oauth2_session_data(config):
     return authz_session_data
 
 
-def authorization_code_enabled(org_url, oauth2_config):
-    """
-    Determine if authorization code grant is enabled.
-
-    Returns True if the dict key is in authorization server info, and False otherwise,
-    """
-    try:
-        if "authorization_code" not in oauth2_config["grant_types_supported"]:
-            return False
-    except (KeyError, ValueError) as e:
-        logger.error(f"No grant types supported on {org_url}:{str(e)}")
-        sys.exit(1)
-
-    return True
-
-
-def get_oauth2_configuration(url=None):
+def get_oauth2_configuration(config):
     """Get authorization server configuration data from Okta instance.
 
     :param url: URL of the Okta org
     :return: dict of conguration values
     """
-    url = f"{url}/.well-known/oauth-authorization-server"
+    url = f"{config.okta['org']}/.well-known/oauth-authorization-server"
     headers = {"accept": "application/json"}
     response = HTTP_client.get(url, headers=headers)
     logger.debug(f"Authorization Server info: {response.json()}")
     # todo: handle errors.
     oauth2_config = response.json()
+    oauth2_config["client_id"] = config.okta["client_id"]
+    oauth2_config["org"] = config.okta["org"]
     validate_oauth2_configuration(oauth2_config)
     return oauth2_config
 
@@ -539,6 +542,8 @@ def validate_oauth2_configuration(oauth2_config):
         "grant_types_supported",
         "response_types_supported",
         "scopes_supported",
+        "client_id",
+        "org",
     }  # the authorization server must have these config elements
     for item in mandadory_oauth2_config_items:
         if item not in oauth2_config:
@@ -587,7 +592,7 @@ def create_authn_cookies(authn_org_url, session_token):
     HTTP_client.add_cookies(cookiejar)  # add cookies
 
 
-def idp_authenticate(config):
+def idp_authenticate(config, oauth2_config, oauth2_session_data):
     "Authenticate user to okta."
     auth_properties = get_auth_properties(userid=config.okta["username"], url=config.okta["org"])
 
@@ -604,12 +609,20 @@ def idp_authenticate(config):
     if is_saml2_authentication(auth_properties):
         # We may loop thru the saml2 servers until
         # we find the authentication server.
-        saml2_authenticate(config, auth_properties)
+        saml2_authenticate(config, auth_properties, oauth2_config, oauth2_session_data)
     elif local_authentication_enabled(auth_properties):
         session_token = local_authenticate(config)
         # authentication sends us a token
         # which we then put in our session cookies
         create_authn_cookies(config.okta["org"], session_token)
+        if oauth2_config is not None:  # meaning our initial okta org was OIE enabled.
+            if authorization_code_enabled(oauth2_config):
+                authorize_code = authorize_request(oauth2_config, oauth2_session_data)
+                get_access_token(oauth2_config, oauth2_session_data, authorize_code)
+            else:
+                logger.warning(
+                    f"Authorization Code is not enabled on {oauth2_config['org']}, skipping oauth2"
+                )
     else:
         logger.error(f"{auth_properties['type']} login via IdP Discovery is not curretly supported")
         sys.exit(1)
@@ -628,16 +641,15 @@ def access_control(config):
     """
     logger.debug(f"access_control({config})")
 
+    oauth2_config = None
+    oauth2_session_data = None
     if oie_enabled(config.okta["org"]):
         logger.debug("OIE enabled")
-        oauth2_config = get_oauth2_configuration(config.okta["org"])
-        oauth2_session_data = generate_oauth2_session_data(config)
+        oauth2_config = get_oauth2_configuration(config)
+        oauth2_session_data = generate_oauth2_session_data(config.okta["org"])
         create_oauth2_cookies(config, oauth2_config, oauth2_session_data)
-        authorize_request(config, oauth2_config, oauth2_session_data)
 
-    idp_authenticate(config)
-
-    return
+    idp_authenticate(config, oauth2_config, oauth2_session_data)
 
 
 def step_up_authenticate(config, state_token):
@@ -683,7 +695,7 @@ def is_saml2_auth(auth_properties):
     return False
 
 
-def saml2_authenticate(config, auth_properties):
+def saml2_authenticate(config, auth_properties, oauth2_config, oauth2_session_data):
     """SAML2 authentication flow.
 
     :param config: Config object
@@ -701,7 +713,7 @@ def saml2_authenticate(config, auth_properties):
 
     # Try to authenticate using the new configuration. This could cause
     # recursive calls, which allows for IdP chaining.
-    idp_authenticate(saml2_config)
+    idp_authenticate(saml2_config, oauth2_config, oauth2_session_data)
 
     # Once we are authenticated, send the SAML request to the IdP.
     # This call requires session cookies.