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 01/14] 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 02/14] 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 From cfeedf20e78a0f8bf4827ac36863eceac1d1e55b Mon Sep 17 00:00:00 2001 From: Tamar Weisskopf Date: Wed, 9 Sep 2026 12:31:50 +0300 Subject: [PATCH 03/14] fix: address code review feedback for Cognito auth Fixes based on Theodor's review: 1. Add wildcard case to get_auth_header - logs error for unrecognized auth_type values (e.g. typos like 'Cognito') 2. Handle None config in get_auth_header to prevent AttributeError when config key is missing 3. Auto-prepend https:// to cognito_domain if missing, update field description to accept both formats 4. Replace broad Exception with specific exception handling: - requests.RequestException for network/HTTP errors - KeyError/ValueError for invalid JSON responses Tests updated to cover all new behaviors (27 tests pass). Co-Authored-By: Claude Sonnet 4.5 --- .../functions/cve_http_output.py | 25 +++++++-- .../functions/tests/test_cognito_auth.py | 52 +++++++++++++++++++ 2 files changed, 73 insertions(+), 4 deletions(-) diff --git a/src/vuln_analysis/functions/cve_http_output.py b/src/vuln_analysis/functions/cve_http_output.py index a37f6e16d..ff267c91b 100644 --- a/src/vuln_analysis/functions/cve_http_output.py +++ b/src/vuln_analysis/functions/cve_http_output.py @@ -100,7 +100,7 @@ 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_domain: str | None = Field(default=None, description="Cognito domain (e.g. https://myapp.auth.us-east-1.amazoncognito.com or 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") @@ -284,7 +284,12 @@ def _fetch_keycloak_token(http_config: CVEHttpOutputConfig) -> str | None: def _fetch_cognito_token(http_config: CVEHttpOutputConfig) -> str | None: - token_url = f"{http_config.cognito_domain}/oauth2/token" + # Ensure domain has https:// scheme (Cognito always uses HTTPS) + domain = http_config.cognito_domain.rstrip("/") + if not domain.startswith("https://"): + domain = f"https://{domain}" + token_url = f"{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() @@ -300,12 +305,18 @@ def _fetch_cognito_token(http_config: CVEHttpOutputConfig) -> str | None: resp = requests.post(token_url, headers=headers, data=data, timeout=30) resp.raise_for_status() return resp.json()["access_token"] - except Exception as e: + except requests.RequestException as e: logger.error("Unable to obtain Cognito access token from %s: %s", token_url, e) return None + except (KeyError, ValueError) as e: + logger.error("Invalid Cognito token response from %s: %s", token_url, e) + return None def get_auth_header(http_config: CVEHttpOutputConfig | None) -> str | None: + if http_config is None: + return None + match http_config.auth_type: case "basic": if http_config.username and http_config.password: @@ -338,7 +349,13 @@ def get_auth_header(http_config: CVEHttpOutputConfig | None) -> str | None: return None token = _fetch_cognito_token(http_config) return f"Bearer {token}" if token else None - case None: + case "disabled" | None: + return None + case _: + logger.error( + "Unrecognized auth_type '%s'. Valid values: basic, bearer, keycloak, cognito, disabled", + http_config.auth_type, + ) return None diff --git a/src/vuln_analysis/functions/tests/test_cognito_auth.py b/src/vuln_analysis/functions/tests/test_cognito_auth.py index 513b11472..e20906fe1 100644 --- a/src/vuln_analysis/functions/tests/test_cognito_auth.py +++ b/src/vuln_analysis/functions/tests/test_cognito_auth.py @@ -160,10 +160,45 @@ def test_returns_none_on_invalid_json_response(self, cognito_config): assert token is None + def test_auto_prepends_https_when_missing(self, cognito_config, mock_cognito_response): + """Should automatically prepend https:// if domain doesn't have it.""" + cognito_config.cognito_domain = "myapp.auth.us-east-1.amazoncognito.com" # No https:// + + with patch('requests.post', return_value=mock_cognito_response) as mock_post: + _fetch_cognito_token(cognito_config) + + actual_url = mock_post.call_args[0][0] + assert actual_url == "https://myapp.auth.us-east-1.amazoncognito.com/oauth2/token" + + def test_preserves_https_when_already_present(self, cognito_config, mock_cognito_response): + """Should not double-prepend https:// if already present.""" + cognito_config.cognito_domain = "https://myapp.auth.us-east-1.amazoncognito.com" + + with patch('requests.post', return_value=mock_cognito_response) as mock_post: + _fetch_cognito_token(cognito_config) + + actual_url = mock_post.call_args[0][0] + assert actual_url == "https://myapp.auth.us-east-1.amazoncognito.com/oauth2/token" + + def test_strips_trailing_slash_from_domain(self, cognito_config, mock_cognito_response): + """Should strip trailing slash from domain to avoid double slashes.""" + cognito_config.cognito_domain = "https://myapp.auth.us-east-1.amazoncognito.com/" + + with patch('requests.post', return_value=mock_cognito_response) as mock_post: + _fetch_cognito_token(cognito_config) + + actual_url = mock_post.call_args[0][0] + assert actual_url == "https://myapp.auth.us-east-1.amazoncognito.com/oauth2/token" + class TestGetAuthHeaderCognito: """Tests for get_auth_header with auth_type='cognito'.""" + def test_returns_none_when_config_is_none(self): + """Should handle None config gracefully without crashing.""" + header = get_auth_header(None) + assert header is None + 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"): @@ -178,6 +213,23 @@ def test_returns_none_when_token_fetch_fails(self, cognito_config): assert header is None + def test_returns_none_for_disabled_auth_type(self, cognito_config): + """Should return None when auth_type is 'disabled'.""" + cognito_config.auth_type = "disabled" + header = get_auth_header(cognito_config) + assert header is None + + def test_logs_error_for_unrecognized_auth_type(self, cognito_config): + """Should log error and return None for typos/invalid auth_type.""" + cognito_config.auth_type = "Cognito" # Capital C typo + + with patch('vuln_analysis.functions.cve_http_output.logger') as mock_logger: + header = get_auth_header(cognito_config) + + assert header is None + mock_logger.error.assert_called_once() + assert "Unrecognized auth_type" in mock_logger.error.call_args[0][0] + @pytest.mark.parametrize("missing_field,config_override", [ ("cognito_domain", {"cognito_domain": None}), ("cognito_client_id", {"cognito_client_id": None}), From b2cd1bbefbe7ba57fb8a2c552fb7c0b13f296a4f Mon Sep 17 00:00:00 2001 From: Tamar Weisskopf Date: Wed, 9 Sep 2026 12:58:46 +0300 Subject: [PATCH 04/14] test: improve config validation test clarity - Specify auth_type='disabled' in test to clarify Cognito fields are optional when NOT using Cognito auth - Add test verifying required field validation when auth_type='cognito' 28 tests now pass. --- .../functions/tests/test_cognito_auth.py | 17 ++++++++++++++++- 1 file changed, 16 insertions(+), 1 deletion(-) diff --git a/src/vuln_analysis/functions/tests/test_cognito_auth.py b/src/vuln_analysis/functions/tests/test_cognito_auth.py index e20906fe1..90b01c306 100644 --- a/src/vuln_analysis/functions/tests/test_cognito_auth.py +++ b/src/vuln_analysis/functions/tests/test_cognito_auth.py @@ -252,10 +252,11 @@ 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.""" + """Cognito fields should default to None when auth_type is not 'cognito'.""" config = CVEHttpOutputConfig( url="https://api.example.com", endpoint="/api/v1/reports", + auth_type="disabled", # NOT cognito mlops_config=MLOpsConfig(mlops_url="https://mlops.example.com"), ) @@ -264,6 +265,20 @@ def test_cognito_fields_are_optional_with_defaults(self): assert config.cognito_client_id is None assert config.cognito_client_secret is None + def test_cognito_fields_required_when_auth_type_is_cognito(self): + """When auth_type='cognito', get_auth_header should validate required fields.""" + config = CVEHttpOutputConfig( + url="https://api.example.com", + endpoint="/api/v1/reports", + auth_type="cognito", + # Cognito fields are None (missing) + mlops_config=MLOpsConfig(mlops_url="https://mlops.example.com"), + ) + + # Runtime validation should return None when required fields are missing + header = get_auth_header(config) + assert header is None + def test_cognito_fields_accept_and_store_values(self): """Cognito fields should accept string values when provided.""" config = CVEHttpOutputConfig( From 7a5bf08cc0ab998c1631d2554a7a180261983f55 Mon Sep 17 00:00:00 2001 From: Tamar Weisskopf Date: Wed, 9 Sep 2026 13:03:02 +0300 Subject: [PATCH 05/14] test: remove redundant Pydantic field storage test Removed test_cognito_fields_accept_and_store_values - testing that Pydantic stores values is the framework's responsibility, not ours. 27 focused tests remain. --- .../functions/tests/test_cognito_auth.py | 17 ----------------- 1 file changed, 17 deletions(-) diff --git a/src/vuln_analysis/functions/tests/test_cognito_auth.py b/src/vuln_analysis/functions/tests/test_cognito_auth.py index 90b01c306..edfbbcf19 100644 --- a/src/vuln_analysis/functions/tests/test_cognito_auth.py +++ b/src/vuln_analysis/functions/tests/test_cognito_auth.py @@ -279,23 +279,6 @@ def test_cognito_fields_required_when_auth_type_is_cognito(self): header = get_auth_header(config) assert header 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'] From c36975c89bface7187afd54d2ceacbc1b5455ef0 Mon Sep 17 00:00:00 2001 From: Tamar Weisskopf Date: Wed, 9 Sep 2026 13:09:28 +0300 Subject: [PATCH 06/14] chore: remove redundant comment Remove comment explaining https:// prepending as the code is self-explanatory. --- src/vuln_analysis/functions/cve_http_output.py | 1 - 1 file changed, 1 deletion(-) diff --git a/src/vuln_analysis/functions/cve_http_output.py b/src/vuln_analysis/functions/cve_http_output.py index ff267c91b..62448a1f4 100644 --- a/src/vuln_analysis/functions/cve_http_output.py +++ b/src/vuln_analysis/functions/cve_http_output.py @@ -284,7 +284,6 @@ def _fetch_keycloak_token(http_config: CVEHttpOutputConfig) -> str | None: def _fetch_cognito_token(http_config: CVEHttpOutputConfig) -> str | None: - # Ensure domain has https:// scheme (Cognito always uses HTTPS) domain = http_config.cognito_domain.rstrip("/") if not domain.startswith("https://"): domain = f"https://{domain}" From 047f230d92b48e0a041e03b3aca41417fc323a73 Mon Sep 17 00:00:00 2001 From: Tamar Weisskopf Date: Wed, 9 Sep 2026 13:22:04 +0300 Subject: [PATCH 07/14] test: clarify Cognito field defaults test name Rename test to reflect that fields default to None (not specific to when auth_type is disabled). The same pattern applies to all auth types - fields are optional until that specific auth_type is used. --- src/vuln_analysis/functions/tests/test_cognito_auth.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/vuln_analysis/functions/tests/test_cognito_auth.py b/src/vuln_analysis/functions/tests/test_cognito_auth.py index edfbbcf19..e5de35f60 100644 --- a/src/vuln_analysis/functions/tests/test_cognito_auth.py +++ b/src/vuln_analysis/functions/tests/test_cognito_auth.py @@ -251,12 +251,12 @@ def test_returns_none_when_required_config_missing(self, cognito_config, missing class TestCognitoConfigFields: """Tests for Cognito-related configuration fields.""" - def test_cognito_fields_are_optional_with_defaults(self): - """Cognito fields should default to None when auth_type is not 'cognito'.""" + def test_cognito_fields_have_none_defaults(self): + """Cognito fields should default to None (optional until auth_type='cognito').""" config = CVEHttpOutputConfig( url="https://api.example.com", endpoint="/api/v1/reports", - auth_type="disabled", # NOT cognito + auth_type="disabled", mlops_config=MLOpsConfig(mlops_url="https://mlops.example.com"), ) From bbcc1c8509665ce44c345d6e633c1e2083b8d29d Mon Sep 17 00:00:00 2001 From: Tamar Weisskopf Date: Wed, 16 Sep 2026 01:01:40 +0300 Subject: [PATCH 08/14] fix: graceful degradation for missing cve_http_output config Add try/except guards around builder.get_function_config(HTTP_OUTPUT_AGENT_CONFIG) in clone_and_deps, generate_vdbs, and cve_segmentation stages. When cve_http_output is not registered or config is missing, falls back to default SA token authentication instead of crashing. This prevents the three core pipeline stages from breaking when HTTP output config is absent. Addresses review feedback from tmihalac on PR #331. --- src/vuln_analysis/functions/cve_clone_and_deps.py | 8 ++++++-- src/vuln_analysis/functions/cve_generate_vdbs.py | 8 ++++++-- src/vuln_analysis/functions/cve_segmentation.py | 8 ++++++-- 3 files changed, 18 insertions(+), 6 deletions(-) diff --git a/src/vuln_analysis/functions/cve_clone_and_deps.py b/src/vuln_analysis/functions/cve_clone_and_deps.py index f9cbf954d..4f396e45b 100644 --- a/src/vuln_analysis/functions/cve_clone_and_deps.py +++ b/src/vuln_analysis/functions/cve_clone_and_deps.py @@ -103,8 +103,12 @@ async def _arun(message: ExploitIqInput) -> ExploitIqEngineInput: message.scan.id, ) - http_output_config = builder.get_function_config(HTTP_OUTPUT_AGENT_CONFIG) - auth_header = get_auth_header(http_output_config) + try: + http_output_config = builder.get_function_config(HTTP_OUTPUT_AGENT_CONFIG) + auth_header = get_auth_header(http_output_config) + except Exception: + # cve_http_output not registered or config missing - use default SA token + auth_header = None 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( diff --git a/src/vuln_analysis/functions/cve_generate_vdbs.py b/src/vuln_analysis/functions/cve_generate_vdbs.py index fb858bf36..0da5caee8 100644 --- a/src/vuln_analysis/functions/cve_generate_vdbs.py +++ b/src/vuln_analysis/functions/cve_generate_vdbs.py @@ -222,8 +222,12 @@ 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) - http_output_config = builder.get_function_config(HTTP_OUTPUT_AGENT_CONFIG) - auth_header = get_auth_header(http_output_config) + try: + http_output_config = builder.get_function_config(HTTP_OUTPUT_AGENT_CONFIG) + auth_header = get_auth_header(http_output_config) + except Exception: + # cve_http_output not registered or config missing - use default SA token + auth_header = None 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 diff --git a/src/vuln_analysis/functions/cve_segmentation.py b/src/vuln_analysis/functions/cve_segmentation.py index 2bca05fb4..70bf1a9a5 100644 --- a/src/vuln_analysis/functions/cve_segmentation.py +++ b/src/vuln_analysis/functions/cve_segmentation.py @@ -225,8 +225,12 @@ async def _arun(state: ExploitIqEngineInput) -> ExploitIqEngineInput: message.scan.id, ) - http_output_config = builder.get_function_config(HTTP_OUTPUT_AGENT_CONFIG) - auth_header = get_auth_header(http_output_config) + try: + http_output_config = builder.get_function_config(HTTP_OUTPUT_AGENT_CONFIG) + auth_header = get_auth_header(http_output_config) + except Exception: + # cve_http_output not registered or config missing - use default SA token + auth_header = None 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, From 208f75d6c06d8e9aa1c3d9f5a934e026717387c7 Mon Sep 17 00:00:00 2001 From: Tamar Weisskopf Date: Wed, 16 Sep 2026 01:36:05 +0300 Subject: [PATCH 09/14] fix: improve exception handling and add logging - Replace bare 'except Exception' with specific exceptions (KeyError, ValueError, AttributeError) - Add logging when falling back to SA token authentication - Provides visibility into when fallback occurs in production --- src/vuln_analysis/functions/cve_clone_and_deps.py | 7 +++++-- src/vuln_analysis/functions/cve_generate_vdbs.py | 7 +++++-- src/vuln_analysis/functions/cve_segmentation.py | 7 +++++-- 3 files changed, 15 insertions(+), 6 deletions(-) diff --git a/src/vuln_analysis/functions/cve_clone_and_deps.py b/src/vuln_analysis/functions/cve_clone_and_deps.py index 4f396e45b..77d40b893 100644 --- a/src/vuln_analysis/functions/cve_clone_and_deps.py +++ b/src/vuln_analysis/functions/cve_clone_and_deps.py @@ -106,8 +106,11 @@ async def _arun(message: ExploitIqInput) -> ExploitIqEngineInput: try: http_output_config = builder.get_function_config(HTTP_OUTPUT_AGENT_CONFIG) auth_header = get_auth_header(http_output_config) - except Exception: - # cve_http_output not registered or config missing - use default SA token + except (KeyError, ValueError, AttributeError) as e: + logger.info( + "HTTP output config unavailable, falling back to SA token authentication: %s", + e + ) auth_header = None with http_auth_header_context(auth_header), credential_context(message.credential_id): # Configure RPM manager for IMAGE analysis diff --git a/src/vuln_analysis/functions/cve_generate_vdbs.py b/src/vuln_analysis/functions/cve_generate_vdbs.py index 0da5caee8..667a3d482 100644 --- a/src/vuln_analysis/functions/cve_generate_vdbs.py +++ b/src/vuln_analysis/functions/cve_generate_vdbs.py @@ -225,8 +225,11 @@ async def _arun(message: ExploitIqInput) -> ExploitIqEngineInput: try: http_output_config = builder.get_function_config(HTTP_OUTPUT_AGENT_CONFIG) auth_header = get_auth_header(http_output_config) - except Exception: - # cve_http_output not registered or config missing - use default SA token + except (KeyError, ValueError, AttributeError) as e: + logger.info( + "HTTP output config unavailable, falling back to SA token authentication: %s", + e + ) auth_header = None with http_auth_header_context(auth_header), credential_context(message.credential_id): logger.debug("_arun: credential_context entered, credential_id=%r", message.credential_id) diff --git a/src/vuln_analysis/functions/cve_segmentation.py b/src/vuln_analysis/functions/cve_segmentation.py index 70bf1a9a5..1e0d1db42 100644 --- a/src/vuln_analysis/functions/cve_segmentation.py +++ b/src/vuln_analysis/functions/cve_segmentation.py @@ -228,8 +228,11 @@ async def _arun(state: ExploitIqEngineInput) -> ExploitIqEngineInput: try: http_output_config = builder.get_function_config(HTTP_OUTPUT_AGENT_CONFIG) auth_header = get_auth_header(http_output_config) - except Exception: - # cve_http_output not registered or config missing - use default SA token + except (KeyError, ValueError, AttributeError) as e: + logger.info( + "HTTP output config unavailable, falling back to SA token authentication: %s", + e + ) auth_header = None with http_auth_header_context(auth_header), credential_context(message.credential_id): vdb_code_path, vdb_doc_path = await asyncio.to_thread( From 0142a627de375e6625f60ea41d5fef6672b7a51f Mon Sep 17 00:00:00 2001 From: Tamar Weisskopf Date: Wed, 16 Sep 2026 13:40:52 +0300 Subject: [PATCH 10/14] refactor: extract auth resolution to helper and improve error handling Address review feedback from tmihalac: 1. Extract resolve_http_auth_header() helper to eliminate code duplication across three pipeline stages (clone_and_deps, generate_vdbs, segmentation) 2. Narrow exception handling to KeyError only (not ValueError/AttributeError) to avoid swallowing programming bugs 3. Separate get_function_config error handling from get_auth_header call to ensure accurate error messages 4. Add warning when configured auth type fails to produce token (e.g. Cognito failure), distinguishing it from intentionally disabled auth 5. Use logger.warning (not info) for auth downgrade from configured identity to SA token - this is a trust-boundary change that needs visibility Addresses review comments: - Important: Code duplication (discussion_r4024823753) - Important: Exception too broad and logs too quietly (discussion_r4024823757) --- .../functions/cve_clone_and_deps.py | 12 ++------- .../functions/cve_generate_vdbs.py | 12 ++------- .../functions/cve_http_output.py | 25 +++++++++++++++++++ .../functions/cve_segmentation.py | 12 ++------- 4 files changed, 31 insertions(+), 30 deletions(-) diff --git a/src/vuln_analysis/functions/cve_clone_and_deps.py b/src/vuln_analysis/functions/cve_clone_and_deps.py index 77d40b893..d7ace8a1c 100644 --- a/src/vuln_analysis/functions/cve_clone_and_deps.py +++ b/src/vuln_analysis/functions/cve_clone_and_deps.py @@ -32,7 +32,7 @@ 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, http_auth_header_context -from vuln_analysis.functions.cve_http_output import get_auth_header, HTTP_OUTPUT_AGENT_CONFIG +from vuln_analysis.functions.cve_http_output import resolve_http_auth_header from exploit_iq_commons.utils.dep_tree import detect_ecosystem from exploit_iq_commons.utils.git_utils import resolve_path_to_manifest @@ -103,15 +103,7 @@ async def _arun(message: ExploitIqInput) -> ExploitIqEngineInput: message.scan.id, ) - try: - http_output_config = builder.get_function_config(HTTP_OUTPUT_AGENT_CONFIG) - auth_header = get_auth_header(http_output_config) - except (KeyError, ValueError, AttributeError) as e: - logger.info( - "HTTP output config unavailable, falling back to SA token authentication: %s", - e - ) - auth_header = None + auth_header = resolve_http_auth_header(builder) 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( diff --git a/src/vuln_analysis/functions/cve_generate_vdbs.py b/src/vuln_analysis/functions/cve_generate_vdbs.py index 667a3d482..111accb69 100644 --- a/src/vuln_analysis/functions/cve_generate_vdbs.py +++ b/src/vuln_analysis/functions/cve_generate_vdbs.py @@ -30,7 +30,7 @@ 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, http_auth_header_context -from vuln_analysis.functions.cve_http_output import get_auth_header, HTTP_OUTPUT_AGENT_CONFIG +from vuln_analysis.functions.cve_http_output import resolve_http_auth_header 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 @@ -222,15 +222,7 @@ 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) - try: - http_output_config = builder.get_function_config(HTTP_OUTPUT_AGENT_CONFIG) - auth_header = get_auth_header(http_output_config) - except (KeyError, ValueError, AttributeError) as e: - logger.info( - "HTTP output config unavailable, falling back to SA token authentication: %s", - e - ) - auth_header = None + auth_header = resolve_http_auth_header(builder) 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 diff --git a/src/vuln_analysis/functions/cve_http_output.py b/src/vuln_analysis/functions/cve_http_output.py index 62448a1f4..4fa606092 100644 --- a/src/vuln_analysis/functions/cve_http_output.py +++ b/src/vuln_analysis/functions/cve_http_output.py @@ -312,6 +312,31 @@ def _fetch_cognito_token(http_config: CVEHttpOutputConfig) -> str | None: return None +def resolve_http_auth_header(builder) -> str | None: + """ + Resolve HTTP authentication header from cve_http_output config. + + Returns None if config is missing (falls back to SA token) or if auth is disabled. + Logs warnings when configured auth fails to produce a token. + """ + try: + http_output_config = builder.get_function_config(HTTP_OUTPUT_AGENT_CONFIG) + except KeyError as e: + logger.info("HTTP output config not registered, using SA token authentication: %s", e) + return None + + auth_header = get_auth_header(http_output_config) + + # Warn if auth was configured but failed to produce a token + if auth_header is None and http_output_config.auth_type not in ("disabled", None): + logger.warning( + "Auth type '%s' configured but failed to produce token, falling back to SA token", + http_output_config.auth_type + ) + + return auth_header + + def get_auth_header(http_config: CVEHttpOutputConfig | None) -> str | None: if http_config is None: return None diff --git a/src/vuln_analysis/functions/cve_segmentation.py b/src/vuln_analysis/functions/cve_segmentation.py index 1e0d1db42..88fcb6847 100644 --- a/src/vuln_analysis/functions/cve_segmentation.py +++ b/src/vuln_analysis/functions/cve_segmentation.py @@ -37,7 +37,7 @@ from exploit_iq_commons.logging.loggers_factory import LoggingFactory, trace_id 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.functions.cve_http_output import resolve_http_auth_header from vuln_analysis.tools.tool_names import ToolNames logger = LoggingFactory.get_agent_logger(__name__) @@ -225,15 +225,7 @@ async def _arun(state: ExploitIqEngineInput) -> ExploitIqEngineInput: message.scan.id, ) - try: - http_output_config = builder.get_function_config(HTTP_OUTPUT_AGENT_CONFIG) - auth_header = get_auth_header(http_output_config) - except (KeyError, ValueError, AttributeError) as e: - logger.info( - "HTTP output config unavailable, falling back to SA token authentication: %s", - e - ) - auth_header = None + auth_header = resolve_http_auth_header(builder) 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, From 3cffe0a84216c03c90b927c36f9aa312d7afa7e9 Mon Sep 17 00:00:00 2001 From: Tamar Weisskopf Date: Wed, 16 Sep 2026 13:49:46 +0300 Subject: [PATCH 11/14] test: add integration tests for http_auth_header_context Address review feedback (discussion_r4024823747): Tests for http_auth_header_context manager: - Uses context header when set (SA token not called) - Falls back to SA token when context is None - ContextVar resets properly after context exit (concurrency-safe) - Nested contexts use innermost value and restore correctly - Supports both Bearer and Basic auth header formats Tests for resolve_http_auth_header helper: - Returns None when config missing (KeyError handled) - Returns None when auth disabled (no warnings) - Returns Bearer token on Cognito success - Logs WARNING when configured auth fails (not INFO) - Handles basic/keycloak auth types Coverage: - http_auth_header_context: 6 tests (context behavior, reset, nesting) - resolve_http_auth_header: 9 tests (config missing, auth types, warnings) Addresses: Critical review item about untested integration --- .../tests/test_resolve_http_auth_header.py | 188 ++++++++++++++++++ .../tools/tests/test_credential_client.py | 152 ++++++++++++++ 2 files changed, 340 insertions(+) create mode 100644 src/vuln_analysis/functions/tests/test_resolve_http_auth_header.py diff --git a/src/vuln_analysis/functions/tests/test_resolve_http_auth_header.py b/src/vuln_analysis/functions/tests/test_resolve_http_auth_header.py new file mode 100644 index 000000000..e643811cf --- /dev/null +++ b/src/vuln_analysis/functions/tests/test_resolve_http_auth_header.py @@ -0,0 +1,188 @@ +# 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. + +"""Integration tests for resolve_http_auth_header() helper function.""" + +from unittest.mock import Mock, patch +import pytest + +from vuln_analysis.functions.cve_http_output import ( + resolve_http_auth_header, + CVEHttpOutputConfig, + MLOpsConfig, +) + + +class TestResolveHttpAuthHeader: + """Tests for resolve_http_auth_header() helper used by pipeline stages.""" + + def test_returns_none_when_config_missing(self, caplog): + """When get_function_config raises KeyError, should log info and return None.""" + builder = Mock() + builder.get_function_config.side_effect = KeyError("cve_http_output not found") + + result = resolve_http_auth_header(builder) + + assert result is None + assert "HTTP output config not registered" in caplog.text + assert "using SA token authentication" in caplog.text + + def test_returns_none_when_auth_disabled(self): + """When auth_type is 'disabled', should return None without warnings.""" + builder = Mock() + config = CVEHttpOutputConfig( + url="https://api.example.com", + endpoint="/api/v1/reports", + auth_type="disabled", + mlops_config=MLOpsConfig(mlops_url="https://mlops.example.com"), + ) + builder.get_function_config.return_value = config + + result = resolve_http_auth_header(builder) + + assert result is None + + def test_returns_none_when_auth_type_none(self): + """When auth_type is None, should return None without warnings.""" + builder = Mock() + config = CVEHttpOutputConfig( + url="https://api.example.com", + endpoint="/api/v1/reports", + auth_type=None, + mlops_config=MLOpsConfig(mlops_url="https://mlops.example.com"), + ) + builder.get_function_config.return_value = config + + result = resolve_http_auth_header(builder) + + assert result is None + + @patch("vuln_analysis.functions.cve_http_output._fetch_cognito_token") + def test_returns_bearer_token_on_cognito_success(self, mock_fetch): + """When Cognito auth succeeds, should return Bearer token.""" + builder = Mock() + config = 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"), + ) + builder.get_function_config.return_value = config + mock_fetch.return_value = "cognito-access-token-xyz" + + result = resolve_http_auth_header(builder) + + assert result == "Bearer cognito-access-token-xyz" + + @patch("vuln_analysis.functions.cve_http_output._fetch_cognito_token") + def test_warns_when_cognito_fails(self, mock_fetch, caplog): + """When Cognito token fetch fails, should log warning and return None.""" + builder = Mock() + config = 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"), + ) + builder.get_function_config.return_value = config + mock_fetch.return_value = None # Simulate failure + + result = resolve_http_auth_header(builder) + + assert result is None + assert "Auth type 'cognito' configured but failed to produce token" in caplog.text + assert "falling back to SA token" in caplog.text + + def test_warns_when_basic_auth_missing_credentials(self, caplog): + """When basic auth is configured but credentials missing, should warn.""" + builder = Mock() + config = CVEHttpOutputConfig( + url="https://api.example.com", + endpoint="/api/v1/reports", + auth_type="basic", + username=None, # Missing + password=None, # Missing + mlops_config=MLOpsConfig(mlops_url="https://mlops.example.com"), + ) + builder.get_function_config.return_value = config + + result = resolve_http_auth_header(builder) + + assert result is None + assert "Auth type 'basic' configured but failed to produce token" in caplog.text + + def test_returns_basic_auth_header_on_success(self): + """When basic auth credentials are provided, should return Basic header.""" + builder = Mock() + config = CVEHttpOutputConfig( + url="https://api.example.com", + endpoint="/api/v1/reports", + auth_type="basic", + username="admin", + password="secret", + mlops_config=MLOpsConfig(mlops_url="https://mlops.example.com"), + ) + builder.get_function_config.return_value = config + + result = resolve_http_auth_header(builder) + + assert result is not None + assert result.startswith("Basic ") + + def test_no_warning_when_disabled_returns_none(self, caplog): + """When auth_type='disabled' returns None, should not warn (this is expected).""" + builder = Mock() + config = CVEHttpOutputConfig( + url="https://api.example.com", + endpoint="/api/v1/reports", + auth_type="disabled", + mlops_config=MLOpsConfig(mlops_url="https://mlops.example.com"), + ) + builder.get_function_config.return_value = config + + with caplog.at_level("WARNING"): + result = resolve_http_auth_header(builder) + + assert result is None + # Should NOT have any warnings + assert len(caplog.records) == 0 + + @patch("vuln_analysis.functions.cve_http_output._fetch_keycloak_token") + def test_warns_when_keycloak_fails(self, mock_fetch, caplog): + """When Keycloak token fetch fails, should log warning.""" + builder = Mock() + config = CVEHttpOutputConfig( + url="https://api.example.com", + endpoint="/api/v1/reports", + auth_type="keycloak", + keycloak_server="https://keycloak.example.com", + keycloak_realm="myrealm", + client_id="client123", + client_secret="secret456", + mlops_config=MLOpsConfig(mlops_url="https://mlops.example.com"), + ) + builder.get_function_config.return_value = config + mock_fetch.return_value = None # Simulate failure + + result = resolve_http_auth_header(builder) + + assert result is None + assert "Auth type 'keycloak' configured but failed to produce token" in caplog.text diff --git a/src/vuln_analysis/tools/tests/test_credential_client.py b/src/vuln_analysis/tools/tests/test_credential_client.py index d088e6749..cb888b4a3 100644 --- a/src/vuln_analysis/tools/tests/test_credential_client.py +++ b/src/vuln_analysis/tools/tests/test_credential_client.py @@ -10,6 +10,7 @@ CredentialNotFoundError, DecryptionError, fetch_and_decrypt_credential, + http_auth_header_context, ) # --------------------------------------------------------------------------- @@ -247,3 +248,154 @@ def test_secret_not_logged(self, mock_get, mock_ca, caplog): assert secret not in record.getMessage(), ( f"Secret value leaked into log message: {record.getMessage()}" ) + + +# --------------------------------------------------------------------------- +# Tests: http_auth_header_context integration +# --------------------------------------------------------------------------- + +class TestHttpAuthHeaderContext: + """Tests for http_auth_header_context manager and integration with fetch_and_decrypt_credential.""" + + @patch("exploit_iq_commons.utils.credential_client._validate_ca_bundle", return_value=False) + @patch("exploit_iq_commons.utils.credential_client.requests.get") + def test_uses_context_header_when_set(self, mock_get, mock_ca): + """When auth header is set via context, it should be used instead of SA token.""" + mock_get.return_value = _mock_http_ok(_make_response("token")) + + with http_auth_header_context("Bearer cognito-token-xyz"): + fetch_and_decrypt_credential( + credential_id="cred-uuid-ctx", + jwt_token="sa-token-should-not-be-used", + backend_url="https://backend.example.com", + encryption_key=_ENCRYPTION_KEY, + ) + + # Verify the context header was used, NOT the SA token + mock_get.assert_called_once_with( + "https://backend.example.com/api/v1/credentials/cred-uuid-ctx", + headers={"Authorization": "Bearer cognito-token-xyz"}, + timeout=10, + verify=False, + ) + + @patch("exploit_iq_commons.utils.credential_client._validate_ca_bundle", return_value=False) + @patch("exploit_iq_commons.utils.credential_client.requests.get") + @patch("exploit_iq_commons.utils.credential_client._resolve_jwt_token") + def test_resolve_jwt_not_called_when_context_set(self, mock_resolve, mock_get, mock_ca): + """_resolve_jwt_token should not be called when auth header is provided via context.""" + mock_get.return_value = _mock_http_ok(_make_response("token")) + mock_resolve.return_value = "sa-token" + + with http_auth_header_context("Bearer cognito-token"): + fetch_and_decrypt_credential( + credential_id="cred-uuid", + jwt_token=None, + backend_url="https://backend.example.com", + encryption_key=_ENCRYPTION_KEY, + ) + + # _resolve_jwt_token should NOT have been called + mock_resolve.assert_not_called() + + @patch("exploit_iq_commons.utils.credential_client._validate_ca_bundle", return_value=False) + @patch("exploit_iq_commons.utils.credential_client.requests.get") + @patch("exploit_iq_commons.utils.credential_client._resolve_jwt_token") + def test_falls_back_to_jwt_when_context_none(self, mock_resolve, mock_get, mock_ca): + """When context is None, should fall back to SA token via _resolve_jwt_token.""" + mock_get.return_value = _mock_http_ok(_make_response("token")) + mock_resolve.return_value = "sa-token-123" + + with http_auth_header_context(None): + fetch_and_decrypt_credential( + credential_id="cred-uuid", + jwt_token=None, + backend_url="https://backend.example.com", + encryption_key=_ENCRYPTION_KEY, + ) + + # _resolve_jwt_token SHOULD have been called + mock_resolve.assert_called_once_with(None) + + # SA token should be used + mock_get.assert_called_once_with( + "https://backend.example.com/api/v1/credentials/cred-uuid", + headers={"Authorization": "Bearer sa-token-123"}, + timeout=10, + verify=False, + ) + + @patch("exploit_iq_commons.utils.credential_client._validate_ca_bundle", return_value=False) + @patch("exploit_iq_commons.utils.credential_client.requests.get") + def test_context_reset_after_exit(self, mock_get, mock_ca): + """ContextVar should be reset after exiting the context manager.""" + mock_get.return_value = _mock_http_ok(_make_response("token")) + + # Set context + with http_auth_header_context("Bearer temp-token"): + fetch_and_decrypt_credential( + credential_id="cred-1", + jwt_token="jwt", + backend_url="https://backend.example.com", + encryption_key=_ENCRYPTION_KEY, + ) + # Inside context: temp-token is used + assert mock_get.call_args[1]["headers"]["Authorization"] == "Bearer temp-token" + + # After context exit: should fall back to JWT + mock_get.reset_mock() + fetch_and_decrypt_credential( + credential_id="cred-2", + jwt_token="jwt-after-exit", + backend_url="https://backend.example.com", + encryption_key=_ENCRYPTION_KEY, + ) + # Should use the JWT token, not the temp token + assert mock_get.call_args[1]["headers"]["Authorization"] == "Bearer jwt-after-exit" + + @patch("exploit_iq_commons.utils.credential_client._validate_ca_bundle", return_value=False) + @patch("exploit_iq_commons.utils.credential_client.requests.get") + def test_nested_contexts_use_innermost(self, mock_get, mock_ca): + """Nested contexts should use the innermost value.""" + mock_get.return_value = _mock_http_ok(_make_response("token")) + + with http_auth_header_context("Bearer outer-token"): + with http_auth_header_context("Bearer inner-token"): + fetch_and_decrypt_credential( + credential_id="cred-nested", + jwt_token="jwt", + backend_url="https://backend.example.com", + encryption_key=_ENCRYPTION_KEY, + ) + # Should use inner token + assert mock_get.call_args[1]["headers"]["Authorization"] == "Bearer inner-token" + + # After inner exit, should revert to outer + mock_get.reset_mock() + fetch_and_decrypt_credential( + credential_id="cred-outer", + jwt_token="jwt", + backend_url="https://backend.example.com", + encryption_key=_ENCRYPTION_KEY, + ) + assert mock_get.call_args[1]["headers"]["Authorization"] == "Bearer outer-token" + + @patch("exploit_iq_commons.utils.credential_client._validate_ca_bundle", return_value=False) + @patch("exploit_iq_commons.utils.credential_client.requests.get") + def test_basic_auth_header_format(self, mock_get, mock_ca): + """Context should support Basic auth format (not just Bearer).""" + import base64 + creds = base64.b64encode(b"user:pass").decode() + basic_header = f"Basic {creds}" + + mock_get.return_value = _mock_http_ok(_make_response("token")) + + with http_auth_header_context(basic_header): + fetch_and_decrypt_credential( + credential_id="cred-basic", + jwt_token="jwt", + backend_url="https://backend.example.com", + encryption_key=_ENCRYPTION_KEY, + ) + + assert mock_get.call_args[1]["headers"]["Authorization"] == basic_header From 30108bbefa0fd5e6cd392d386c16b4111bddbd88 Mon Sep 17 00:00:00 2001 From: Tamar Weisskopf Date: Thu, 17 Sep 2026 00:28:32 +0300 Subject: [PATCH 12/14] fix: raise AuthHeaderError instead of silent None on auth failure Replace silent None returns in get_auth_header() with AuthHeaderError exceptions to prevent configured auth failures from silently downgrading to unauthenticated requests. Callers (output_to_http, resolve_http_auth_header) catch the exception and handle appropriately: abort sending or fall back to SA token with a warning. Co-Authored-By: Claude Sonnet 4.5 --- .../functions/cve_http_output.py | 54 +++++++++++------- .../functions/tests/test_cognito_auth.py | 37 ++++++------ .../tests/test_resolve_http_auth_header.py | 57 +++++++++---------- 3 files changed, 76 insertions(+), 72 deletions(-) diff --git a/src/vuln_analysis/functions/cve_http_output.py b/src/vuln_analysis/functions/cve_http_output.py index 4fa606092..d6c1539c8 100644 --- a/src/vuln_analysis/functions/cve_http_output.py +++ b/src/vuln_analysis/functions/cve_http_output.py @@ -191,7 +191,12 @@ async def _arun(message: ExploitIqOutput) -> ExploitIqOutput: #logger.info(f"Saved JSON output to {json_file}") headers = {'Content-type': 'application/json', 'traceId': trace_id.get()} - auth_header = get_auth_header(config) + try: + auth_header = get_auth_header(config) + except AuthHeaderError as e: + logger.error("Authentication failed, cannot send output: %s", e) + return message + if auth_header is not None: headers['Authorization'] = auth_header verify = config.verify_path if config.verify_path else True @@ -325,19 +330,27 @@ def resolve_http_auth_header(builder) -> str | None: logger.info("HTTP output config not registered, using SA token authentication: %s", e) return None - auth_header = get_auth_header(http_output_config) - - # Warn if auth was configured but failed to produce a token - if auth_header is None and http_output_config.auth_type not in ("disabled", None): + try: + return get_auth_header(http_output_config) + except AuthHeaderError as e: logger.warning( - "Auth type '%s' configured but failed to produce token, falling back to SA token", - http_output_config.auth_type + "Auth type '%s' failed, falling back to SA token: %s", + http_output_config.auth_type, + e ) + return None + - return auth_header +class AuthHeaderError(RuntimeError): + """Raised when a configured auth type fails to produce a valid header.""" def get_auth_header(http_config: CVEHttpOutputConfig | None) -> str | None: + """Return an Authorization header value, or None when auth is disabled. + + Raises AuthHeaderError when auth is configured but cannot produce a + valid header (missing credentials, token fetch failure, etc.). + """ if http_config is None: return None @@ -347,16 +360,15 @@ def get_auth_header(http_config: CVEHttpOutputConfig | None) -> str | None: credentials = f"{http_config.username}:{http_config.password}" encoded_creds = base64.b64encode(credentials.encode('utf-8')).decode('utf-8') return f"Basic {encoded_creds}" - return None + raise AuthHeaderError("Basic auth requires username and password") case "keycloak": if not all([http_config.keycloak_server, http_config.keycloak_realm, http_config.client_id, http_config.client_secret]): - logger.error("Keycloak auth requires keycloak_server, keycloak_realm, client_id, client_secret") - return None + raise AuthHeaderError("Keycloak auth requires keycloak_server, keycloak_realm, client_id, client_secret") token = _fetch_keycloak_token(http_config) if token: return f"Bearer {token}" - return None + raise AuthHeaderError("Failed to fetch Keycloak access token") case "bearer": if http_config.token: return f"Bearer {http_config.token}" @@ -365,22 +377,22 @@ def get_auth_header(http_config: CVEHttpOutputConfig | None) -> str | None: with open(http_config.token_path, 'r') as file: return f"Bearer {file.read().strip()}" except Exception as e: - logger.warn(f"Unable to read OAuth token: {e}") - return None + raise AuthHeaderError(f"Unable to read OAuth token from {http_config.token_path}: {e}") from e + raise AuthHeaderError("Bearer auth requires token or token_path") 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 + raise AuthHeaderError("Cognito auth requires cognito_domain, cognito_client_id, and cognito_client_secret") token = _fetch_cognito_token(http_config) - return f"Bearer {token}" if token else None + if token: + return f"Bearer {token}" + raise AuthHeaderError("Failed to fetch Cognito access token") case "disabled" | None: return None case _: - logger.error( - "Unrecognized auth_type '%s'. Valid values: basic, bearer, keycloak, cognito, disabled", - http_config.auth_type, + raise AuthHeaderError( + f"Unrecognized auth_type '{http_config.auth_type}'. " + "Valid values: basic, bearer, keycloak, cognito, disabled" ) - return None def _http_params_override(http_config: CVEHttpOutputConfig, mlops_config: MLOpsConfig,http_headers) -> dict[str, Any]: diff --git a/src/vuln_analysis/functions/tests/test_cognito_auth.py b/src/vuln_analysis/functions/tests/test_cognito_auth.py index e5de35f60..10807fac2 100644 --- a/src/vuln_analysis/functions/tests/test_cognito_auth.py +++ b/src/vuln_analysis/functions/tests/test_cognito_auth.py @@ -21,6 +21,7 @@ import requests from vuln_analysis.functions.cve_http_output import ( + AuthHeaderError, CVEHttpOutputConfig, MLOpsConfig, _fetch_cognito_token, @@ -206,12 +207,11 @@ def test_returns_bearer_token_on_success(self, 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.""" + def test_raises_when_token_fetch_fails(self, cognito_config): + """Should raise AuthHeaderError 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 + with pytest.raises(AuthHeaderError, match="Failed to fetch Cognito access token"): + get_auth_header(cognito_config) def test_returns_none_for_disabled_auth_type(self, cognito_config): """Should return None when auth_type is 'disabled'.""" @@ -219,16 +219,12 @@ def test_returns_none_for_disabled_auth_type(self, cognito_config): header = get_auth_header(cognito_config) assert header is None - def test_logs_error_for_unrecognized_auth_type(self, cognito_config): - """Should log error and return None for typos/invalid auth_type.""" + def test_raises_for_unrecognized_auth_type(self, cognito_config): + """Should raise AuthHeaderError for typos/invalid auth_type.""" cognito_config.auth_type = "Cognito" # Capital C typo - with patch('vuln_analysis.functions.cve_http_output.logger') as mock_logger: - header = get_auth_header(cognito_config) - - assert header is None - mock_logger.error.assert_called_once() - assert "Unrecognized auth_type" in mock_logger.error.call_args[0][0] + with pytest.raises(AuthHeaderError, match="Unrecognized auth_type"): + get_auth_header(cognito_config) @pytest.mark.parametrize("missing_field,config_override", [ ("cognito_domain", {"cognito_domain": None}), @@ -236,15 +232,15 @@ def test_logs_error_for_unrecognized_auth_type(self, cognito_config): ("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.""" + def test_raises_when_required_config_missing(self, cognito_config, missing_field, config_override): + """Should raise AuthHeaderError when required fields 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) + with pytest.raises(AuthHeaderError, match="Cognito auth requires"): + get_auth_header(cognito_config) - assert header is None mock_fetch.assert_not_called() # Should not attempt fetch with incomplete config @@ -266,7 +262,7 @@ def test_cognito_fields_have_none_defaults(self): assert config.cognito_client_secret is None def test_cognito_fields_required_when_auth_type_is_cognito(self): - """When auth_type='cognito', get_auth_header should validate required fields.""" + """When auth_type='cognito', get_auth_header should raise if required fields are missing.""" config = CVEHttpOutputConfig( url="https://api.example.com", endpoint="/api/v1/reports", @@ -275,9 +271,8 @@ def test_cognito_fields_required_when_auth_type_is_cognito(self): mlops_config=MLOpsConfig(mlops_url="https://mlops.example.com"), ) - # Runtime validation should return None when required fields are missing - header = get_auth_header(config) - assert header is None + with pytest.raises(AuthHeaderError, match="Cognito auth requires"): + get_auth_header(config) def test_auth_type_field_documents_cognito(self): """The auth_type field description should mention 'cognito' as a valid option.""" diff --git a/src/vuln_analysis/functions/tests/test_resolve_http_auth_header.py b/src/vuln_analysis/functions/tests/test_resolve_http_auth_header.py index e643811cf..bd6e024f8 100644 --- a/src/vuln_analysis/functions/tests/test_resolve_http_auth_header.py +++ b/src/vuln_analysis/functions/tests/test_resolve_http_auth_header.py @@ -28,7 +28,8 @@ class TestResolveHttpAuthHeader: """Tests for resolve_http_auth_header() helper used by pipeline stages.""" - def test_returns_none_when_config_missing(self, caplog): + @patch("vuln_analysis.functions.cve_http_output.logger") + def test_returns_none_when_config_missing(self, mock_logger): """When get_function_config raises KeyError, should log info and return None.""" builder = Mock() builder.get_function_config.side_effect = KeyError("cve_http_output not found") @@ -36,8 +37,9 @@ def test_returns_none_when_config_missing(self, caplog): result = resolve_http_auth_header(builder) assert result is None - assert "HTTP output config not registered" in caplog.text - assert "using SA token authentication" in caplog.text + mock_logger.info.assert_called_once() + log_msg = mock_logger.info.call_args[0][0] + assert "HTTP output config not registered" in log_msg def test_returns_none_when_auth_disabled(self): """When auth_type is 'disabled', should return None without warnings.""" @@ -54,21 +56,6 @@ def test_returns_none_when_auth_disabled(self): assert result is None - def test_returns_none_when_auth_type_none(self): - """When auth_type is None, should return None without warnings.""" - builder = Mock() - config = CVEHttpOutputConfig( - url="https://api.example.com", - endpoint="/api/v1/reports", - auth_type=None, - mlops_config=MLOpsConfig(mlops_url="https://mlops.example.com"), - ) - builder.get_function_config.return_value = config - - result = resolve_http_auth_header(builder) - - assert result is None - @patch("vuln_analysis.functions.cve_http_output._fetch_cognito_token") def test_returns_bearer_token_on_cognito_success(self, mock_fetch): """When Cognito auth succeeds, should return Bearer token.""" @@ -89,8 +76,9 @@ def test_returns_bearer_token_on_cognito_success(self, mock_fetch): assert result == "Bearer cognito-access-token-xyz" + @patch("vuln_analysis.functions.cve_http_output.logger") @patch("vuln_analysis.functions.cve_http_output._fetch_cognito_token") - def test_warns_when_cognito_fails(self, mock_fetch, caplog): + def test_warns_when_cognito_fails(self, mock_fetch, mock_logger): """When Cognito token fetch fails, should log warning and return None.""" builder = Mock() config = CVEHttpOutputConfig( @@ -108,10 +96,13 @@ def test_warns_when_cognito_fails(self, mock_fetch, caplog): result = resolve_http_auth_header(builder) assert result is None - assert "Auth type 'cognito' configured but failed to produce token" in caplog.text - assert "falling back to SA token" in caplog.text + mock_logger.warning.assert_called_once() + log_msg = mock_logger.warning.call_args[0][0] % mock_logger.warning.call_args[0][1:] + assert "cognito" in log_msg + assert "falling back to SA token" in log_msg - def test_warns_when_basic_auth_missing_credentials(self, caplog): + @patch("vuln_analysis.functions.cve_http_output.logger") + def test_warns_when_basic_auth_missing_credentials(self, mock_logger): """When basic auth is configured but credentials missing, should warn.""" builder = Mock() config = CVEHttpOutputConfig( @@ -127,7 +118,10 @@ def test_warns_when_basic_auth_missing_credentials(self, caplog): result = resolve_http_auth_header(builder) assert result is None - assert "Auth type 'basic' configured but failed to produce token" in caplog.text + mock_logger.warning.assert_called_once() + log_msg = mock_logger.warning.call_args[0][0] % mock_logger.warning.call_args[0][1:] + assert "basic" in log_msg + assert "falling back to SA token" in log_msg def test_returns_basic_auth_header_on_success(self): """When basic auth credentials are provided, should return Basic header.""" @@ -147,7 +141,8 @@ def test_returns_basic_auth_header_on_success(self): assert result is not None assert result.startswith("Basic ") - def test_no_warning_when_disabled_returns_none(self, caplog): + @patch("vuln_analysis.functions.cve_http_output.logger") + def test_no_warning_when_disabled_returns_none(self, mock_logger): """When auth_type='disabled' returns None, should not warn (this is expected).""" builder = Mock() config = CVEHttpOutputConfig( @@ -158,15 +153,14 @@ def test_no_warning_when_disabled_returns_none(self, caplog): ) builder.get_function_config.return_value = config - with caplog.at_level("WARNING"): - result = resolve_http_auth_header(builder) + result = resolve_http_auth_header(builder) assert result is None - # Should NOT have any warnings - assert len(caplog.records) == 0 + mock_logger.warning.assert_not_called() + @patch("vuln_analysis.functions.cve_http_output.logger") @patch("vuln_analysis.functions.cve_http_output._fetch_keycloak_token") - def test_warns_when_keycloak_fails(self, mock_fetch, caplog): + def test_warns_when_keycloak_fails(self, mock_fetch, mock_logger): """When Keycloak token fetch fails, should log warning.""" builder = Mock() config = CVEHttpOutputConfig( @@ -185,4 +179,7 @@ def test_warns_when_keycloak_fails(self, mock_fetch, caplog): result = resolve_http_auth_header(builder) assert result is None - assert "Auth type 'keycloak' configured but failed to produce token" in caplog.text + mock_logger.warning.assert_called_once() + log_msg = mock_logger.warning.call_args[0][0] % mock_logger.warning.call_args[0][1:] + assert "keycloak" in log_msg + assert "falling back to SA token" in log_msg From 0da736a9b1b217af9b8fe12a1ee09f2e3e0e675d Mon Sep 17 00:00:00 2001 From: Tamar Weisskopf Date: Thu, 17 Sep 2026 00:45:56 +0300 Subject: [PATCH 13/14] fix: let AuthHeaderError propagate from output_to_http MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Auth failure in the delivery stage should fail the pipeline visibly, not silently return the message as if delivery succeeded. Pipeline stages (clone, vdb, segmentation) still catch and fall back to SA token via resolve_http_auth_header — that graceful degradation is correct for credential fetching, but not for the final report delivery. Co-Authored-By: Claude Sonnet 4.5 --- src/vuln_analysis/functions/cve_http_output.py | 7 +------ 1 file changed, 1 insertion(+), 6 deletions(-) diff --git a/src/vuln_analysis/functions/cve_http_output.py b/src/vuln_analysis/functions/cve_http_output.py index d6c1539c8..6d668d913 100644 --- a/src/vuln_analysis/functions/cve_http_output.py +++ b/src/vuln_analysis/functions/cve_http_output.py @@ -191,12 +191,7 @@ async def _arun(message: ExploitIqOutput) -> ExploitIqOutput: #logger.info(f"Saved JSON output to {json_file}") headers = {'Content-type': 'application/json', 'traceId': trace_id.get()} - try: - auth_header = get_auth_header(config) - except AuthHeaderError as e: - logger.error("Authentication failed, cannot send output: %s", e) - return message - + auth_header = get_auth_header(config) if auth_header is not None: headers['Authorization'] = auth_header verify = config.verify_path if config.verify_path else True From eaf22554ea16213b861a5f0ee342681d08ed9cf2 Mon Sep 17 00:00:00 2001 From: Tamar Weisskopf Date: Thu, 17 Sep 2026 12:25:53 +0300 Subject: [PATCH 14/14] fix: propagate HTTP delivery failures from output_to_http Re-raise after logging in the outer except block so the pipeline sees the failure instead of silently returning message as success. Co-Authored-By: Claude Opus 4.6 --- src/vuln_analysis/functions/cve_http_output.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/vuln_analysis/functions/cve_http_output.py b/src/vuln_analysis/functions/cve_http_output.py index 6d668d913..dc095cefb 100644 --- a/src/vuln_analysis/functions/cve_http_output.py +++ b/src/vuln_analysis/functions/cve_http_output.py @@ -226,6 +226,7 @@ async def _arun(message: ExploitIqOutput) -> ExploitIqOutput: logger.error('Unable to send job to MLOps API at %s. Error: %s', mlops_url, mlops_e) except Exception as e: logger.error('Unable to send output response to %s. Error: %s', payload.url, e) + raise else: logger.info('Successfully sent output to %s', payload.url)