Client credentials flow working

Author: ethho

datajoint/axon.py

@@ -6,6 +6,11 @@
 import sys
 import flask
 import webbrowser
+import requests_oauthlib
+import oauthlib
+from oauthlib.oauth2 import BackendApplicationClient
+from requests_oauthlib import OAuth2Session
+import requests
 import urllib
 import http.client
 import botocore
@@ -249,4 +254,69 @@ def refreshable_session(self) -> boto3.Session:
         session._credentials = refreshable_credentials
         autorefresh_session = boto3.Session(botocore_session=session)
 
-        return autorefresh_session
+        return autorefresh_session
+
+def get_s3_client(
+    aws_account_id: str,
+    s3_role: str,
+    auth_client_id: str,
+    auth_client_secret: str = None,
+    bearer_token: str = None,
+    well_known_url: str = "https://keycloak-qa.datajoint.io/realms/datajoint/.well-known/openid-configuration",
+):
+    """
+    Get S3 client with the given credentials.
+
+    Parameters
+    ----------
+    aws_account_id : str
+        AWS account ID
+
+    s3_role : str
+        S3 role
+
+    auth_client_id : str
+        Auth client ID
+
+    auth_client_secret : str (optional)
+        Auth client secret
+
+    bearer_token : str (optional)
+        Bearer token
+
+    well_known_url : str (optional)
+        Well-known URL for the OpenID configuration
+
+    Returns
+    -------
+    boto3.client
+        S3 client
+    """
+    # Get token URL from well-known URL
+    well_known_resp = requests.get(well_known_url)
+    assert well_known_resp.status_code == 200, f"Failed to get well-known URL: {well_known_url}"
+    well_known_data = well_known_resp.json()
+    token_url = well_known_data.get("token_endpoint")
+    assert token_url, f"Token URL not found in well-known data: {well_known_data=}"
+
+    # Client credentials flow
+    token = _client_credentials_flow(
+        auth_client_id, auth_client_secret, token_url
+    )
+
+    #
+
+
+def _client_credentials_flow(client_id, client_secret, token_url):
+    client = BackendApplicationClient(client_id=client_id)
+    oauth = OAuth2Session(client=client)
+    try:
+        return oauth.fetch_token(
+            token_url=token_url,
+            client_id=client_id,
+            client_secret=client_secret,
+        )
+    except oauthlib.oauth2.rfc6749.errors.UnauthorizedClientError as e:
+        msg = f"Error getting OAuth2 client: {e.description}"
+        log.error(msg)
+        raise ValueError(msg) from e

pyproject.toml

@@ -45,12 +45,15 @@ classifiers = [
 test = [
   "pytest",
   "pytest-cov",
+  "pytest-dotenv",
   "black==24.2.0",
   "flake8",
   "moto[s3]>=4.2.13",
 ]
 axon = [
   "boto3",
+  "requests_oauthlib",
+  "requests",
 ]
 
 [project.urls]

tests/test_axon.py

@@ -1,21 +1,106 @@
-from datajoint.axon import Session
+import os
+from datajoint.axon import Session, get_s3_client
+import json
 import pytest
 import boto3
 from moto import mock_aws
+import dotenv
+dotenv.load_dotenv(dotenv.find_dotenv())
 
 
 @pytest.fixture
 def moto_account_id():
     """Default account ID for moto"""
     return "123456789012"
 
+
+@pytest.fixture
+def keycloak_client_secret():
+    secret = os.getenv("OAUTH_CLIENT_SECRET")
+    if not secret:
+        pytest.skip("No client secret found")
+    else:
+        return secret
+
+
+@pytest.fixture
+def keycloak_client_id():
+    return os.getenv("OAUTH_CLIENT_ID", "works")
+
+
+@pytest.fixture(scope="function")
+def aws_credentials():
+    """Mocked AWS Credentials for moto."""
+    os.environ["AWS_ACCESS_KEY_ID"] = "testing"
+    os.environ["AWS_SECRET_ACCESS_KEY"] = "testing"
+    os.environ["AWS_SECURITY_TOKEN"] = "testing"
+    os.environ["AWS_SESSION_TOKEN"] = "testing"
+    os.environ["AWS_DEFAULT_REGION"] = "us-east-1"
+
+
+@pytest.fixture(scope="function")
+def s3_client(aws_credentials):
+    """
+    Return a mocked S3 client
+    """
+    with mock_aws():
+        yield boto3.client("s3", region_name="us-east-1")
+
+
+@pytest.fixture(scope="function")
+def iam_client(aws_credentials):
+    """
+    Return a mocked S3 client
+    """
+    with mock_aws():
+        yield boto3.client("iam", region_name="us-east-1")
+
+
+@pytest.fixture
+def s3_policy(iam_client):
+    """Create a policy with S3 read access using boto3."""
+    policy_doc = {
+        "Version": "2012-10-17",
+        "Statement": [
+            {
+                "Effect": "Allow",
+                "Action": "s3:GetObject",
+                "Resource": "arn:aws:s3:::mybucket/*",
+            }
+        ],
+    }
+    return iam_client.create_policy(
+        PolicyName="test-policy",
+        Path="/",
+        PolicyDocument=json.dumps(policy_doc),
+        Description="Test policy",
+    )
+
+@pytest.fixture
+def s3_role(moto_account_id, s3_policy):
+    """Create a mock role and policy document for testing"""
+    return "123456789012"
+
+
 @mock_aws
+@pytest.mark.skip
 class TestSession:
-    def test_can_init(self):
+    def test_can_init(self, s3_role, keycloak_client_id, keycloak_client_secret, moto_account_id):
         session = Session(
             aws_account_id=moto_account_id,
-            s3_role="test-role",
-            auth_client_id="test-client-id",
-            auth_client_secret="test-client-secret",
+            s3_role=s3_role,
+            auth_client_id=keycloak_client_id,
+            auth_client_secret=keycloak_client_secret,
         )
         assert session.bearer_token, "Bearer token not set"
+
+def test_get_s3_client(s3_role, keycloak_client_id, keycloak_client_secret, moto_account_id):
+    client = get_s3_client(
+        auth_client_id=keycloak_client_id,
+        auth_client_secret=keycloak_client_secret,
+        aws_account_id=moto_account_id,
+        s3_role=s3_role,
+        bearer_token=None,
+    )
+    assert client
+