Updates the client to work using Refresh Tokens instead of clientId and Secret

pyproject.toml

@@ -24,6 +24,7 @@ dependencies = [
     "certifi",
     "confluent-kafka >= 2.2.0",
     "requests",
+    "pyjwt"
 ]
 requires-python = ">=3.9"
 dynamic = [ "version" ]

src/gcn_kafka/core.py

@@ -10,7 +10,7 @@
 from .oidc import set_oauth_cb
 
 
-def get_config(mode, config, **kwargs):
+def get_config(mode, scope, config, **kwargs):
     # Merge configuration from user.
     config = update_config(config, **kwargs)
 
@@ -25,20 +25,15 @@
 
     domain = config.pop("domain", "gcn.nasa.gov")
     client_id = config.pop("client_id", None)
-    client_secret = config.pop("client_secret", None)
-
     config.setdefault("bootstrap.servers", f"kafka.{domain}")
 
     if client_id:
-        # Configure authentication and authorization using OpenID Connect.
         config.setdefault("sasl.mechanisms", "OAUTHBEARER")
-        config.setdefault("sasl.oauthbearer.method", "oidc")
         config.setdefault("sasl.oauthbearer.client.id", client_id)
-        if client_secret:
-            config.setdefault("sasl.oauthbearer.client.secret", client_secret)
         config.setdefault(
             "sasl.oauthbearer.token.endpoint.url", f"https://auth.{domain}/oauth2/token"
         )
+        config.setdefault("sasl.oauthbearer.scope", scope)
 
     if mode == "consumer" and not config.get("group.id"):
         config["group.id"] = str(uuid4())
@@ -59,9 +54,9 @@
 class Producer(confluent_kafka.Producer):
     def __init__(
         self,
+        scope: Optional[str] = None,  # Maybe these should be required?
         config: Optional[Mapping[str, Any]] = None,
         client_id: Optional[str] = None,
-        client_secret: Optional[str] = None,
         domain: Optional[
             Union[
                 Literal["gcn.nasa.gov"],
@@ -74,9 +69,9 @@
         super().__init__(
             get_config(
                 "producer",
+                scope,
                 config,
                 client_id=client_id,
-                client_secret=client_secret,
                 domain=domain,
                 **kwargs,
             )
@@ -89,9 +84,9 @@
 class Consumer(confluent_kafka.Consumer):
     def __init__(
         self,
+        scope: Optional[str] = None,  # Maybe these should be required?
         config: Optional[Mapping[str, Any]] = None,
         client_id: Optional[str] = None,
-        client_secret: Optional[str] = None,
         domain: Optional[
             Union[
                 Literal["gcn.nasa.gov"],
@@ -104,9 +99,9 @@
         super().__init__(
             get_config(
                 "consumer",
+                scope,
                 config,
                 client_id=client_id,
-                client_secret=client_secret,
                 domain=domain,
                 **kwargs,
             )
@@ -121,7 +116,6 @@
         self,
         config: Optional[Mapping[str, Any]] = None,
         client_id: Optional[str] = None,
-        client_secret: Optional[str] = None,
         domain: Optional[
             Union[
                 Literal["gcn.nasa.gov"],
@@ -135,8 +129,8 @@
             get_config(
                 "admin",
                 config,
+                scope="gcn.nasa.gov/kafka-admin",
                 client_id=client_id,
-                client_secret=client_secret,
                 domain=domain,
                 **kwargs,
             )

src/gcn_kafka/oidc.py

@@ -1,5 +1,7 @@
 # SPDX-License-Identifier: CC0-1.0
+from pathlib import Path
 
+import jwt
 from authlib.integrations.requests_client import OAuth2Session
 
 
@@ -13,18 +15,25 @@ def set_oauth_cb(config):
     Meanwhile, this is a pure Python implementation of the refresh token
     callback.
     """
-    if config.pop("sasl.oauthbearer.method", None) != "oidc":
-        return
 
     client_id = config.pop("sasl.oauthbearer.client.id")
-    client_secret = config.pop("sasl.oauthbearer.client.secret", None)
-    scope = config.pop("sasl.oauthbearer.scope", None)
-    token_endpoint = config.pop("sasl.oauthbearer.token.endpoint.url")
 
-    session = OAuth2Session(client_id, client_secret, scope=scope)
-
-    def oauth_cb(*_, **__):
-        token = session.fetch_token(token_endpoint, grant_type="client_credentials")
-        return token["access_token"], token["expires_at"]
-
-    config["oauth_cb"] = oauth_cb
+    client = OAuth2Session(client_id=client_id)
+    scope = config.pop("sasl.oauthbearer.scope")
+    url = config.pop("sasl.oauthbearer.token.endpoint.url")
+
+    def refresh_cognito_tokens():
+        home = Path.home()
+        with open(home.joinpath(".gcn", scope.replace("/", "_")), "r") as file:
+            token = file.read()
+        newToken = client.refresh_token(url, token)
+        return newToken
+
+    def oauthbearer_token_refresh_cb(*_, **__):
+        token_info = refresh_cognito_tokens()
+        jwt_token = token_info["access_token"]
+        decoded = jwt.decode(jwt_token, options={"verify_signature": False})
+        exp = decoded["exp"]
+        return jwt_token, exp
+
+    config["oauth_cb"] = oauthbearer_token_refresh_cb

test/test_oidc.py

@@ -1,33 +1,46 @@
-from unittest.mock import MagicMock
+from pathlib import Path
+from unittest.mock import MagicMock, mock_open, patch
 
-from gcn_kafka import oidc
+import jwt
 
 
-def test_no_oidc():
-    config = {}
-    oidc.set_oauth_cb(config)
-    assert config == {}
-
-
-def test_oidc(monkeypatch):
-    mock_session_class = MagicMock()
-    monkeypatch.setattr(oidc, "OAuth2Session", mock_session_class)
+def test_oidc_with_file(monkeypatch):
+    mock_fetch_token = MagicMock(
+        return_value={
+            "access_token": jwt.encode({"exp": 1234567890}, "secret", algorithm="HS256")
+        }
+    )
+    mock_session_instance = MagicMock(refresh_token=mock_fetch_token)
+    mock_session_class = MagicMock(return_value=mock_session_instance)
+    monkeypatch.setattr("gcn_kafka.oidc.OAuth2Session", mock_session_class)
 
+    # Setup config
     config = {
-        "sasl.oauthbearer.method": "oidc",
         "sasl.oauthbearer.client.id": "client_id",
-        "sasl.oauthbearer.client.secret": "client_secret",
-        "sasl.oauthbearer.scope": "scope",
-        "sasl.oauthbearer.token.endpoint.url": "token_endpoint",
+        "sasl.oauthbearer.scope": "my/scope",
+        "sasl.oauthbearer.token.endpoint.url": "https://example.com/token",
     }
-    oidc.set_oauth_cb(config)
 
-    oauth_cb = config.pop("oauth_cb")
-    assert config == {}
-    mock_session_class.assert_called_once_with(
-        "client_id", "client_secret", scope="scope"
-    )
-    oauth_cb()
-    mock_session_class.return_value.fetch_token.assert_called_once_with(
-        "token_endpoint", grant_type="client_credentials"
-    )
+    fake_home = Path("/fake/home")
+    fake_token = "fake-refresh-token"
+    with patch("pathlib.Path.home", return_value=fake_home):
+        with patch("builtins.open", mock_open(read_data=fake_token)) as m_open:
+            from gcn_kafka.oidc import set_oauth_cb
+
+            set_oauth_cb(config)
+
+            oauth_cb = config.pop("oauth_cb")
+            assert config == {}
+
+            token, exp = oauth_cb()
+
+            # Check file was opened at the correct path
+            expected_file = fake_home.joinpath(".gcn", "my_scope")
+            m_open.assert_called_once_with(expected_file, "r")
+
+            # Check OAuth2Session was used correctly
+            mock_fetch_token.assert_called_once_with(
+                "https://example.com/token", "fake-refresh-token"
+            )
+            assert isinstance(token, str)
+            assert exp == jwt.decode(token, options={"verify_signature": False})["exp"]