From e31ad2c5e84cb6dc826255b7f3ab2e1efee1d69d Mon Sep 17 00:00:00 2001 From: Tamar Weisskopf <98809100+TamarW0@users.noreply.github.com> Date: Tue, 4 Aug 2026 15:58:59 +0300 Subject: [PATCH 1/2] feat: Add AWS Cognito as OIDC Identity Provider (#325) --------- Signed-off-by: Zvi Grinberg Co-authored-by: Tamar Weisskopf Co-authored-by: Zvi Grinberg --- kustomize/base/exploit-iq-config.yml | 4 +++ .../utils/credential_client.py | 21 +++++++++-- .../configs/config-http-openai.yml | 4 +++ .../functions/cve_clone_and_deps.py | 10 ++++-- .../functions/cve_generate_vdbs.py | 7 ++-- .../functions/cve_http_output.py | 36 ++++++++++++++++++- .../functions/cve_segmentation.py | 7 ++-- 7 files changed, 79 insertions(+), 10 deletions(-) diff --git a/kustomize/base/exploit-iq-config.yml b/kustomize/base/exploit-iq-config.yml index 9128bd77c..921ceef05 100644 --- a/kustomize/base/exploit-iq-config.yml +++ b/kustomize/base/exploit-iq-config.yml @@ -166,6 +166,10 @@ functions: verify_path: /app/certs/service-ca.crt keycloak_server: ${KC_SERVER} keycloak_realm: ${KC_REALM:-quarkus} + cognito_domain: ${COGNITO_DOMAIN} + cognito_scope: ${COGNITO_SCOPE} + cognito_client_id: ${COGNITO_CLIENT_ID} + cognito_client_secret: ${COGNITO_CLIENT_SECRET} client_id: ${KC_CLIENT_ID:-exploit-iq-client} client_secret: ${KC_CLIENT_SECRET} verify_path_keycloak: ${VERIFY_PATH_KEYCLOAK} diff --git a/src/exploit_iq_commons/utils/credential_client.py b/src/exploit_iq_commons/utils/credential_client.py index 775f418db..306e8e9a2 100644 --- a/src/exploit_iq_commons/utils/credential_client.py +++ b/src/exploit_iq_commons/utils/credential_client.py @@ -32,6 +32,7 @@ AES_256_KEY_SIZE_BYTES = 32 _credential_id_ctx: ContextVar[str | None] = ContextVar("credential_id", default=None) +_http_auth_header_ctx: ContextVar[str | None] = ContextVar("http_auth_header", default=None) @contextmanager @@ -52,6 +53,18 @@ def credential_context(credential_id: str | None) -> Generator[None]: _credential_id_ctx.reset(token) +@contextmanager +def http_auth_header_context(auth_header: str | None) -> Generator[None]: + """Make a pre-resolved Authorization header available to + fetch_and_decrypt_credential via ContextVar. When set, the header is + used instead of the SA token / JWT fallback.""" + token = _http_auth_header_ctx.set(auth_header) + try: + yield + finally: + _http_auth_header_ctx.reset(token) + + def _resolve_jwt_token(jwt_token: str | None) -> str: """ Resolve JWT token for authenticating with the credential backend. @@ -181,9 +194,13 @@ def fetch_and_decrypt_credential( RuntimeError Unexpected HTTP status or network error. """ - resolved_token = _resolve_jwt_token(jwt_token) + auth_header = _http_auth_header_ctx.get() + if auth_header is None: + resolved_token = _resolve_jwt_token(jwt_token) + auth_header = f"Bearer {resolved_token}" + url = f"{backend_url.rstrip('/')}/api/v1/credentials/{credential_id}" - headers = {"Authorization": f"Bearer {resolved_token}"} + headers = {"Authorization": auth_header} logger.info("Fetching credential: credential_id=%s", credential_id) diff --git a/src/vuln_analysis/configs/config-http-openai.yml b/src/vuln_analysis/configs/config-http-openai.yml index ff247e888..b132e35ab 100644 --- a/src/vuln_analysis/configs/config-http-openai.yml +++ b/src/vuln_analysis/configs/config-http-openai.yml @@ -159,6 +159,10 @@ functions: auth_type: ${AUTH_TYPE:-disabled} keycloak_server: ${KC_SERVER:-http://localhost:8180} keycloak_realm: ${KC_REALM:-quarkus} + cognito_domain: ${COGNITO_DOMAIN} + cognito_scope: ${COGNITO_SCOPE} + cognito_client_id: ${COGNITO_CLIENT_ID} + cognito_client_secret: ${COGNITO_CLIENT_SECRET} client_id: ${KC_CLIENT_ID:-exploit-iq-client} client_secret: ${KC_CLIENT_SECRET:-example-credentials} verify_path_keycloak: ${VERIFY_PATH_KEYCLOAK} diff --git a/src/vuln_analysis/functions/cve_clone_and_deps.py b/src/vuln_analysis/functions/cve_clone_and_deps.py index 29ec8ed00..f9cbf954d 100644 --- a/src/vuln_analysis/functions/cve_clone_and_deps.py +++ b/src/vuln_analysis/functions/cve_clone_and_deps.py @@ -31,10 +31,13 @@ from exploit_iq_commons.data_models.common import AnalysisType from exploit_iq_commons.logging.loggers_factory import LoggingFactory, trace_id -from exploit_iq_commons.utils.credential_client import credential_context +from exploit_iq_commons.utils.credential_client import credential_context, http_auth_header_context +from vuln_analysis.functions.cve_http_output import get_auth_header, HTTP_OUTPUT_AGENT_CONFIG from exploit_iq_commons.utils.dep_tree import detect_ecosystem from exploit_iq_commons.utils.git_utils import resolve_path_to_manifest + + logger = LoggingFactory.get_agent_logger(__name__) @@ -81,7 +84,6 @@ async def clone_and_deps(config: CVECloneAndDepsConfig, builder: Builder): git_directory=config.base_git_dir, pickle_cache_directory=config.base_pickle_dir, ) - async def _arun(message: ExploitIqInput) -> ExploitIqEngineInput: """ Clone repositories and install dependencies. @@ -101,7 +103,9 @@ async def _arun(message: ExploitIqInput) -> ExploitIqEngineInput: message.scan.id, ) - with credential_context(message.credential_id): + http_output_config = builder.get_function_config(HTTP_OUTPUT_AGENT_CONFIG) + auth_header = get_auth_header(http_output_config) + with http_auth_header_context(auth_header), credential_context(message.credential_id): # Configure RPM manager for IMAGE analysis if message.image.analysis_type == AnalysisType.IMAGE and isinstance( sbom_infos, ManualSBOMInfoInput diff --git a/src/vuln_analysis/functions/cve_generate_vdbs.py b/src/vuln_analysis/functions/cve_generate_vdbs.py index 1f946013f..fb858bf36 100644 --- a/src/vuln_analysis/functions/cve_generate_vdbs.py +++ b/src/vuln_analysis/functions/cve_generate_vdbs.py @@ -29,7 +29,8 @@ from exploit_iq_commons.data_models.common import AnalysisType from exploit_iq_commons.logging.loggers_factory import LoggingFactory, trace_id -from exploit_iq_commons.utils.credential_client import credential_context +from exploit_iq_commons.utils.credential_client import credential_context, http_auth_header_context +from vuln_analysis.functions.cve_http_output import get_auth_header, HTTP_OUTPUT_AGENT_CONFIG from exploit_iq_commons.utils.dep_tree import Ecosystem, detect_ecosystem from exploit_iq_commons.utils.git_utils import resolve_path_to_manifest from vuln_analysis.tools.tool_names import ToolNames @@ -221,7 +222,9 @@ async def _arun(message: ExploitIqInput) -> ExploitIqEngineInput: trace_id.set(message.scan.id) logger.debug("_arun: received credential_id=%r scan_id=%s", message.credential_id, message.scan.id) # Build VDBs (credential_id is propagated via async context) - with credential_context(message.credential_id): + http_output_config = builder.get_function_config(HTTP_OUTPUT_AGENT_CONFIG) + auth_header = get_auth_header(http_output_config) + with http_auth_header_context(auth_header), credential_context(message.credential_id): logger.debug("_arun: credential_context entered, credential_id=%r", message.credential_id) # When ignore_code_embedding is True, also skip doc VDBs vdb_source_infos = ( diff --git a/src/vuln_analysis/functions/cve_http_output.py b/src/vuln_analysis/functions/cve_http_output.py index d2dde55a2..a37f6e16d 100644 --- a/src/vuln_analysis/functions/cve_http_output.py +++ b/src/vuln_analysis/functions/cve_http_output.py @@ -36,6 +36,8 @@ import os import re +HTTP_OUTPUT_AGENT_CONFIG = "cve_http_output" + if TYPE_CHECKING: from vuln_analysis.data_models.output import ExploitIqOutput, FailureReport @@ -89,7 +91,7 @@ class CVEHttpOutputConfig(FunctionBaseConfig, name="cve_http_output"): """ url: str = Field(description="URL to send CVE workflow output") endpoint: str = Field(description="Endpoint to send CVE workflow output") - auth_type: str = Field(default="disabled", description="Type of auth - bearer, basic, keycloak or disabled") + auth_type: str = Field(default="disabled", description="Type of auth - bearer, basic, keycloak, cognito or disabled") token: str | None = Field(default=None, description="Token to authenticate when sending CVE workflow output") token_path: str | None = Field(default=None, description="Path to token file containing auth token") verify_path: str | None = Field(default=None, description="Path to certificate to validate the token key found in ") @@ -98,6 +100,10 @@ class CVEHttpOutputConfig(FunctionBaseConfig, name="cve_http_output"): keycloak_server: str | None = Field(default=None, description="Keycloak server URL (e.g. https://keycloak.example.com)") keycloak_realm: str | None = Field(default=None, description="Keycloak realm name") verify_path_keycloak: str | None = Field(default=None, description="Path to ca to validate the certificate of keycloak instance ") + cognito_domain: str | None = Field(default=None, description="Cognito domain (e.g. myapp.auth.us-east-1.amazoncognito.com)") + cognito_scope: str | None = Field(default=None, description="Cognito custom scope (e.g. api/read)") + cognito_client_id: str | None = Field(default=None, description="OAuth2 client ID for Cognito M2M authentication") + cognito_client_secret: str | None = Field(default=None, description="OAuth2 client secret for Cognito M2M authentication") client_id: str | None = Field(default=None, description="OAuth2 client ID for keycloak authentication") client_secret: str | None = Field(default=None, description="OAuth2 client secret for keycloak authentication") failure_endpoint: str = Field(default="/api/v1/reports/failed", @@ -277,6 +283,28 @@ def _fetch_keycloak_token(http_config: CVEHttpOutputConfig) -> str | None: return None +def _fetch_cognito_token(http_config: CVEHttpOutputConfig) -> str | None: + token_url = f"{http_config.cognito_domain}/oauth2/token" + # Cognito requires Basic auth header for client_credentials + credentials = base64.b64encode( + f"{http_config.cognito_client_id}:{http_config.cognito_client_secret}".encode() + ).decode() + headers = { + "Content-Type": "application/x-www-form-urlencoded", + "Authorization": f"Basic {credentials}", + } + data = {"grant_type": "client_credentials"} + if http_config.cognito_scope: + data["scope"] = http_config.cognito_scope + try: + resp = requests.post(token_url, headers=headers, data=data, timeout=30) + resp.raise_for_status() + return resp.json()["access_token"] + except Exception as e: + logger.error("Unable to obtain Cognito access token from %s: %s", token_url, e) + return None + + def get_auth_header(http_config: CVEHttpOutputConfig | None) -> str | None: match http_config.auth_type: case "basic": @@ -304,6 +332,12 @@ def get_auth_header(http_config: CVEHttpOutputConfig | None) -> str | None: except Exception as e: logger.warn(f"Unable to read OAuth token: {e}") return None + case "cognito": + if not all([http_config.cognito_domain, http_config.cognito_client_id, http_config.cognito_client_secret]): + logger.error("Cognito auth requires cognito_domain, cognito_client_id, and cognito_client_secret") + return None + token = _fetch_cognito_token(http_config) + return f"Bearer {token}" if token else None case None: return None diff --git a/src/vuln_analysis/functions/cve_segmentation.py b/src/vuln_analysis/functions/cve_segmentation.py index de75f94c8..2bca05fb4 100644 --- a/src/vuln_analysis/functions/cve_segmentation.py +++ b/src/vuln_analysis/functions/cve_segmentation.py @@ -36,7 +36,8 @@ from pydantic import Field from exploit_iq_commons.logging.loggers_factory import LoggingFactory, trace_id -from exploit_iq_commons.utils.credential_client import credential_context +from exploit_iq_commons.utils.credential_client import credential_context, http_auth_header_context +from vuln_analysis.functions.cve_http_output import get_auth_header, HTTP_OUTPUT_AGENT_CONFIG from vuln_analysis.tools.tool_names import ToolNames logger = LoggingFactory.get_agent_logger(__name__) @@ -224,7 +225,9 @@ async def _arun(state: ExploitIqEngineInput) -> ExploitIqEngineInput: message.scan.id, ) - with credential_context(message.credential_id): + http_output_config = builder.get_function_config(HTTP_OUTPUT_AGENT_CONFIG) + auth_header = get_auth_header(http_output_config) + with http_auth_header_context(auth_header), credential_context(message.credential_id): vdb_code_path, vdb_doc_path = await asyncio.to_thread( embedder.build_vdbs, source_infos, From 7402e3acb8168547750a0248fc9684667da2d2b3 Mon Sep 17 00:00:00 2001 From: Tamar Weisskopf Date: Wed, 9 Sep 2026 04:19:36 +0300 Subject: [PATCH 2/2] test: add unit tests for AWS Cognito authentication Comprehensive tests for OAuth2 client_credentials flow, token fetching, and configuration validation. 21 tests pass covering success/error paths and edge cases. Relates to: TC-5942 Co-Authored-By: Claude Sonnet 4.5 --- .../functions/tests/test_cognito_auth.py | 235 ++++++++++++++++++ 1 file changed, 235 insertions(+) create mode 100644 src/vuln_analysis/functions/tests/test_cognito_auth.py diff --git a/src/vuln_analysis/functions/tests/test_cognito_auth.py b/src/vuln_analysis/functions/tests/test_cognito_auth.py new file mode 100644 index 000000000..513b11472 --- /dev/null +++ b/src/vuln_analysis/functions/tests/test_cognito_auth.py @@ -0,0 +1,235 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Unit tests for AWS Cognito authentication in cve_http_output module.""" + +import base64 +from unittest.mock import Mock, patch +import pytest +import requests + +from vuln_analysis.functions.cve_http_output import ( + CVEHttpOutputConfig, + MLOpsConfig, + _fetch_cognito_token, + get_auth_header, +) + + +@pytest.fixture +def cognito_config(): + """Fixture providing a valid Cognito configuration.""" + return CVEHttpOutputConfig( + url="https://api.example.com", + endpoint="/api/v1/reports", + auth_type="cognito", + cognito_domain="https://myapp.auth.us-east-1.amazoncognito.com", + cognito_client_id="client123", + cognito_client_secret="secret456", + mlops_config=MLOpsConfig(mlops_url="https://mlops.example.com"), + ) + + +@pytest.fixture +def mock_cognito_response(): + """Fixture providing a successful mock Cognito response.""" + mock_response = Mock() + mock_response.json.return_value = {"access_token": "test-access-token"} + mock_response.raise_for_status = Mock() + return mock_response + + +class TestFetchCognitoToken: + """Tests for _fetch_cognito_token function - core OAuth2 client_credentials flow.""" + + def test_constructs_correct_token_url(self, cognito_config, mock_cognito_response): + """Token URL should be {domain}/oauth2/token.""" + with patch('requests.post', return_value=mock_cognito_response) as mock_post: + _fetch_cognito_token(cognito_config) + + assert mock_post.call_args[0][0] == "https://myapp.auth.us-east-1.amazoncognito.com/oauth2/token" + + def test_sends_basic_auth_with_client_credentials(self, cognito_config, mock_cognito_response): + """Client ID and secret should be base64-encoded in Authorization header.""" + with patch('requests.post', return_value=mock_cognito_response) as mock_post: + _fetch_cognito_token(cognito_config) + + headers = mock_post.call_args[1]['headers'] + expected_creds = base64.b64encode(b"client123:secret456").decode() + assert headers['Authorization'] == f"Basic {expected_creds}" + assert headers['Content-Type'] == "application/x-www-form-urlencoded" + + def test_sends_client_credentials_grant_type(self, cognito_config, mock_cognito_response): + """Request data should contain grant_type=client_credentials.""" + with patch('requests.post', return_value=mock_cognito_response) as mock_post: + _fetch_cognito_token(cognito_config) + + data = mock_post.call_args[1]['data'] + assert data['grant_type'] == "client_credentials" + + def test_includes_scope_when_configured(self, cognito_config, mock_cognito_response): + """Custom scope should be included in request when set.""" + cognito_config.cognito_scope = "api/read api/write" + + with patch('requests.post', return_value=mock_cognito_response) as mock_post: + _fetch_cognito_token(cognito_config) + + data = mock_post.call_args[1]['data'] + assert data['scope'] == "api/read api/write" + + def test_omits_scope_when_not_configured(self, cognito_config, mock_cognito_response): + """Scope should not be in request when cognito_scope is None.""" + with patch('requests.post', return_value=mock_cognito_response) as mock_post: + _fetch_cognito_token(cognito_config) + + data = mock_post.call_args[1]['data'] + assert 'scope' not in data + + def test_returns_access_token_on_success(self, cognito_config): + """Should extract and return access_token from JSON response.""" + mock_resp = Mock() + mock_resp.json.return_value = { + "access_token": "eyJhbGci...", + "token_type": "Bearer", + "expires_in": 3600 + } + mock_resp.raise_for_status = Mock() + + with patch('requests.post', return_value=mock_resp): + token = _fetch_cognito_token(cognito_config) + + assert token == "eyJhbGci..." + + def test_uses_30_second_timeout(self, cognito_config, mock_cognito_response): + """Request should have a 30-second timeout to prevent hanging.""" + with patch('requests.post', return_value=mock_cognito_response) as mock_post: + _fetch_cognito_token(cognito_config) + + assert mock_post.call_args[1]['timeout'] == 30 + + @pytest.mark.parametrize("error", [ + requests.exceptions.HTTPError("401 Unauthorized"), + requests.exceptions.ConnectionError("Network unreachable"), + requests.exceptions.Timeout("Request timeout"), + ]) + def test_returns_none_on_request_errors(self, cognito_config, error): + """Any request exception should return None instead of raising.""" + with patch('requests.post') as mock_post: + if isinstance(error, requests.exceptions.HTTPError): + mock_resp = Mock() + mock_resp.raise_for_status.side_effect = error + mock_post.return_value = mock_resp + else: + mock_post.side_effect = error + + token = _fetch_cognito_token(cognito_config) + + assert token is None + + def test_returns_none_on_missing_access_token_in_response(self, cognito_config): + """Should return None if response doesn't contain access_token key.""" + mock_resp = Mock() + mock_resp.json.return_value = {"token_type": "Bearer"} # Missing access_token + mock_resp.raise_for_status = Mock() + + with patch('requests.post', return_value=mock_resp): + token = _fetch_cognito_token(cognito_config) + + assert token is None + + def test_returns_none_on_invalid_json_response(self, cognito_config): + """Should return None if JSON parsing fails.""" + mock_resp = Mock() + mock_resp.json.side_effect = ValueError("Invalid JSON") + mock_resp.raise_for_status = Mock() + + with patch('requests.post', return_value=mock_resp): + token = _fetch_cognito_token(cognito_config) + + assert token is None + + +class TestGetAuthHeaderCognito: + """Tests for get_auth_header with auth_type='cognito'.""" + + def test_returns_bearer_token_on_success(self, cognito_config): + """Should return Bearer header when token fetch succeeds.""" + with patch('vuln_analysis.functions.cve_http_output._fetch_cognito_token', return_value="test-token"): + header = get_auth_header(cognito_config) + + assert header == "Bearer test-token" + + def test_returns_none_when_token_fetch_fails(self, cognito_config): + """Should return None when _fetch_cognito_token returns None.""" + with patch('vuln_analysis.functions.cve_http_output._fetch_cognito_token', return_value=None): + header = get_auth_header(cognito_config) + + assert header is None + + @pytest.mark.parametrize("missing_field,config_override", [ + ("cognito_domain", {"cognito_domain": None}), + ("cognito_client_id", {"cognito_client_id": None}), + ("cognito_client_secret", {"cognito_client_secret": None}), + ("cognito_domain", {"cognito_domain": ""}), # Empty string treated as missing + ]) + def test_returns_none_when_required_config_missing(self, cognito_config, missing_field, config_override): + """Should validate required fields and return None if any are missing.""" + for key, value in config_override.items(): + setattr(cognito_config, key, value) + + with patch('vuln_analysis.functions.cve_http_output._fetch_cognito_token') as mock_fetch: + header = get_auth_header(cognito_config) + + assert header is None + mock_fetch.assert_not_called() # Should not attempt fetch with incomplete config + + +class TestCognitoConfigFields: + """Tests for Cognito-related configuration fields.""" + + def test_cognito_fields_are_optional_with_defaults(self): + """All Cognito fields should default to None and not be required.""" + config = CVEHttpOutputConfig( + url="https://api.example.com", + endpoint="/api/v1/reports", + mlops_config=MLOpsConfig(mlops_url="https://mlops.example.com"), + ) + + assert config.cognito_domain is None + assert config.cognito_scope is None + assert config.cognito_client_id is None + assert config.cognito_client_secret is None + + def test_cognito_fields_accept_and_store_values(self): + """Cognito fields should accept string values when provided.""" + config = CVEHttpOutputConfig( + url="https://api.example.com", + endpoint="/api/v1/reports", + cognito_domain="https://test.auth.region.amazoncognito.com", + cognito_scope="custom/scope", + cognito_client_id="test-client", + cognito_client_secret="test-secret", + mlops_config=MLOpsConfig(mlops_url="https://mlops.example.com"), + ) + + assert config.cognito_domain == "https://test.auth.region.amazoncognito.com" + assert config.cognito_scope == "custom/scope" + assert config.cognito_client_id == "test-client" + assert config.cognito_client_secret == "test-secret" + + def test_auth_type_field_documents_cognito(self): + """The auth_type field description should mention 'cognito' as a valid option.""" + field_info = CVEHttpOutputConfig.model_fields['auth_type'] + assert "cognito" in field_info.description