From d7409d2231a7f10e324b9bd320159ddd3d326bea Mon Sep 17 00:00:00 2001
From: pcmxgti <16561338+pcmxgti@users.noreply.github.com>
Date: Fri, 1 Dec 2023 12:07:31 -0500
Subject: [PATCH] Clean up AuthZ cookie generation

---
 tests/unit/test_okta.py |  9 ++++++-
 tokendito/okta.py       | 52 +++++++++++++++++++++++++++--------------
 2 files changed, 43 insertions(+), 18 deletions(-)

diff --git a/tests/unit/test_okta.py b/tests/unit/test_okta.py
index 8a32ebe..e9587dd 100644
--- a/tests/unit/test_okta.py
+++ b/tests/unit/test_okta.py
@@ -549,6 +549,9 @@ def test_create_authz_cookies():
         "org": "acme",
         "authorization_endpoint": "pytesturl",
         "token_endpoint": "tokeneurl",
+        "nonce": "pytest",
+        "issuer": "pytest",
+        "ln": "pytest",
     }
     assert okta.create_authz_cookies(pytest_oauth2_config, pytest_oauth2_session_data) is None
     from tokendito import okta
@@ -708,7 +711,9 @@ def test_get_oauth2_configuration(mocker):
         "grant_types_supported": "authorization_code",
         "request_parameter_supported": "pytest",
     }
-    pytest_config = Config(okta={"client_id": "test_client_id", "org": "acme"})
+    pytest_config = Config(
+        okta={"client_id": "test_client_id", "org": "acme", "username": "pytest"}
+    )
     mocker.patch.object(HTTP_client, "get", return_value=response)
     assert okta.get_oauth2_configuration(pytest_config)["org"] == "acme"
 
@@ -739,6 +744,8 @@ def test_validate_oauth2_configuration():
         "scopes_supported": "pytest",
         "response_types_supported": "code",
         "request_parameter_supported": "pytest",
+        "ln": "pytest",
+        "nonce": "pytest",
     }
     assert okta.validate_oauth2_configuration(pytest_oauth2_config) is None
 
diff --git a/tokendito/okta.py b/tokendito/okta.py
index 541f135..e40b7dd 100644
--- a/tokendito/okta.py
+++ b/tokendito/okta.py
@@ -17,7 +17,7 @@
 import re
 import sys
 import time
-from urllib.parse import urlencode
+from urllib.parse import quote
 from urllib.parse import urlparse
 
 import bs4
@@ -194,20 +194,33 @@ def create_authz_cookies(oauth2_config, oauth2_session_data):
 
     Needed for SAML2 flow for OIE.
     """
-    session_token = HTTP_client.session.cookies.get("sessionToken")
     try:
         oauth2_url = f"{oauth2_config['org']}/oauth2/v1"
         oauth2_config_reformatted = {
             "responseType": "code",
             "state": oauth2_session_data["state"],
+            "nonce": oauth2_config["nonce"],
+            "scopes": [
+                "openid",
+                "profile",
+                "email",
+                "okta.users.read.self",
+                "okta.users.manage.self",
+                "okta.internal.enduser.read",
+                "okta.internal.enduser.manage",
+                "okta.enduser.dashboard.read",
+                "okta.enduser.dashboard.manage",
+            ],
             "clientId": oauth2_config["client_id"],
-            "authorizeUrl": oauth2_config["authorization_endpoint"],
-            "tokenUrl": oauth2_config["token_endpoint"],
-            "scope": "openid",
-            "sessionToken": session_token,
-            "userInfoUrl": f"{oauth2_url}/userinfo",
-            "revokeUrl": f"{oauth2_url}/revoke",
-            "logoutUrl": f"{oauth2_url}/logout",
+            "urls": {
+                "issuer": oauth2_config["issuer"],
+                "authorizeUrl": oauth2_config["authorization_endpoint"],
+                "userInfoUrl": f"{oauth2_url}/userinfo",
+                "tokenUrl": oauth2_config["token_endpoint"],
+                "revokeUrl": f"{oauth2_url}/revoke",
+                "logoutUrl": f"{oauth2_url}/logout",
+            },
+            "ignoreSignature": False,
         }
     except KeyError as e:
         logger.error(f"Missing key in config:{e}")
@@ -217,11 +230,13 @@ def create_authz_cookies(oauth2_config, oauth2_session_data):
     domain = urlparse(oauth2_config["org"]).netloc
     cookiejar.set(
         "okta-oauth-redirect-params",
-        f"{{{urlencode(oauth2_config_reformatted)}}}",
+        quote(json.dumps(oauth2_config_reformatted, separators=(",", ":")), safe="{}:[]/"),
         domain=domain,
         path="/",
     )
     cookiejar.set("okta-oauth-state", oauth2_session_data["state"], domain=domain, path="/")
+    cookiejar.set("okta-oauth-nonce", oauth2_config["nonce"], domain=domain, path="/")
+    cookiejar.set("ln", oauth2_config["ln"], domain=domain, path="/")
     HTTP_client.add_cookies(cookiejar)  # add cookies
 
 
@@ -267,17 +282,16 @@ def send_saml_response(config, saml_response):
             "content-type": "application/json",
         }
         response = HTTP_client.get(
-            # myurl, allow_redirects=False, params={"stateToken": state_token}
             f"{config.okta['org']}/login/token/redirect",
             params=params,
             headers=headers,
         )
-        logger.warning(
-            f"""
-            State token from {url}: {state_token}. TODO: need to go from this state token
-            to an idx cookies.
-            """
-        )
+        if "idx" not in response.cookies:
+            logger.error(
+                f"Session cookie idx for {config.okta['org']} not found. Please file a bug."
+            )
+            logger.debug(f"Response: {response.text}")
+            sys.exit(2)
 
 
 def get_session_token(config, primary_auth, headers):
@@ -595,6 +609,8 @@ def get_oauth2_configuration(config):
     oauth2_config = response.json()
     oauth2_config["org"] = config.okta["org"]
     oauth2_config["client_id"] = get_client_id(config)
+    oauth2_config["ln"] = config.okta["username"]
+    oauth2_config["nonce"] = "TODO"
     validate_oauth2_configuration(oauth2_config)
     return oauth2_config
 
@@ -614,6 +630,8 @@ def validate_oauth2_configuration(oauth2_config):
         "scopes_supported",
         "client_id",
         "org",
+        "ln",
+        "nonce",
     }  # the authorization server must have these config elements
     for item in mandadory_oauth2_config_items:
         if item not in oauth2_config: