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..d7ace8a1c 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 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 + + 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,8 @@ async def _arun(message: ExploitIqInput) -> ExploitIqEngineInput: message.scan.id, ) - with credential_context(message.credential_id): + 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( 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..111accb69 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 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 @@ -221,7 +222,8 @@ 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): + 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 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..dc095cefb 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. 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") 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", @@ -220,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) @@ -277,23 +284,87 @@ def _fetch_keycloak_token(http_config: CVEHttpOutputConfig) -> str | None: return None +def _fetch_cognito_token(http_config: CVEHttpOutputConfig) -> str | None: + 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() + ).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 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 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 + + try: + return get_auth_header(http_output_config) + except AuthHeaderError as e: + logger.warning( + "Auth type '%s' failed, falling back to SA token: %s", + http_output_config.auth_type, + e + ) + return None + + +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 + match http_config.auth_type: case "basic": if http_config.username and http_config.password: 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}" @@ -302,10 +373,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 - case 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]): + raise AuthHeaderError("Cognito auth requires cognito_domain, cognito_client_id, and cognito_client_secret") + token = _fetch_cognito_token(http_config) + if token: + return f"Bearer {token}" + raise AuthHeaderError("Failed to fetch Cognito access token") + case "disabled" | None: return None + case _: + raise AuthHeaderError( + f"Unrecognized auth_type '{http_config.auth_type}'. " + "Valid values: basic, bearer, keycloak, cognito, disabled" + ) def _http_params_override(http_config: CVEHttpOutputConfig, mlops_config: MLOpsConfig,http_headers) -> dict[str, Any]: diff --git a/src/vuln_analysis/functions/cve_segmentation.py b/src/vuln_analysis/functions/cve_segmentation.py index de75f94c8..88fcb6847 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 resolve_http_auth_header from vuln_analysis.tools.tool_names import ToolNames logger = LoggingFactory.get_agent_logger(__name__) @@ -224,7 +225,8 @@ async def _arun(state: ExploitIqEngineInput) -> ExploitIqEngineInput: message.scan.id, ) - with credential_context(message.credential_id): + 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, source_infos, 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..10807fac2 --- /dev/null +++ b/src/vuln_analysis/functions/tests/test_cognito_auth.py @@ -0,0 +1,280 @@ +# 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 ( + AuthHeaderError, + 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 + + 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"): + header = get_auth_header(cognito_config) + + assert header == "Bearer test-token" + + 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): + 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'.""" + cognito_config.auth_type = "disabled" + header = get_auth_header(cognito_config) + assert header is None + + 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 pytest.raises(AuthHeaderError, match="Unrecognized auth_type"): + get_auth_header(cognito_config) + + @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_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: + with pytest.raises(AuthHeaderError, match="Cognito auth requires"): + get_auth_header(cognito_config) + + mock_fetch.assert_not_called() # Should not attempt fetch with incomplete config + + +class TestCognitoConfigFields: + """Tests for Cognito-related configuration fields.""" + + 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", + 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_required_when_auth_type_is_cognito(self): + """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", + auth_type="cognito", + # Cognito fields are None (missing) + mlops_config=MLOpsConfig(mlops_url="https://mlops.example.com"), + ) + + 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.""" + field_info = CVEHttpOutputConfig.model_fields['auth_type'] + assert "cognito" in field_info.description 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..bd6e024f8 --- /dev/null +++ b/src/vuln_analysis/functions/tests/test_resolve_http_auth_header.py @@ -0,0 +1,185 @@ +# 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.""" + + @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") + + result = resolve_http_auth_header(builder) + + assert result is None + 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.""" + 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 + + @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.logger") + @patch("vuln_analysis.functions.cve_http_output._fetch_cognito_token") + 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( + 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 + 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 + + @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( + 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 + 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.""" + 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 ") + + @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( + 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 + 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, mock_logger): + """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 + 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 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