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

@@ -30,10 +30,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__)
 
 
@@ -80,7 +83,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.
@@ -100,7 +102,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

src/vuln_analysis/functions/cve_generate_vdbs.py

@@ -28,7 +28,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
@@ -220,7 +221,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 = (

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
 

src/vuln_analysis/functions/cve_segmentation.py

@@ -35,7 +35,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 exploit_iq_commons.utils.dep_tree import Ecosystem
 from vuln_analysis.tools.tool_names import ToolNames
 
@@ -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 = embedder.build_vdbs(
                     source_infos,
                     config.ignore_code_embedding,