feat: Add AWS Cognito as OIDC Identity Provider (#325)

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}

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)
 

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}

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

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 = (

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]:

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,