diff --git a/pyproject.toml b/pyproject.toml index afa662b..31a8e28 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -24,6 +24,7 @@ dependencies = [ "certifi", "confluent-kafka >= 2.2.0", "requests", + "pyjwt" ] requires-python = ">=3.9" dynamic = [ "version" ] diff --git a/src/gcn_kafka/core.py b/src/gcn_kafka/core.py index c350088..6ab883c 100644 --- a/src/gcn_kafka/core.py +++ b/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 @@ def get_config(mode, config, **kwargs): 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 @@ def update_config(config, **kwargs): 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 @@ def __init__( super().__init__( get_config( "producer", + scope, config, client_id=client_id, - client_secret=client_secret, domain=domain, **kwargs, ) @@ -89,9 +84,9 @@ def __init__( 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 @@ def __init__( super().__init__( get_config( "consumer", + scope, config, client_id=client_id, - client_secret=client_secret, domain=domain, **kwargs, ) @@ -121,7 +116,6 @@ def __init__( 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 @@ def __init__( get_config( "admin", config, + scope="gcn.nasa.gov/kafka-admin", client_id=client_id, - client_secret=client_secret, domain=domain, **kwargs, ) diff --git a/src/gcn_kafka/oidc.py b/src/gcn_kafka/oidc.py index 4a985ba..7ec853d 100644 --- a/src/gcn_kafka/oidc.py +++ b/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 diff --git a/test/test_oidc.py b/test/test_oidc.py index d33b3ef..faa7b5e 100644 --- a/test/test_oidc.py +++ b/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"]