diff --git a/specifications/SPEC_PLATFORM_SERVICE.md b/specifications/SPEC_PLATFORM_SERVICE.md index a0b198ac2..da5f4daf5 100644 --- a/specifications/SPEC_PLATFORM_SERVICE.md +++ b/specifications/SPEC_PLATFORM_SERVICE.md @@ -409,14 +409,15 @@ The Platform module provides foundational services but does not directly expose ### 6.1 Configuration Parameters -| Parameter | Type | Default | Description | Required | -| ------------------------------- | ---- | --------------------------------- | --------------------------- | -------- | -| `api_root` | str | `https://platform.aignostics.com` | Base URL of Aignostics API | Yes | -| `audience` | str | Environment-specific | OAuth audience claim | Yes | -| `scope` | str | `offline_access` | OAuth scopes required | Yes | -| `cache_dir` | str | User cache directory | Directory for token storage | No | -| `request_timeout_seconds` | int | 30 | API request timeout | No | -| `authorization_backoff_seconds` | int | 3 | Retry backoff time | No | +| Parameter | Type | Default | Description | Required | +| ------------------------------- | ---- | --------------------------------- | ---------------------------- | -------- | +| `api_root` | str | `https://platform.aignostics.com` | Base URL of Aignostics API | Yes | +| `audience` | str | Environment-specific | OAuth audience claim | Yes | +| `scope` | str | `offline_access` | OAuth scopes required | Yes | +| `organization_id` | str | None | Auth0 organization for OAuth | No | +| `cache_dir` | str | User cache directory | Directory for token storage | No | +| `request_timeout_seconds` | int | 30 | API request timeout | No | +| `authorization_backoff_seconds` | int | 3 | Retry backoff time | No | ### 6.2 Environment Variables @@ -432,6 +433,7 @@ The Platform module provides foundational services but does not directly expose | `AIGNOSTICS_DEVICE_URL` | Custom device authorization URL | `https://custom.auth0.com/oauth/device/code` | | `AIGNOSTICS_JWS_JSON_URL` | Custom JWS key set URL | `https://custom.auth0.com/.well-known/jwks.json` | | `AIGNOSTICS_CLIENT_ID_INTERACTIVE` | Interactive flow client ID | `interactive_client_123` | +| `AIGNOSTICS_ORGANIZATION_ID` | Auth0 organization for OAuth | `my-organization` | | `AIGNOSTICS_REFRESH_TOKEN` | Long-lived refresh token | `refresh_token_value` | | `AIGNOSTICS_CACHE_DIR` | Custom cache directory | `/custom/cache/path` | | `AIGNOSTICS_REQUEST_TIMEOUT_SECONDS` | API request timeout | `60` | @@ -492,7 +494,7 @@ The Platform module provides foundational services but does not directly expose ### 9.1 Key Algorithms and Business Logic -- **PKCE Flow**: OAuth 2.0 Authorization Code flow with Proof Key for Code Exchange for enhanced security in public clients +- **PKCE Flow**: OAuth 2.0 Authorization Code flow with Proof Key for Code Exchange for enhanced security in public clients. Supports optional `AIGNOSTICS_ORGANIZATION_ID` parameter for Auth0 organization-specific authentication flows - **Token Caching**: File-based token persistence with expiration tracking and automatic cleanup - **Health Monitoring**: Multi-layer health checks including public endpoint availability and authenticated API access diff --git a/src/aignostics/platform/_authentication.py b/src/aignostics/platform/_authentication.py index d76830079..3b8b101c4 100644 --- a/src/aignostics/platform/_authentication.py +++ b/src/aignostics/platform/_authentication.py @@ -333,10 +333,18 @@ def _perform_authorization_code_with_pkce_flow() -> str: # noqa: C901 redirect_uri=settings().redirect_uri, pkce="S256", ) + auth_params = { + "access_type": "offline", + "audience": settings().audience, + } + organization_id = settings().organization_id + + if organization_id: + auth_params["organization"] = organization_id + authorization_url, _ = session.authorization_url( settings().authorization_base_url, - access_type="offline", - audience=settings().audience, + **auth_params, ) authentication_result = AuthenticationResult() diff --git a/src/aignostics/platform/_settings.py b/src/aignostics/platform/_settings.py index bbaece783..5b51c50a6 100644 --- a/src/aignostics/platform/_settings.py +++ b/src/aignostics/platform/_settings.py @@ -206,6 +206,10 @@ def profile_edit_url(self) -> str: ] client_id_interactive: Annotated[str, Field(description="OAuth client ID for interactive flows")] + organization_id: Annotated[ + str | None, Field(description="Optional Auth0 organization ID parameter for the /authorize OAuth endpoint") + ] = None + @computed_field # type: ignore[prop-decorator] @property def tenant_domain(self) -> str: diff --git a/tests/aignostics/platform/authentication_test.py b/tests/aignostics/platform/authentication_test.py index d31dbfe15..874969ba5 100644 --- a/tests/aignostics/platform/authentication_test.py +++ b/tests/aignostics/platform/authentication_test.py @@ -61,6 +61,7 @@ def mock_settings() -> MagicMock: settings.auth_retry_attempts = 3 settings.auth_jwk_set_cache_ttl = 300 settings.refresh_token = None + settings.organization_id = None mock_settings.return_value = settings yield mock_settings @@ -339,25 +340,22 @@ def test_can_open_browser_false(record_property) -> None: class TestAuthorizationCodeFlow: """Test cases for the authorization code flow with PKCE.""" - @pytest.mark.unit @staticmethod - def test_perform_authorization_code_flow_success(record_property, mock_settings) -> None: - """Test successful authorization code flow with PKCE.""" - record_property("tested-item-id", "SPEC-PLATFORM-SERVICE") - # Mock OAuth session + def _run_pkce_flow() -> tuple[str | None, MagicMock]: + """Set up common PKCE flow mocks, run the flow, and return (token, mock_session). + + Returns: + tuple: The returned token and the OAuth2Session mock for further assertions. + """ mock_session = MagicMock(spec=OAuth2Session) mock_session.authorization_url.return_value = ("https://test.auth/authorize?code_challenge=abc", None) mock_session.fetch_token.return_value = {"access_token": "pkce.token"} - # Mock HTTP server mock_server = MagicMock() - - # Setup mocks for the redirect URI parsing mock_redirect_parsed = MagicMock() mock_redirect_parsed.hostname = "localhost" mock_redirect_parsed.port = 8000 - # Create a custom HTTPServer mock implementation that simulates a callback class MockHTTPServer: def __init__(self, *args, **kwargs) -> None: pass @@ -368,7 +366,6 @@ def __enter__(self) -> MagicMock: def __exit__(self, *args) -> None: pass - # Create a mock for the auth result mock_auth_result = MagicMock() mock_auth_result.token = "pkce.token" # noqa: S105 - Test credential mock_auth_result.error = None @@ -379,20 +376,19 @@ def __exit__(self, *args) -> None: patch("urllib.parse.urlparse", return_value=mock_redirect_parsed), patch("aignostics.platform._authentication.AuthenticationResult", return_value=mock_auth_result), ): - # Simulate a successful server response by making handle_request set the token - def handle_request_side_effect(): - # This simulates what the HTTP handler would do on success - mock_auth_result.token = "pkce.token" # noqa: S105 - Test credential - - mock_server.handle_request.side_effect = handle_request_side_effect - - # Call the function under test + mock_server.handle_request.side_effect = lambda: None token = _perform_authorization_code_with_pkce_flow() - # Assertions - assert token == "pkce.token" # noqa: S105 - Test credential - mock_server.handle_request.assert_called_once() - mock_session.authorization_url.assert_called_once() + return token, mock_session + + @pytest.mark.unit + @staticmethod + def test_perform_authorization_code_flow_success(record_property, mock_settings) -> None: + """Test successful authorization code flow with PKCE.""" + record_property("tested-item-id", "SPEC-PLATFORM-SERVICE") + token, mock_session = TestAuthorizationCodeFlow._run_pkce_flow() + assert token == "pkce.token" # noqa: S105 - Test credential + mock_session.authorization_url.assert_called_once() @pytest.mark.unit @staticmethod @@ -466,6 +462,27 @@ def handle_request_side_effect(): with pytest.raises(RuntimeError, match=AUTHENTICATION_FAILED): _perform_authorization_code_with_pkce_flow() + @pytest.mark.unit + @staticmethod + def test_perform_authorization_code_flow_with_organization(record_property, mock_settings) -> None: + """Test authorization code flow includes organization parameter when set.""" + record_property("tested-item-id", "SPEC-PLATFORM-SERVICE") + mock_settings.return_value.organization_id = "test-org" + token, mock_session = TestAuthorizationCodeFlow._run_pkce_flow() + assert token == "pkce.token" # noqa: S105 + assert mock_session.authorization_url.call_args[1].get("organization") == "test-org" + + @pytest.mark.unit + @staticmethod + def test_perform_authorization_code_flow_without_organization(record_property, mock_settings) -> None: + """Test authorization code flow omits organization parameter when unset.""" + record_property("tested-item-id", "SPEC-PLATFORM-SERVICE") + # organization_id is None by default in mock_settings fixture + assert mock_settings.return_value.organization_id is None + token, mock_session = TestAuthorizationCodeFlow._run_pkce_flow() + assert token == "pkce.token" # noqa: S105 + assert "organization" not in mock_session.authorization_url.call_args[1] + class TestDeviceFlow: """Test cases for the device flow authentication."""