diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 43575ac5..37a727a1 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -42,6 +42,14 @@ jobs: run: | uv pip install --python .venv -e ".[all]" --group dev + - name: Run AWS Secrets Manager tests + run: | + .venv/bin/pytest tests/test_aws_secrets_manager.py -v + + - name: Run managed configuration and AWS profile adapter tests + run: | + .venv/bin/pytest tests/test_managed_configuration.py tests/test_aws_profile_adapter.py -v + - name: Run log handling tests run: | .venv/bin/pytest tests/test_log_handlers.py -v diff --git a/AGENTS.md b/AGENTS.md new file mode 100644 index 00000000..c03e646b --- /dev/null +++ b/AGENTS.md @@ -0,0 +1,23 @@ +# Repository guidance + +## Managed collector job queue API + +Managed collector job queue endpoints must use `/api/v2/collie-job-queue` in +public URL paths. Do not use `collector` in this endpoint path: ad blockers may +block URLs containing that term. This applies to all managed queue operations. +This URL naming requirement does not change internal terminology. Keep names +such as `CollectorJob`, collector-related Python identifiers, job keys, models, +and database concepts unchanged unless a separate requirement calls for it. + +## BloodHound API contracts + +When adding or changing a BloodHound or BloodHound Enterprise (BHE) API endpoint, +verify its path, HTTP method, parameters, request body, and response against the +applicable API contract. Refer to the [BloodHound repository](https://github.com/SpecterOps/BloodHound) +and [BloodHound Enterprise repository](https://github.com/SpecterOps/bloodhound-enterprise) +for their contracts and implementation details. If an endpoint is being +developed concurrently and is not yet in the published contract, check the +relevant repository's implementation or mock for the agreed behavior and +coordinate any contract changes with that work. Update existing OpenHound API +specifications or endpoint documentation when present; do not invent a new +specification if the repository has none. diff --git a/example-configurations/bloodhound-enterprise/.dlt-example/config.toml b/example-configurations/bloodhound-enterprise/.dlt-example/config.toml index 8da76f80..370cf514 100644 --- a/example-configurations/bloodhound-enterprise/.dlt-example/config.toml +++ b/example-configurations/bloodhound-enterprise/.dlt-example/config.toml @@ -1,4 +1,14 @@ # Example configuration: https://bloodhound.specterops.io/openhound/enterprise#full-configuration-example +[openhound] +managed = false + +# Managed deployments read BHE HMAC credentials from AWS Secrets Manager rather +# than secrets.toml. The named secret must be a JSON object with token_id and +# token_key fields. poll_interval controls the BHE scheduler cadence in seconds. +# [openhound.aws_secrets_manager] +# secret_name = "openhound/bhe-credentials" +# poll_interval = 30 + [runtime] http_show_error_body = true log_cli_level = "WARNING" @@ -10,8 +20,9 @@ request_max_attempts = 15 request_backoff_factor = 1.3 request_max_retry_delay = 900 -# Default: logs are set to human readable text -# To switch to structured JSON instead, uncomment the line below (must be uppercase "JSON") +# Self-managed mode defaults to human-readable text; managed mode defaults to +# structured JSON. To override either default, uncomment the line below (DLT +# requires the value to be uppercase "JSON"). # log_format = "JSON" # BHE destination URL (non-sensitive; tokens go in secrets.toml) diff --git a/pyproject.toml b/pyproject.toml index 064a7106..db615b6c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -19,6 +19,8 @@ dependencies = [ "jinja2>=3.1.6", "types-requests==2.33.0.20260503", "pydantic-extra-types>=2.11.1", + "boto3>=1.43.90", + "filelock>=3.32.0", ] [project.optional-dependencies] diff --git a/src/openhound/config.py b/src/openhound/config.py new file mode 100644 index 00000000..598c310c --- /dev/null +++ b/src/openhound/config.py @@ -0,0 +1,76 @@ +"""OpenHound managed-mode configuration.""" + +from dataclasses import dataclass + +import dlt + + +@dataclass(frozen=True) +class ModeDefaults: + """Configuration defaults that differ between OpenHound modes.""" + + log_format: str + + +@dataclass(frozen=True) +class ManagedBHEConfig: + """Managed-mode settings for BHE credentials and scheduler polling.""" + + secret_name: str + poll_interval: int + + +_UNMANAGED_DEFAULTS = ModeDefaults(log_format="text") +_MANAGED_DEFAULTS = ModeDefaults(log_format="json") +DEFAULT_MANAGED_BHE_POLL_INTERVAL = 30 +MAX_MANAGED_BHE_POLL_INTERVAL = 599 + + +def is_managed() -> bool: + """Return whether OpenHound is configured to run in managed mode. + + The value is resolved on every call so all platform code observes the current + DLT configuration. Missing configuration defaults to unmanaged mode. Invalid + boolean values raise DLT's configuration coercion error. + """ + + configured = dlt.config.get("openhound.managed", bool) + return configured if configured is not None else False + + +def get_mode_defaults() -> ModeDefaults: + """Return the defaults profile for the configured mode.""" + + return _MANAGED_DEFAULTS if is_managed() else _UNMANAGED_DEFAULTS + + +def get_managed_bhe_config() -> ManagedBHEConfig: + """Return AWS secret and scheduler settings used by managed BHE clients. + + ``openhound.aws_secrets_manager.secret_name`` is deliberately required in + managed mode: there is no safe default credential source. The scheduler + cadence defaults to 30 seconds, matching the historical fixed interval. + """ + + secret_name = dlt.config.get("openhound.aws_secrets_manager.secret_name", str) + if not secret_name: + raise ValueError( + "Managed mode requires " + "openhound.aws_secrets_manager.secret_name to be configured" + ) + + poll_interval = dlt.config.get( + "openhound.aws_secrets_manager.poll_interval", int + ) + if poll_interval is None: + poll_interval = DEFAULT_MANAGED_BHE_POLL_INTERVAL + if not 0 < poll_interval <= MAX_MANAGED_BHE_POLL_INTERVAL: + raise ValueError( + "openhound.aws_secrets_manager.poll_interval must be between 1 and " + f"{MAX_MANAGED_BHE_POLL_INTERVAL} seconds" + ) + + return ManagedBHEConfig( + secret_name=secret_name, + poll_interval=poll_interval, + ) diff --git a/src/openhound/core/clients/aws_secrets_manager.py b/src/openhound/core/clients/aws_secrets_manager.py new file mode 100644 index 00000000..7f1f7a5d --- /dev/null +++ b/src/openhound/core/clients/aws_secrets_manager.py @@ -0,0 +1,758 @@ +import json +import logging +import re +from collections.abc import Iterable, Mapping +from typing import Any, Protocol + +import boto3 +from botocore.config import Config +from botocore.exceptions import ( + BotoCoreError, + ClientError, + ConfigNotFound, + CredentialRetrievalError, + HTTPClientError, + InvalidConfigError, + MetadataRetrievalError, + NoCredentialsError, + NoRegionError, + ParamValidationError, + PartialCredentialsError, + ProfileNotFound, + TokenRetrievalError, +) + +logger = logging.getLogger(__name__) + +MAX_BATCH_SIZE = 20 +_PROVIDER = "aws.secretsmanager" +_PROVIDER_ERROR_TYPES = ( + ConfigNotFound, + CredentialRetrievalError, + InvalidConfigError, + MetadataRetrievalError, + NoCredentialsError, + NoRegionError, + PartialCredentialsError, + ProfileNotFound, + TokenRetrievalError, +) + +# Bounded defaults for AWS control-plane requests. Client construction reuses +# this policy; injected clients retain their caller-supplied configuration. +DEFAULT_AWS_CLIENT_CONFIG = Config( + connect_timeout=10, + read_timeout=20, + retries={"mode": "standard", "total_max_attempts": 4}, +) + + +def _configure_safe_sdk_logging() -> None: + """Silence AWS SDK logs process-wide; they can contain requests and secrets. + + Existing children may have explicit DEBUG levels or their own handlers. + Namespace sinks also catch newly created children, including DEBUG children. + OpenHound's separate logger retains safe request/result context. + """ + for namespace in ("boto3", "botocore"): + sdk_logger = logging.getLogger(namespace) + sdk_logger.setLevel(logging.CRITICAL + 1) + sdk_logger.handlers.clear() + sdk_logger.addHandler(logging.NullHandler()) + sdk_logger.propagate = False + for name, child_logger in tuple(logging.Logger.manager.loggerDict.items()): + if name.startswith(("boto3.", "botocore.")) and isinstance( + child_logger, logging.Logger + ): + child_logger.disabled = True + # Ancestor handlers still receive descendant records even when the + # ancestor is disabled. Route new descendants to the namespace sink. + child_logger.handlers.clear() + child_logger.propagate = True + + +class SecretsManagerClient(Protocol): + def get_secret_value(self, *, SecretId: str) -> Mapping[str, Any]: ... + + def batch_get_secret_value( + self, *, SecretIdList: list[str] + ) -> Mapping[str, Any]: ... + + +class SecretRetrievalError(Exception): + """Base class for errors retrieving a secret.""" + + failure_summary = "AWS secret retrieval failed" + + def __init__(self, message: str, **context: Any) -> None: + super().__init__(message) + self.context = context + # Only the SDK classifier supplies diagnostics suitable for propagation. + self.safe_diagnostics: tuple[str, dict[str, Any]] | None = None + + +class SecretNotFoundError(SecretRetrievalError): + """The requested secret does not exist.""" + + failure_summary = "Credential secret not found" + + +class SecretPermissionError(SecretRetrievalError): + """The configured AWS identity cannot read the requested secret.""" + + failure_summary = "AWS secret access denied" + + +class AWSConfigurationError(SecretRetrievalError): + """The AWS SDK or Secrets Manager client is not configured correctly.""" + + failure_summary = "AWS credentials unavailable" + + +class SecretRequestError(SecretRetrievalError): + """AWS rejected a request for an otherwise unclassified reason.""" + + +class SecretBatchError(SecretRetrievalError): + """One or more requested secrets could not be retrieved.""" + + def __init__(self, failures: Mapping[str, SecretRetrievalError]) -> None: + self.failures = dict(failures) + super().__init__( + "; ".join( + f"secret {secret_id}: {error}" for secret_id, error in failures.items() + ) + ) + + +SecretValue = str | dict[str, Any] + + +class AWSSecretsManager: + def __init__(self, client: SecretsManagerClient | None = None) -> None: + self._client = client + + def get_secret(self, secret_id: str) -> SecretValue: + """Retrieve one text or JSON-object secret.""" + try: + client = self._get_client() + except SecretRetrievalError as error: + _log_failure("get_secret_value", error, 1, 1) + raise + return self._retrieve_secret(client, secret_id) + + def get_secrets(self, secret_ids: Iterable[str]) -> dict[str, SecretValue]: + """Retrieve any number of secrets in AWS-supported batch sizes.""" + if isinstance(secret_ids, str): + raise TypeError("secret_ids must be an iterable of secret identifiers") + requested_ids = list(secret_ids) + if not requested_ids: + return {} + + try: + client = self._get_client() + except SecretRetrievalError as error: + failures = {secret_id: error for secret_id in requested_ids} + _log_failure( + "batch_get_secret_value", + error, + len(requested_ids), + len(failures), + ) + raise SecretBatchError(failures) from None + return self._get_secrets_in_batches(client, requested_ids) + + def _get_secrets_in_batches( + self, client: SecretsManagerClient, secret_ids: list[str] + ) -> dict[str, SecretValue]: + values: dict[str, SecretValue] = {} + failures: dict[str, SecretRetrievalError] = {} + for chunk_index, chunk in enumerate(_chunks(secret_ids, MAX_BATCH_SIZE)): + _log_debug("batch_get_secret_value", "attempt", len(chunk), chunk_index + 1) + try: + response = client.batch_get_secret_value(SecretIdList=chunk) + except ( + ClientError, + BotoCoreError, + FileNotFoundError, + PermissionError, + RuntimeError, + ) as error: + classified = _classify_aws_error(error, region=_client_region(client)) + chunk_values: dict[str, SecretValue] = {} + chunk_failures = {secret_id: classified for secret_id in chunk} + else: + chunk_values, chunk_failures = _parse_batch_response( + response, chunk, _client_region(client) + ) + + values.update(chunk_values) + failures.update(chunk_failures) + if chunk_failures: + _log_failure( + "batch_get_secret_value", + SecretBatchError(chunk_failures), + len(chunk), + len(chunk_failures), + chunk_index + 1, + ) + else: + _log_debug( + "batch_get_secret_value", "success", len(chunk), chunk_index + 1 + ) + + if failures: + raise SecretBatchError(failures) from None + + return { + secret_id: values[secret_id] + for secret_id in secret_ids + if secret_id in values + } + + def _get_client(self) -> SecretsManagerClient: + _configure_safe_sdk_logging() + if self._client is not None: + return self._client + try: + return boto3.client( + "secretsmanager", + config=DEFAULT_AWS_CLIENT_CONFIG, + ) + except ( + ClientError, + BotoCoreError, + FileNotFoundError, + PermissionError, + RuntimeError, + ) as error: + raise _classify_aws_error(error) from None + + def _retrieve_secret( + self, client: SecretsManagerClient, secret_id: str + ) -> SecretValue: + _log_debug("get_secret_value", "attempt", 1) + try: + response = client.get_secret_value(SecretId=secret_id) + value = _parse_secret_value(response) + except ( + ClientError, + BotoCoreError, + FileNotFoundError, + PermissionError, + RuntimeError, + ) as error: + classified = _classify_aws_error( + error, secret_id=secret_id, region=_client_region(client) + ) + _log_failure("get_secret_value", classified, 1, 1) + raise classified from None + except SecretRetrievalError as error: + _log_failure("get_secret_value", error, 1, 1) + raise + + _log_debug("get_secret_value", "success", 1) + return value + + +def _chunks(values: list[str], size: int) -> Iterable[list[str]]: + for start in range(0, len(values), size): + yield values[start : start + size] + + +def _client_region(client: SecretsManagerClient) -> str | None: + return getattr(getattr(client, "meta", None), "region_name", None) + + +def _parse_secret_value(response: Mapping[str, Any]) -> SecretValue: + secret_string = response.get("SecretString") + if not isinstance(secret_string, str): + # boto3 already base64-decodes SecretBinary. BHE's secret writer uses + # this representation for JSON documents. + secret_binary = response.get("SecretBinary") + if not isinstance(secret_binary, bytes): + raise SecretRequestError( + "AWS Secrets Manager response omitted a text payload" + ) + try: + secret_string = secret_binary.decode("utf-8") + except UnicodeDecodeError: + raise SecretRequestError("AWS secret payload is not UTF-8 text") from None + + if secret_string.lstrip().startswith("{"): + try: + parsed = json.loads(secret_string) + except json.JSONDecodeError: + return secret_string + if isinstance(parsed, dict): + return parsed + + return secret_string + + +def _parse_batch_response( + response: Mapping[str, Any], + requested_ids: list[str], + region: str | None = None, +) -> tuple[dict[str, SecretValue], dict[str, SecretRetrievalError]]: + values: dict[str, SecretValue] = {} + failures: dict[str, SecretRetrievalError] = {} + for entry in response.get("SecretValues", []): + secret_id = next( + ( + entry.get(field) + for field in ("Name", "ARN") + if entry.get(field) in requested_ids + ), + None, + ) + if not isinstance(secret_id, str): + continue + try: + values[secret_id] = _parse_secret_value(entry) + except SecretRetrievalError as value_error: + failures[secret_id] = value_error + + for entry in response.get("Errors", []): + secret_id = entry.get("SecretId") + if isinstance(secret_id, str) and secret_id in requested_ids: + failures[secret_id] = _classify_aws_error( + ClientError( + { + "Error": { + "Code": entry.get("ErrorCode"), + "Message": entry.get("Message", ""), + }, + "ResponseMetadata": response.get("ResponseMetadata", {}), + }, + "BatchGetSecretValue", + ), + secret_id=secret_id, + region=region, + ) + + for secret_id in requested_ids: + if secret_id not in values and secret_id not in failures: + failures[secret_id] = SecretRequestError( + "AWS batch response omitted a requested secret" + ) + + return values, failures + + +def _classify_aws_error( + error: ClientError + | BotoCoreError + | FileNotFoundError + | PermissionError + | RuntimeError, + *, + secret_id: str | None = None, + region: str | None = None, +) -> SecretRetrievalError: + """Map AWS SDK/provider failures to payload-free retrieval exceptions.""" + context: dict[str, Any] = {"sdk_exception_type": type(error).__name__} + location = "" + if secret_id is not None: + context["secret_id"] = secret_id + location += f" for secret {secret_id}" + if region is not None: + context["region"] = region + location += f" in {region}" + if isinstance(error, RuntimeError): + classified = _classify_runtime_error(error, context) + elif isinstance(error, (FileNotFoundError, PermissionError)): + classified = _classify_file_error(error, context, location) + elif isinstance(error, ClientError): + classified = _classify_client_error(error, context, location) + elif isinstance(error, _PROVIDER_ERROR_TYPES): + classified = _classify_provider_error(error, context) + else: + classified = _classify_transport_error(error, context, location) + classified.safe_diagnostics = _safe_aws_diagnostics(error, classified) + return classified + + +def _safe_aws_diagnostics( + error: Exception, classified: SecretRetrievalError +) -> tuple[str, dict[str, Any]]: + """Keep actionable causes without copying SDK messages, URLs or secret IDs.""" + context = classified.context + safe: dict[str, Any] = {"sdk_exception_type": type(error).__name__} + reason = "AWS request failed; check endpoint, network and TLS configuration." + if isinstance(error, ClientError): + operation = context.get("aws_operation") + if operation not in { + "GetSecretValue", + "BatchGetSecretValue", + "GetCallerIdentity", + "AssumeRole", + "AssumeRoleWithWebIdentity", + "AssumeRoleWithSAML", + "GetSessionToken", + "GetRoleCredentials", + "CreateToken", + }: + operation = "AWS operation" + code = context.get("aws_error_code") + if code not in { + "AccessDenied", + "AccessDeniedException", + "AuthorizationError", + "UnauthorizedOperation", + "ResourceNotFoundException", + "ExpiredToken", + "ExpiredTokenException", + "UnrecognizedClientException", + "InvalidClientTokenId", + "InvalidSignatureException", + "SignatureDoesNotMatch", + "DecryptionFailure", + "ThrottlingException", + "InternalServiceError", + }: + code = "Unknown" + safe.update(aws_operation=operation, aws_error_code=code) + reason = f"{operation} returned {code}." + if code == "DecryptionFailure" or "kms:Decrypt" in str( + error.response.get("Error", {}).get("Message", "") + ): + reason += " Check KMS key permissions and kms:Decrypt." + elif operation.startswith(("AssumeRole", "GetSessionToken")): + reason += ( + " Check workload identity configuration and role trust conditions." + ) + elif isinstance(classified, SecretPermissionError): + resource = ( + "secret" + if operation in {"GetSecretValue", "BatchGetSecretValue"} + else "resource" + ) + reason += f" Check the runtime AWS identity and {resource} permissions." + elif isinstance(classified, SecretNotFoundError): + resource = ( + "secret" + if operation in {"GetSecretValue", "BatchGetSecretValue"} + else "resource" + ) + reason += f" Check the {resource} identifier and runtime AWS region." + elif isinstance(classified, AWSConfigurationError): + reason += " Check runtime AWS credential validity, signing and expiration." + elif type(error) in (NoRegionError, NoCredentialsError, RuntimeError): + # These classifier branches produce fixed, payload-free explanations. + reason = str(classified) + elif isinstance(error, (FileNotFoundError, PermissionError)): + reason = "Runtime AWS credential file access failed; check paths, mounts and permissions." + elif isinstance(error, _PROVIDER_ERROR_TYPES): + reason = "Runtime AWS credential configuration failed; check the credential provider and workload identity." + elif isinstance(error, ParamValidationError): + reason = "AWS SDK parameter validation failed; check credential acquisition and secret retrieval parameters." + provider = ( + error.kwargs.get("provider") if isinstance(error, BotoCoreError) else None + ) + if provider in ("container-role", "custom-process"): + safe["credential_provider"] = provider + transport = context.get("transport_exception_type") + if transport in { + "ReadTimeoutError", + "ConnectTimeoutError", + "EndpointConnectionError", + "SSLError", + "TimeoutError", + "ConnectionError", + "ValueError", + }: + safe["transport_exception_type"] = transport + region = context.get("region") + if isinstance(region, str) and re.fullmatch(r"[a-z]{2}(?:-[a-z]+)+-\d+", region): + safe["region"] = region + request_id = context.get("RequestId") + if isinstance(request_id, str) and re.fullmatch( + r"[0-9a-fA-F]{8}(?:-[0-9a-fA-F]{4}){3}-[0-9a-fA-F]{12}", request_id + ): + safe["RequestId"] = request_id + for key, maximum in ( + ("HTTPStatusCode", 599), + ("RetryAttempts", 100), + ("credential_endpoint_http_status", 599), + ): + value = context.get(key) + if type(value) is int and 0 <= value <= maximum: + safe[key] = value + for key in ("MaxAttemptsReached", "RetryQuotaReached"): + if type(context.get(key)) is bool: + safe[key] = context[key] + if type(error) not in (NoRegionError, NoCredentialsError, RuntimeError): + reason = f"{reason} ({type(error).__name__})" + return reason, safe + + +def _classify_runtime_error( + error: RuntimeError, context: dict[str, Any] +) -> SecretRetrievalError: + # Botocore uses a plain RuntimeError for this credential refresh failure. + expired_refresh_message = ( + "Credentials were refreshed, but the refreshed credentials are still expired." + ) + if str(error) == expired_refresh_message: + return AWSConfigurationError( + f"{expired_refresh_message} Check the credential provider and system clock.", + **context, + ) + return SecretRequestError( + f"AWS secret retrieval failed ({type(error).__name__})", + **context, + ) + + +def _classify_file_error( + error: FileNotFoundError | PermissionError, context: dict[str, Any], location: str +) -> AWSConfigurationError: + reason = error.strerror or ( + "File not found" + if isinstance(error, FileNotFoundError) + else "Permission denied" + ) + if error.filename is not None: + context["file_path"] = error.filename + reason += f": {error.filename}" + if error.errno is not None: + context["os_errno"] = error.errno + return AWSConfigurationError( + f"AWS file access{location} failed ({type(error).__name__}): {reason}. " + "Check the configured file path, mounts and permissions.", + **context, + ) + + +def _classify_client_error( + error: ClientError, context: dict[str, Any], location: str +) -> SecretRetrievalError: + code = error.response.get("Error", {}).get("Code", "Unknown") + message = error.response.get("Error", {}).get("Message") or "" + if code in {"SignatureDoesNotMatch", "InvalidSignatureException"}: + # Signature diagnostics may echo canonical requests and session tokens. + message = ( + "AWS could not validate the request signature; " + "check signing configuration, credentials and system clock" + ) + operation = error.operation_name + context.update(aws_error_code=code, aws_operation=operation) + details = _aws_response_details(error.response.get("ResponseMetadata", {}), context) + credential_operation = operation in { + "AssumeRole", + "AssumeRoleWithWebIdentity", + "AssumeRoleWithSAML", + "GetSessionToken", + "GetRoleCredentials", + "CreateToken", + } + error_type: type[SecretRetrievalError] = SecretRequestError + guidance = "" + if credential_operation: + error_type = AWSConfigurationError + service = ( + "STS " + if operation.startswith("AssumeRole") or operation == "GetSessionToken" + else "AWS " + ) + action = f"Cannot acquire AWS credentials: {service}{operation}" + if operation == "AssumeRoleWithWebIdentity" and code in { + "AccessDenied", + "AccessDeniedException", + }: + guidance = ( + " Check workload identity configuration and role trust conditions." + ) + else: + # Only label the service when the operation identifies it unambiguously. + service = ( + "Secrets Manager" + if operation in {"GetSecretValue", "BatchGetSecretValue"} + else "AWS" + ) + action = f"{service} {operation}" + if code == "ResourceNotFoundException": + error_type = SecretNotFoundError + guidance = " Check the secret identifier and region." + elif code in { + "AccessDenied", + "AccessDeniedException", + "AuthorizationError", + "UnauthorizedOperation", + }: + error_type = SecretPermissionError + guidance = " Check the AWS identity and resource permissions." + elif code in { + "UnrecognizedClientException", + "InvalidClientTokenId", + "ExpiredToken", + "ExpiredTokenException", + "InvalidSignatureException", + "SignatureDoesNotMatch", + }: + error_type = AWSConfigurationError + guidance = " Check AWS credential validity and expiration." + explanation = f"{action}{location} returned {code}" + if message: + explanation += f": {message.rstrip('.')}" + return error_type( + f"{explanation}.{guidance}" + + (" " + ", ".join(details) + "." if details else ""), + **context, + ) + + +def _aws_response_details( + metadata: Mapping[str, Any], context: dict[str, Any] +) -> list[str]: + """Preserve only diagnostic response metadata in the message and context.""" + details = [] + for key, label in ( + ("RequestId", "Request ID"), + ("HTTPStatusCode", "HTTP"), + ("RetryAttempts", "SDK retries"), + ): + if metadata.get(key) not in (None, ""): + context[key] = metadata[key] + details.append(f"{label} {metadata[key]}") + for key, label in ( + ("MaxAttemptsReached", "SDK maximum attempts reached"), + ("RetryQuotaReached", "SDK retry quota exhausted"), + ): + if key in metadata: + context[key] = metadata[key] + if metadata[key] is True: + details.append(label) + return details + + +def _classify_provider_error( + error: BotoCoreError, context: dict[str, Any] +) -> AWSConfigurationError: + provider = error.kwargs.get("provider") + if isinstance(error, MetadataRetrievalError) or ( + isinstance(error, CredentialRetrievalError) and provider == "container-role" + ): + # Container metadata errors can embed the entire credential response body. + provider_message = error.kwargs.get("error_msg", "") + status = re.search( + r"Received non 200 response (\d{3}) from container metadata:", + provider_message, + ) + explanation = "AWS credential metadata retrieval failed" + if provider is not None: + explanation += f" from provider {provider}" + if status is not None: + context["credential_endpoint_http_status"] = int(status[1]) + explanation += f": credential endpoint returned HTTP {status[1]}" + elif provider_message.removeprefix("Error retrieving metadata: ") == ( + "Unable to parse JSON returned from container metadata services" + ): + explanation += ( + ": unable to parse JSON returned from the credential endpoint" + ) + else: + # The SDK wraps transport errors in metadata/provider exceptions. + # Preserve the underlying type without echoing URLs or response bodies. + cause = error.__cause__ or error.__context__ + if isinstance(cause, MetadataRetrievalError): + cause = cause.__cause__ or cause.__context__ + if isinstance(cause, Exception): + context["transport_exception_type"] = type(cause).__name__ + explanation += f": credential endpoint failure ({type(cause).__name__})" + explanation += ( + ". Check the metadata endpoint and workload credential configuration." + ) + elif isinstance(error, CredentialRetrievalError) and provider == "custom-process": + # credential_process failures can embed helper stderr. + explanation = "AWS credential_process helper failed to provide credentials. Check the helper command and its configuration (provider custom-process)." + elif isinstance(error, NoRegionError): + explanation = "AWS region is missing. Set AWS_DEFAULT_REGION or configure region in the selected AWS profile." + elif isinstance(error, NoCredentialsError): + explanation = f"{error}. Check AWS credential provider configuration (workload role, environment or selected AWS profile)." + else: + explanation = str(error) + return AWSConfigurationError(f"{explanation} ({type(error).__name__})", **context) + + +def _classify_transport_error( + error: BotoCoreError, context: dict[str, Any], location: str +) -> SecretRequestError: + if isinstance(error, ParamValidationError): + # Validation reports can echo WebIdentityToken and other authentication inputs. + explanation = "AWS SDK parameter validation failed; check the parameters for secret retrieval or credential acquisition." + # Retain known SDK parameter/rule fields, excluding the supplied value. + valid_types = r"" + r"(?:, )*" + details = re.findall( + r"^(Invalid (?:length|value|type) for (?:document )?parameter [\w.\[\]]+), " + r"value: [^\n]*, (?:type: , )?" + rf"(valid min (?:length|value): -?\d+(?:\.\d+)?|valid types: {valid_types})$", + error.kwargs.get("report", ""), + re.MULTILINE, + ) + if details: + explanation = "AWS SDK parameter validation failed: " + "; ".join( + f"{parameter}, {constraint}" for parameter, constraint in details + ) + elif type(error) is HTTPClientError: + # The generic HTTP wrapper echoes arbitrary nested exceptions; specific + # SDK transport errors retain useful endpoint/TLS/timeout explanations. + explanation = "AWS HTTP transport failed." + transport_error = error.kwargs.get("error") + if isinstance(transport_error, Exception): + context["transport_exception_type"] = type(transport_error).__name__ + explanation = ( + f"AWS HTTP transport failed: {type(transport_error).__name__}." + ) + else: + explanation = str(error) + return SecretRequestError( + f"AWS secret retrieval{location} failed ({type(error).__name__}): {explanation}", + **context, + ) + + +def _log_debug( + operation: str, + outcome: str, + secret_count: int, + batch_number: int | None = None, +) -> None: + extra: dict[str, str | int] = { + "provider": _PROVIDER, + "operation": operation, + "outcome": outcome, + "secret_count": secret_count, + } + if batch_number is not None: + extra["batch_number"] = batch_number + message = ( + "Secret retrieval attempt" + if outcome == "attempt" + else "Secret retrieval result" + ) + logger.debug(message, extra=extra) + + +def _log_failure( + operation: str, + error: SecretRetrievalError, + secret_count: int, + failure_count: int, + batch_number: int | None = None, +) -> None: + extra: dict[str, str | int] = { + "provider": _PROVIDER, + "operation": operation, + "outcome": "failure", + "failure_category": type(error).__name__, + "secret_count": secret_count, + "failure_count": failure_count, + } + if batch_number is not None: + extra["batch_number"] = batch_number + message, context = error.safe_diagnostics or (str(error), error.context) + extra.update(context) + # Managed credential callers report the cause at normal logging levels. + logger.debug("AWS %s failed: %s", operation, message, extra=extra) diff --git a/src/openhound/core/clients/bhe_credentials.py b/src/openhound/core/clients/bhe_credentials.py new file mode 100644 index 00000000..73dcfaaf --- /dev/null +++ b/src/openhound/core/clients/bhe_credentials.py @@ -0,0 +1,112 @@ +"""Managed BHE credential retrieval and validation.""" + +import logging +import time +from dataclasses import dataclass +from typing import Protocol +from uuid import UUID + +from .aws_secrets_manager import AWSSecretsManager, SecretRetrievalError, SecretValue + +logger = logging.getLogger(__name__) + + +class SecretReader(Protocol): + def get_secret(self, secret_id: str) -> SecretValue: ... + + +@dataclass(frozen=True) +class BHECredentials: + """The HMAC credentials required by the BHE API client.""" + + token_key: str + token_id: str + + +class InvalidBHECredentialsSecret(ValueError): + """The configured AWS secret cannot be used for BHE authentication.""" + + +class AWSBHECredentials: + """Synchronously load BHE credentials from one AWS Secrets Manager secret.""" + + def __init__( + self, + secret_name: str, + secrets_manager: SecretReader | None = None, + ) -> None: + self.secret_name = secret_name + self._secrets_manager = secrets_manager or AWSSecretsManager() + + def refresh(self) -> BHECredentials: + """Fetch and validate the current credentials. + + This intentionally does no background work. A caller can refresh at + startup or in response to a 401 and continue its original operation as + soon as this method returns. + """ + try: + return self._load_credentials() + except Exception as error: + # BHE request wrappers can replace this exception before the scheduler + # sees it. Report the safe cause here; the caller owns retry decisions. + logger.error( + "Cannot refresh BHE credentials from secret %s: %s.", + self.secret_name, + str(error).rstrip(".") + if isinstance( + error, (SecretRetrievalError, InvalidBHECredentialsSecret) + ) + else f"Unexpected credential reader failure ({type(error).__name__})", + ) + raise + + def _load_credentials(self) -> BHECredentials: + secret = self._secrets_manager.get_secret(self.secret_name) + # TODO: The BHE secret payload format is still in flight. It may be JSON + # with token_id/token_key fields or an encoded binary value to decode + # before parsing; keep this conversion boundary easy to adapt. + if not isinstance(secret, dict): + raise InvalidBHECredentialsSecret( + "BHE credentials secret must be a JSON object containing token_id and token_key" + ) + + token_id = secret.get("token_id") + token_key = secret.get("token_key") + if not isinstance(token_id, str) or not token_id: + raise InvalidBHECredentialsSecret( + "BHE credentials secret must contain a non-empty token_id" + ) + try: + UUID(token_id) + except ValueError: + raise InvalidBHECredentialsSecret( + "BHE credentials secret must contain a valid UUID token_id" + ) from None + if not isinstance(token_key, str) or not token_key: + raise InvalidBHECredentialsSecret( + "BHE credentials secret must contain a non-empty token_key" + ) + return BHECredentials(token_id=token_id, token_key=token_key) + + def poll_until_available(self, interval: int) -> BHECredentials: + """Keep checking the configured secret until it contains usable credentials.""" + + while True: + try: + return self._load_credentials() + except Exception as error: # noqa: BLE001 - bootstrap retries all reader failures + # AWS errors can be transient, and an absent or not-yet-populated + # secret is expected while a managed deployment is provisioning. + logger.warning( + "Cannot load bootstrap BHE credentials from secret %s: %s. " + "Retrying bootstrap retrieval in %s seconds.", + self.secret_name, + str(error).rstrip(".") + if isinstance( + error, (SecretRetrievalError, InvalidBHECredentialsSecret) + ) + else f"Unexpected credential reader failure ({type(error).__name__})", + interval, + ) + time.sleep(interval) diff --git a/src/openhound/core/clients/bloodhound.py b/src/openhound/core/clients/bloodhound.py index c33427bc..e96f529c 100644 --- a/src/openhound/core/clients/bloodhound.py +++ b/src/openhound/core/clients/bloodhound.py @@ -5,6 +5,7 @@ import logging from abc import ABC, abstractmethod from datetime import timedelta +from typing import BinaryIO import requests from dlt.common import json @@ -45,7 +46,7 @@ def request( self, method: str, path: str, - body: bytes | None = None, + body: bytes | BinaryIO | None = None, extra_headers: dict[str, str] | None = None, timeout: tuple[float, float] | None = None, ): ... @@ -55,7 +56,7 @@ def _request( method: str, path: str, headers: dict[str, str], - body: bytes | None = None, + body: bytes | BinaryIO | None = None, extra_headers: dict[str, str] | None = None, timeout: tuple[float, float] | None = None, ): @@ -169,7 +170,7 @@ def request( self, method: str, path: str, - body: bytes | None = None, + body: bytes | BinaryIO | None = None, extra_headers: dict[str, str] | None = None, timeout: tuple[float, float] | None = None, ): @@ -187,8 +188,15 @@ def request( # HMAC Part 3 digester = hmac.new(digester.digest(), None, hashlib.sha256) - if body is not None: + if isinstance(body, bytes): digester.update(body) + elif body is not None: + position = body.tell() + try: + while chunk := body.read(1024 * 1024): + digester.update(chunk) + finally: + body.seek(position) sig = base64.b64encode(digester.digest()).decode() headers = { @@ -217,7 +225,7 @@ def request( self, method: str, path: str, - body: bytes | None = None, + body: bytes | BinaryIO | None = None, extra_headers: dict[str, str] | None = None, timeout: tuple[float, float] | None = None, ): diff --git a/src/openhound/core/clients/bloodhound_enterprise.py b/src/openhound/core/clients/bloodhound_enterprise.py index 78de5e7d..f9c9e908 100644 --- a/src/openhound/core/clients/bloodhound_enterprise.py +++ b/src/openhound/core/clients/bloodhound_enterprise.py @@ -6,19 +6,25 @@ import math import socket import time +from collections.abc import Callable from enum import Enum from pathlib import Path -from typing import Callable, TypeVar +from typing import Any, BinaryIO, TypeVar +from urllib.parse import quote +from uuid import UUID -import openhound import requests + +from openhound.core.clients.bhe_credentials import BHECredentials from openhound.core.clients.bloodhound import BloodHound, BloodHoundHTTPError from openhound.core.clients.models.jobs import ( + ArtifactUploadSession, + CollectorJob, + CollectorJobsAvailable, JobsAvailable, JobsCurrent, JobsEnd, JobStart, - ArtifactUploadSession, ManagementAvailable, ManagementOperationResult, ManagementOperationStatus, @@ -32,16 +38,90 @@ class JobStatus(str, Enum): FAILED = "failed" +class ManagedJobOutcome(str, Enum): + SUCCEEDED = "succeeded" + FAILED = "failed" + + +MANAGED_JOB_END_MAX_ATTEMPTS = 4 +MANAGED_JOB_END_RETRY_DELAY_SECONDS = 2 +MANAGED_JOB_END_CONNECT_TIMEOUT_SECONDS = 10 +MANAGED_JOB_END_READ_TIMEOUT_SECONDS = 120 +MANAGED_JOB_LIFECYCLE_CONNECT_TIMEOUT_SECONDS = 10 +MANAGED_JOB_LIFECYCLE_READ_TIMEOUT_SECONDS = 20 SUPPORT_BUNDLE_PART_SIZE = 8 * 1024 * 1024 # 8 MiB SUPPORT_BUNDLE_MAX_RETRIES = 3 SUPPORT_BUNDLE_RETRY_DELAY_SECONDS = 2 SUPPORT_BUNDLE_CONNECT_TIMEOUT_SECONDS = 10 SUPPORT_BUNDLE_READ_TIMEOUT_SECONDS = 120 +MANAGED_COLLECTOR_JOB_AVAILABLE_CONNECT_TIMEOUT_SECONDS = 10 +MANAGED_COLLECTOR_JOB_AVAILABLE_READ_TIMEOUT_SECONDS = 20 +MANAGED_JOB_CLAIM_MAX_ATTEMPTS = 4 +MANAGED_JOB_CLAIM_RETRY_DELAY_SECONDS = 2 +MANAGED_JOB_CLAIM_CONNECT_TIMEOUT_SECONDS = 10 +MANAGED_JOB_CLAIM_READ_TIMEOUT_SECONDS = 20 T = TypeVar("T") +def _safe_managed_request_error(error: Exception) -> Exception: + """Preserve actionable status/type without exposing provider response content.""" + if isinstance(error, BloodHoundHTTPError): + return BloodHoundHTTPError("Managed collector request failed", error.code) + if isinstance(error, requests.RequestException): + return type(error)("Managed collector request failed") + return RuntimeError("Managed collector request failed") + + class BloodHoundEnterprise(BloodHound): + def __init__( + self, + token_key: str, + token_id: str, + bhe_uri: str = "http://localhost:8000", + credential_refresh: Callable[[], BHECredentials] | None = None, + ) -> None: + super().__init__(token_key=token_key, token_id=token_id, bhe_uri=bhe_uri) + self._credential_refresh = credential_refresh + + def request( + self, + method: str, + path: str, + body: bytes | BinaryIO | None = None, + extra_headers: dict[str, str] | None = None, + timeout: tuple[float, float] | None = None, + ): + """Make a BHE request, refreshing managed credentials once after a 401.""" + position = ( + body.tell() if body is not None and not isinstance(body, bytes) else None + ) + try: + return super().request( + method=method, + path=path, + body=body, + extra_headers=extra_headers, + timeout=timeout, + ) + except BloodHoundHTTPError as error: + if error.code != 401 or self._credential_refresh is None: + raise + + if body is not None and not isinstance(body, bytes) and position is not None: + body.seek(position) + logger.info("BHE authentication failed; refreshing managed credentials.") + credentials = self._credential_refresh() + self.token_key = credentials.token_key + self.token_id = credentials.token_id + return super().request( + method=method, + path=path, + body=body, + extra_headers=extra_headers, + timeout=timeout, + ) + @property def jobs_available(self) -> JobsAvailable: path = "/api/v2/jobs/available" @@ -54,6 +134,88 @@ def jobs_current(self) -> JobsCurrent: response = self.request(method="GET", path=path) return JobsCurrent.model_validate(response.json()) + def available_collector_jobs(self, job_key: str) -> CollectorJobsAvailable: + """Return the next job from BHE's managed collector queue. + + This endpoint is used only by managed OpenHound deployments and is not + part of the standard open-source/self-hosted collector workflow. + """ + encoded_job_key = quote(f"eq:{job_key}", safe="") + path = ( + "/api/v2/collie-job-queue/available?limit=1&job_key=" + f"{encoded_job_key}&sort_by=-priority&sort_by=created_at" + ) + logger.debug( + "Polling managed collector job queue.", + extra={"endpoint": path, "job_key": job_key}, + ) + try: + response = self.request( + method="GET", + path=path, + timeout=( + MANAGED_COLLECTOR_JOB_AVAILABLE_CONNECT_TIMEOUT_SECONDS, + MANAGED_COLLECTOR_JOB_AVAILABLE_READ_TIMEOUT_SECONDS, + ), + ) + except (BloodHoundHTTPError, requests.RequestException) as error: + logger.error( + "Managed collector job queue request failed.", + extra={"endpoint": path, "job_key": job_key}, + ) + raise _safe_managed_request_error(error) from None + try: + return CollectorJobsAvailable.model_validate(response.json()) + except (ValueError, TypeError, KeyError): + raise RuntimeError("Invalid managed collector queue response") from None + + def claim_managed_job(self, job_id: str) -> CollectorJob: + """Claim job from BHE's managed collector job queue. + + This is used only by managed OpenHound deployments and is not + part of the standard open-source/self-hosted collector workflow. + """ + path = f"/api/v2/collie-job-queue/{job_id}/claim" + for attempt in range(1, MANAGED_JOB_CLAIM_MAX_ATTEMPTS + 1): + context = { + "endpoint": path, + "job_id": job_id, + "attempt": attempt, + "max_attempts": MANAGED_JOB_CLAIM_MAX_ATTEMPTS, + } + logger.debug("Sending managed collector claim request.", extra=context) + try: + response = self.request( + method="POST", + path=path, + timeout=( + MANAGED_JOB_CLAIM_CONNECT_TIMEOUT_SECONDS, + MANAGED_JOB_CLAIM_READ_TIMEOUT_SECONDS, + ), + ) + except Exception as error: # noqa: BLE001 - redact provider errors + retryable = self._is_transient_request_error( + error, (requests.ConnectionError, requests.Timeout) + ) + logger.debug( + "Managed collector claim request failed.", + extra={ + **context, + "retryable": retryable, + "status_code": getattr(error, "code", None), + }, + ) + if not retryable or attempt == MANAGED_JOB_CLAIM_MAX_ATTEMPTS: + raise _safe_managed_request_error(error) from None + time.sleep(MANAGED_JOB_CLAIM_RETRY_DELAY_SECONDS) + continue + try: + return CollectorJob.model_validate(response.json()["data"]["job"]) + except (ValueError, TypeError, KeyError): + raise RuntimeError("Invalid managed collector claim response") from None + + raise AssertionError("Managed collector claim retry loop exited unexpectedly.") + def start_job(self, job_id: int) -> JobStart: path = "/api/v2/jobs/start" body = json.dumps({"id": job_id}) @@ -67,6 +229,122 @@ def end_job(self, status: JobStatus, message: str) -> JobsEnd: response = self.request(method="POST", path=path, body=job_content.encode()) return JobsEnd.model_validate(response.json()) + def start_collector_job(self, job_id: str) -> None: + """Move a claimed managed collector job to running.""" + self.request( + method="POST", + path=f"/api/v2/collie-job-queue/{job_id}/start", + timeout=( + MANAGED_JOB_LIFECYCLE_CONNECT_TIMEOUT_SECONDS, + MANAGED_JOB_LIFECYCLE_READ_TIMEOUT_SECONDS, + ), + ) + + def heartbeat_collector_job(self, job_id: str) -> None: + """Renew the claim for a running managed collector job.""" + self.request( + method="POST", + path=f"/api/v2/collie-job-queue/{job_id}/heartbeat", + timeout=( + MANAGED_JOB_LIFECYCLE_CONNECT_TIMEOUT_SECONDS, + MANAGED_JOB_LIFECYCLE_READ_TIMEOUT_SECONDS, + ), + ) + + def end_managed_job( + self, + job_id: str, + outcome: ManagedJobOutcome, + failure_message: str | None = None, + metadata: dict[str, Any] | None = None, + ) -> None: + """End a claimed managed collector job.""" + path = f"/api/v2/collie-job-queue/{job_id}/end" + payload: dict[str, Any] = {"outcome": outcome.value} + if failure_message is not None: + payload["failure_message"] = failure_message + if metadata is not None: + payload["metadata"] = metadata + body = json.dumps(payload).encode() + for attempt in range(1, MANAGED_JOB_END_MAX_ATTEMPTS + 1): + log_context = { + "job_id": job_id, + "outcome": outcome.value, + "attempt": attempt, + "max_attempts": MANAGED_JOB_END_MAX_ATTEMPTS, + } + logger.info( + "Attempting to end managed collector job %s (%s/%s).", + job_id, + attempt, + MANAGED_JOB_END_MAX_ATTEMPTS, + extra=log_context, + ) + logger.debug( + "Sending managed collector end-job request.", + extra={ + **log_context, + "endpoint": path, + "has_failure_message": failure_message is not None, + "has_metadata": metadata is not None, + }, + ) + + try: + self.request( + method="POST", + path=path, + body=body, + timeout=( + MANAGED_JOB_END_CONNECT_TIMEOUT_SECONDS, + MANAGED_JOB_END_READ_TIMEOUT_SECONDS, + ), + ) + except Exception as error: # noqa: BLE001 - redact provider errors + if isinstance(error, BloodHoundHTTPError) and error.code == 404: + logger.info( + "Managed collector job is no longer actionable; discarding local state.", + extra=log_context, + ) + return + retryable = self._is_transient_request_error( + error, + (requests.ConnectionError, requests.Timeout), + ) + logger.info( + "Managed collector job %s end attempt failed.", + job_id, + extra={**log_context, "retryable": retryable}, + ) + logger.debug( + "Managed collector job end request failed.", + extra={ + **log_context, + "retryable": retryable, + "status_code": getattr(error, "code", None), + }, + ) + + if not retryable or attempt == MANAGED_JOB_END_MAX_ATTEMPTS: + raise _safe_managed_request_error(error) from None + + logger.debug( + "Retrying managed collector job end request in %s seconds.", + MANAGED_JOB_END_RETRY_DELAY_SECONDS, + extra={ + **log_context, + "retry_delay_seconds": MANAGED_JOB_END_RETRY_DELAY_SECONDS, + }, + ) + time.sleep(MANAGED_JOB_END_RETRY_DELAY_SECONDS) + else: + logger.info( + "Managed collector job %s ended successfully.", + job_id, + extra=log_context, + ) + return + def ingest(self, data: str) -> None: path = "/api/v2/ingest" headers = { @@ -78,6 +356,30 @@ def ingest(self, data: str) -> None: method="POST", path=path, body=compressed_data, extra_headers=headers ) + def upload_collector_job_data( + self, + collector_job_id: UUID | str, + data: bytes | BinaryIO, + content_type: str = "application/json", + content_encoding: str | None = None, + ) -> None: + """Upload one data file before ending a running managed queue job. + + Media type and encoding must match the supplied file. Credentials are + refreshed through the same request path used by other managed calls. + """ + job_id = UUID(str(collector_job_id)) + headers = {"Content-Type": content_type} + if content_encoding is not None: + headers["Content-Encoding"] = content_encoding + self.request( + method="POST", + path=f"/api/v2/collie-job-queue/{job_id}/ingest", + body=data, + extra_headers=headers, + timeout=(10, 120), + ) + @property def management_available(self) -> ManagementAvailable: response = self.request( @@ -218,8 +520,11 @@ def _file_checksum(path: Path) -> str: return digest.hexdigest() @staticmethod - def _is_transient_support_bundle_error(error: Exception) -> bool: - if isinstance(error, requests.RequestException): + def _is_transient_request_error( + error: Exception, + retryable_request_errors: tuple[type[requests.RequestException], ...], + ) -> bool: + if isinstance(error, retryable_request_errors): return True return isinstance(error, BloodHoundHTTPError) and error.code in { 408, @@ -237,7 +542,10 @@ def _retry_support_bundle_request( try: return request() except Exception as error: - if not self._is_transient_support_bundle_error(error): + if not self._is_transient_request_error( + error, + (requests.RequestException,), + ): raise if retry == SUPPORT_BUNDLE_MAX_RETRIES: raise diff --git a/src/openhound/core/clients/models/jobs.py b/src/openhound/core/clients/models/jobs.py index cc218f7e..2da85e74 100644 --- a/src/openhound/core/clients/models/jobs.py +++ b/src/openhound/core/clients/models/jobs.py @@ -1,6 +1,9 @@ -from pydantic import BaseModel from datetime import datetime from enum import StrEnum +from typing import Any +from uuid import UUID + +from pydantic import BaseModel class Job(BaseModel): @@ -53,6 +56,42 @@ class JobsEnd(BaseModel): data: Job +class CollectorJob(BaseModel): + id: UUID + job_schedule_id: UUID | None + job_profile_id: UUID | None + job_type_id: int + job_key: str + version: str + params: dict[str, Any] + scope_client_id: UUID | None + external_secret_key_id: str | None + secret_key_id: str | None + priority: int + status: str + run_at: datetime + unclaimed_deadline_at: datetime + attempts: int + max_attempts: int + last_failure: str | None + claimed_by: UUID | None + claimed_at: datetime | None + claim_expires_at: datetime | None + created_at: datetime + updated_at: datetime + + +class CollectorJobsAvailableData(BaseModel): + jobs: list[CollectorJob] + + +class CollectorJobsAvailable(BaseModel): + count: int + skip: int + limit: int + data: CollectorJobsAvailableData + + class ManagementOperationType(StrEnum): SUPPORT_BUNDLE = "support_bundle" diff --git a/src/openhound/core/logging.py b/src/openhound/core/logging.py index 55ca595c..46630da5 100644 --- a/src/openhound/core/logging.py +++ b/src/openhound/core/logging.py @@ -8,14 +8,255 @@ import re import sys import time +import traceback +from collections.abc import Iterator +from contextlib import contextmanager +from copy import copy from enum import Enum -from logging.handlers import TimedRotatingFileHandler +from logging.handlers import QueueHandler, QueueListener, TimedRotatingFileHandler +from multiprocessing.managers import SyncManager from pathlib import Path +from threading import RLock +from typing import Any import dlt -import openhound from rich.console import Console from rich.logging import RichHandler +from rich.markup import escape + +import openhound +from openhound.config import get_mode_defaults + + +class _ValueRedactor(logging.Filter): + def __init__(self, values: tuple[str, ...]): + super().__init__() + self.values = tuple(value for value in values if value) + self.handlers: set[logging.Handler] = set() + + def _redact(self, value): + if isinstance(value, str): + for secret in self.values: + value = value.replace(secret, "[REDACTED]") + return value + if isinstance(value, dict): + return {self._redact(k): self._redact(v) for k, v in value.items()} + if isinstance(value, (tuple, list)): + return type(value)(self._redact(item) for item in value) + if value is None or isinstance(value, (bool, int, float)): + return value + return self._redact(str(value)) + + def filter(self, record: logging.LogRecord) -> bool: + # Preserve argument types during interpolation, then redact the result. + record.msg = self._redact(record.getMessage()) + record.args = () + if record.exc_info: + record.exc_text = self._redact( + "".join(traceback.format_exception(*record.exc_info)) + ) + # Rich handlers can format the exception again, including locals. + record.exc_info = None + record.__dict__.update( + self._redact( + { + key: value + for key, value in record.__dict__.items() + if key not in ("msg", "args") + } + ) + ) + return True + + +_active_redactors: list[_ValueRedactor] = [] +_worker_handler: "_WorkerQueueHandler | None" = None +_worker_context: dict[str, Any] = {} +_worker_extension: str | None = None + + +class _WorkerQueueHandler(QueueHandler): + def prepare(self, record: logging.LogRecord) -> logging.LogRecord: + # QueueHandler's default prepare folds exceptions into msg and discards + # exc_text. Keep them separate for JSON/text/Rich and avoid duplication. + record = copy(record) + for key, value in _worker_context.items(): + record.__dict__.setdefault(key, value) + if _worker_extension and record.name.split(".")[0] == "dlt": + record._openhound_extension = _worker_extension + # Normalize non-serializable extras and format exceptions in the worker. + # Redact the complete record before even the queue proxy serializes it. + _ValueRedactor(()).filter(record) + for redactor in _active_redactors: + redactor.filter(record) + return record + + +def _install_worker_handlers() -> None: + assert _worker_handler is not None + root = logging.getLogger() + root.handlers[:] = [_worker_handler] + dlt_logger = logging.getLogger("dlt") + dlt_logger.handlers[:] = [_worker_handler] + dlt_logger.propagate = False + dlt_logger.disabled = False + + +def configure_worker_logging(queue: Any, level: int) -> None: + """Executor initializer: workers forward, and only the parent owns files.""" + global _worker_handler + loggers = [logging.getLogger(), *logging.Logger.manager.loggerDict.values()] + handlers = set() + for logger in loggers: + if isinstance(logger, logging.Logger): + handlers.update(logger.handlers) + logger.handlers.clear() + logger.propagate = True + for handler in handlers: + handler.close() + logger_override._file_handlers.clear() + _worker_handler = _WorkerQueueHandler(queue) + logging.getLogger().setLevel(level) + logging.getLogger("dlt").setLevel(level) + # Queue proxy connection diagnostics must never recurse into that same queue. + multiprocessing_logger = logging.getLogger("multiprocessing") + multiprocessing_logger.handlers[:] = [logging.NullHandler()] + multiprocessing_logger.propagate = False + _install_worker_handlers() + + +@contextmanager +def worker_log_context(job_id: int | str, collector: str) -> Iterator[None]: + global _worker_context, _worker_extension + previous, extension = _worker_context, _worker_extension + _worker_context = {"job_id": job_id, "collector_extension": collector} + _worker_extension = None + try: + yield + finally: + _worker_context, _worker_extension = previous, extension + + +def _destination(handler: logging.Handler) -> Any: + if isinstance(handler, logging.FileHandler): + return ("file", handler.baseFilename) + if isinstance(handler, RichHandler): + return ("stream", id(handler.console.file)) + if isinstance(handler, logging.StreamHandler): + return ("stream", id(handler.stream)) + return ("handler", id(handler)) + + +class _ParentWorkerHandler(logging.Handler): + def __init__(self, configuration: "CustomLogger"): + super().__init__() + self.configuration = configuration + + def emit(self, record: logging.LogRecord) -> None: + # Read current destinations so parent-only handlers and reconfiguration + # take effect for workers too. Always include the main destinations, + # including for DLT's normally non-propagating logger. + handlers = list(logging.getLogger().handlers) + logger = logging.getLogger(record.name) + while logger.parent is not None: + handlers.extend(logger.handlers) + if not logger.propagate: + break + logger = logger.parent + extension = getattr(record, "_openhound_extension", None) + if extension: + config = self.configuration + try: + path = config._is_valid_path( + config.base_path, config.base_path / f"ext_{extension}.log" + ) + handlers.append(config._get_file_handler(path)) + except (ValueError, OSError): + # Main destinations still receive the record if an optional + # extension file is unavailable. Never report supplied paths. + sys.stderr.write("OpenHound could not open worker extension log\n") + seen = set() + for handler in handlers: + destination = _destination(handler) + if destination not in seen and record.levelno >= handler.level: + seen.add(destination) + try: + handler.handle(record) + except Exception: # noqa: BLE001 - isolate arbitrary parent handlers + # One custom destination must not kill the listener thread + # and strand every other worker logging destination. + sys.stderr.write("OpenHound worker logging destination failed\n") + + +class WorkerLogRouter: + """Parent-owned logging resources for one scheduler, across pool resets. + + A manager-backed multiprocessing queue has no worker feeder thread or shared + write lock: killing a producer cannot corrupt the listener's input pipe. + Its server and listener are explicitly stopped after all executors exit. + """ + + def __init__(self, context: Any, configuration: "CustomLogger | None" = None): + self.context = context + self.configuration = configuration or logger_override + self.manager: SyncManager | None = None + self.listener: QueueListener | None = None + self._queue: Any = None + + @property + def queue(self) -> Any: + if self.manager is None: + self.manager = SyncManager(ctx=self.context) + try: + self.manager.start() + self._queue = self.manager.Queue() + self.listener = QueueListener( + self._queue, _ParentWorkerHandler(self.configuration) + ) + self.listener.start() + except BaseException: + self.close() + raise + return self._queue + + def close(self) -> None: + # Workers have exited before the sentinel; synchronous proxy puts mean + # every completed log submission is already ahead of it in the queue. + try: + if self.listener is not None: + self.listener.stop() + finally: + self.listener = None + if self.manager is not None and hasattr(self.manager, "shutdown"): + self.manager.shutdown() + self.manager = None + self._queue = None + + +def _protect_handler[H: logging.Handler](handler: H) -> H: + for redactor in _active_redactors: + handler.addFilter(redactor) + redactor.handlers.add(handler) + return handler + + +@contextmanager +def redact_log_values(values: tuple[str, ...]) -> Iterator[None]: + """Protect current and newly configured OpenHound handlers for one worker.""" + redactor = _ValueRedactor(values) + _active_redactors.append(redactor) + loggers = [logging.getLogger(), *logging.Logger.manager.loggerDict.values()] + for logger in loggers: + if isinstance(logger, logging.Logger): + for handler in logger.handlers: + _protect_handler(handler) + try: + yield + finally: + _active_redactors.remove(redactor) + for handler in redactor.handlers: + handler.removeFilter(redactor) + VALID_LEVELS = ["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"] @@ -101,7 +342,11 @@ def shouldRollover(self, record: logging.LogRecord) -> bool: # self.stream = self._open() # size = self.stream.seek(0, os.SEEK_END) - base_file_size = os.path.getsize(self.baseFilename) + base_file_size = ( + os.path.getsize(self.baseFilename) + if os.path.exists(self.baseFilename) + else 0 + ) if base_file_size >= self.max_bytes: self._size_triggered = True return True @@ -145,7 +390,9 @@ def format(self, record: logging.LogRecord) -> str: "openhound_version": openhound.__version__, } - if record.exc_info: + if record.exc_text: + log_data["exception"] = record.exc_text + elif record.exc_info: log_data["exception"] = self.formatException(record.exc_info) for key, value in record.__dict__.items(): @@ -193,7 +440,9 @@ def format(self, record: logging.LogRecord) -> str: extra_str = " ".join(f"{key}={value}" for key, value in extras.items()) log_line = f"{log_line} | {extra_str}" - if record.exc_info: + if record.exc_text: + log_line = f"{log_line}\n{record.exc_text}" + elif record.exc_info: log_line = f"{log_line}\n{self.formatException(record.exc_info)}" return log_line @@ -212,6 +461,10 @@ def format(self, record: logging.LogRecord) -> str: str: A formatted string for Rich logging """ log_fmt = f"time={self.formatTime(record, '%Y-%m-%d %H:%M:%S')}, msg={record.getMessage()} (openhound_version={openhound.__version__})" + if record.exc_text and not record.exc_info: + # The redactor cleared exc_info, so Rich cannot render the original. + # Treat the sanitized traceback (including [REDACTED]) as literal text. + log_fmt = f"{log_fmt}\n{escape(record.exc_text)}" return log_fmt @@ -226,7 +479,7 @@ def __init__( interval: int = 1, cli_level: str = "ERROR", base_path: str | None = None, - log_format: str = "text", + log_format: str | None = None, ): self.level = level self.name = name @@ -237,14 +490,16 @@ def __init__( self.backup_count = backup_count self.interval = interval self.cli_level = cli_level - self.log_format = log_format + self._log_format_override = log_format + self.log_format = self._validate_format(log_format or "text") self.base_path = Path(base_path) if base_path else self.default_platform_path() self.log_file_path: Path | None = None # Share one rotating handler per file across loggers; opening the same file # with multiple handlers breaks rotation on Windows (WinError 32). - self._file_handlers: dict[Path, "RotatingFileHandler"] = {} + self._file_handlers: dict[Path, RotatingFileHandler] = {} + self._file_handler_lock = RLock() self.handlers = { LogMode.CLI: self.cli_handlers, @@ -315,6 +570,11 @@ def set_handler(self, name: str) -> None: Args: name (str): Name of the extension or pipeline to create a specific logger for """ + global _worker_extension + if _worker_handler is not None: + _worker_extension = name + _install_worker_handlers() + return log_file_path = self.base_path / f"ext_{name}.log" valid_path = self._is_valid_path(self.base_path, log_file_path) self.dlt_logger.setLevel(self.root_logger.level) @@ -325,6 +585,10 @@ def set_handler(self, name: str) -> None: def setup(self) -> None: """Set the correct logging handler based on the detected mode""" + if _worker_handler is not None: + _install_worker_handlers() + return + # Get values from DLT config if set dlt_level = dlt.config.get("runtime.log_level", str) dlt_max_bytes = dlt.config.get("runtime.log_max_bytes", int) @@ -342,8 +606,10 @@ def setup(self) -> None: self.cli_level = self._validate_level( dlt_cli_level or self.cli_level, default="ERROR" ) + mode_log_format = get_mode_defaults().log_format self.log_format = self._validate_format( - dlt_log_format or self.log_format, default="text" + dlt_log_format or self._log_format_override or mode_log_format, + default=mode_log_format, ) # Override the base path if log_path is set in DLT config, otherwise use the default platform path @@ -386,6 +652,7 @@ def _build_file_handler(self, file_path: Path) -> "RotatingFileHandler": interval=self.interval, backupCount=self.backup_count, max_bytes=self.max_bytes, + delay=True, ) rotating_file_handler.setFormatter(self._file_formatter()) # This regular expression overrides the default extMatch to recognize both @@ -398,11 +665,18 @@ def _build_file_handler(self, file_path: Path) -> "RotatingFileHandler": def _get_file_handler(self, file_path: Path) -> "RotatingFileHandler": """Return a shared rotating file handler for the path (one open file per path).""" key = Path(file_path).resolve() - handler = self._file_handlers.get(key) - if handler is None: - handler = self._build_file_handler(file_path) - self._file_handlers[key] = handler - return handler + # Parent reconfiguration and the worker listener can request the same + # extension file concurrently. Keep a single rotation owner per path. + with self._file_handler_lock: + handler = self._file_handlers.get(key) + if handler is None: + handler = self._build_file_handler(file_path) + self._file_handlers[key] = handler + else: + # setup() may be called again after configuration changes. Keep a + # cached handler's formatter aligned with the newly resolved mode. + handler.setFormatter(self._file_formatter()) + return _protect_handler(handler) def container_handlers(self, logger: logging.Logger, file_path: Path) -> None: """Set the logging handler/format when running in a container""" @@ -413,7 +687,7 @@ def container_handlers(self, logger: logging.Logger, file_path: Path) -> None: # Output is human-readable text by default; set runtime.log_format = "JSON" for structured JSON. stdout_handler = logging.StreamHandler(sys.stdout) stdout_handler.setFormatter(formatter) - logger.addHandler(stdout_handler) + logger.addHandler(_protect_handler(stdout_handler)) # But also log the same format to a file for persistence and debugging when needed logger.addHandler(self._get_file_handler(file_path)) @@ -433,7 +707,7 @@ def cli_handlers(self, logger: logging.Logger, file_path: Path) -> None: rich_tracebacks=True, ) console_handler.setFormatter(rich_formatter) - logger.addHandler(console_handler) + logger.addHandler(_protect_handler(console_handler)) # But also save the logs to a file using the configured format (text by default) logger.addHandler(self._get_file_handler(file_path)) diff --git a/src/openhound/scheduler/aws_profile.py b/src/openhound/scheduler/aws_profile.py new file mode 100644 index 00000000..00560474 --- /dev/null +++ b/src/openhound/scheduler/aws_profile.py @@ -0,0 +1,525 @@ +"""Adapt BHE's IAM source keys to a direct-key or assume-role AWS profile. + +Only the worker installs AWS/DLT configuration or writes credential material. +The scheduler owns the workspace and its lifetime, including killed workers. + +Managed authentication preflights are a beta compatibility bridge (BED-10368). +Move them with profile generation to the AWS extension once it reports fatal +initial-authentication failures before best-effort resource iteration. +""" + +import logging +import os +import shutil +import stat +import tempfile +import traceback +from collections.abc import Callable, Iterator +from configparser import ConfigParser +from contextlib import closing, contextmanager +from copy import deepcopy +from dataclasses import dataclass, field +from io import StringIO +from pathlib import Path +from typing import Any +from uuid import uuid4 + +import boto3 +import dlt +from botocore.exceptions import BotoCoreError, ClientError +from botocore.session import Session as BotocoreSession +from dlt.common.configuration.container import Container +from dlt.common.configuration.providers.provider import ConfigProvider +from dlt.common.configuration.specs.pluggable_run_context import PluggableRunContext +from filelock import FileLock, Timeout + +from openhound.core.clients.aws_secrets_manager import ( + SecretRetrievalError, + SecretValue, + _classify_aws_error, + _configure_safe_sdk_logging, +) +from openhound.core.clients.models.jobs import CollectorJob +from openhound.core.logging import redact_log_values +from openhound.scheduler.configuration import ( + ManagedConfigurationError, + managed_job_configuration, +) + +logger = logging.getLogger(__name__) + +_ROLE_PROFILE_SETTINGS = {"role_arn", "role_session_name", "external_id"} + + +class AWSAccessKeyError(ManagedConfigurationError): + """An AWS preparation error containing no supplied credential values.""" + + def __init__( + self, + message: str, + summary: str = "Managed AWS credentials could not be prepared", + **context: Any, + ) -> None: + super().__init__(message) + self.summary = summary + self.context = context + + +def validate_secret_reference(value: Any) -> str: + if not isinstance(value, str) or not value.strip(): + raise AWSAccessKeyError( + "Managed AWS job requires a valid secret_key_id reference" + ) + return value + + +@dataclass(frozen=True) +class AWSAccessKeyCredentials: + access_key_id: str = field(repr=False) + secret_access_key: str = field(repr=False) + + def __post_init__(self) -> None: + # BHE's raw-key handoff needs strings for the INI writer, not AWS syntax. + if not isinstance(self.access_key_id, str) or not isinstance( + self.secret_access_key, str + ): + raise AWSAccessKeyError( + "Managed AWS raw credentials require a key identifier and text secret payload" + ) + + +def prepare_aws_credentials( + job: CollectorJob, load_secret: Callable[[str], SecretValue] +) -> dict[str, Any] | AWSAccessKeyCredentials: + """Prepare authoritative claimed inputs without installing worker state. + + An explicit key ID selects the beta raw-key contract. Without it, retain + compatibility with nested credential documents; a raw payload still needs + a key ID and must never fall back to the scheduler's AWS identity. + """ + reference = validate_secret_reference(job.secret_key_id) + raw_key = job.external_secret_key_id is not None + if raw_key: + aws_profile_params(job.params) + try: + credentials = load_secret(reference) + except SecretRetrievalError as error: + # A retrieval exception's type alone does not make its text safe to log. + message, context = error.safe_diagnostics or ( + f"Managed AWS credential retrieval failed ({type(error).__name__})", + {}, + ) + raise AWSAccessKeyError( + f"OpenHound runtime secret retrieval failed: {message}", + error.failure_summary, + **context, + ) from None + except Exception as error: # noqa: BLE001 - arbitrary secret readers + raise AWSAccessKeyError( + f"Credential secret reader failed ({type(error).__name__})", + "Credential secret retrieval failed", + ) from None + if not raw_key and isinstance(credentials, dict): + return credentials # Generic credential-document validation stays shared. + return AWSAccessKeyCredentials(job.external_secret_key_id, credentials) # type: ignore[arg-type] + + +def aws_profile_params( + params: dict[str, Any], *, collector_name: str = "aws" +) -> dict[str, Any]: + """Validate the profile bridge and copy its nested DLT document.""" + if collector_name != "aws": + raise AWSAccessKeyError("Managed AWS credentials require the AWS collector") + if not isinstance(params, dict) or any( + not isinstance(v, dict) for v in params.values() + ): + raise AWSAccessKeyError( + "Managed AWS params must use nested configuration sections" + ) + document = deepcopy(params) + node = document + for section in ("sources", "source", "aws"): + node = node.setdefault(section, {}) + if not isinstance(node, dict): + raise AWSAccessKeyError( + "Managed AWS source configuration must be a nested object" + ) + credentials = node.setdefault("credentials", {}) + if not isinstance(credentials, dict): + raise AWSAccessKeyError( + "Managed AWS credentials configuration must be a nested object" + ) + if credentials.keys() & _ROLE_PROFILE_SETTINGS: + role_arn = credentials.get("role_arn") + # A linked profile without a role can fall back to the worker's identity. + if not isinstance(role_arn, str) or not role_arn.strip(): + raise AWSAccessKeyError( + "Managed AWS role settings require a non-empty role_arn" + ) + return document + + +def _workspace_root() -> Path: + suffix = str(os.getuid()) if hasattr(os, "getuid") else "current-user" + root = Path(tempfile.gettempdir()) / f"openhound-aws-profiles-{suffix}" + root.mkdir(mode=0o700, exist_ok=True) + info = root.lstat() + if not stat.S_ISDIR(info.st_mode) or ( + hasattr(os, "getuid") + and (info.st_uid != os.getuid() or stat.S_IMODE(info.st_mode) != 0o700) + ): + raise AWSAccessKeyError("Managed AWS credential workspace is not private") + return root + + +def _lock(path: Path, timeout: float = 0) -> FileLock: + return FileLock(path, mode=0o600, timeout=timeout, fallback_to_soft=False) + + +def _reap_orphans(root: Path) -> None: + # Caller holds the root lock, also used for creation and worker acquisition. + # OS-backed locks release on process death, unlike PID checks/age limits. + for directory in root.glob("job-*"): + if directory.is_symlink() or not directory.is_dir(): + continue + try: + with _lock(directory / "parent.lock"), _lock(directory / "worker.lock"): + pass + shutil.rmtree(directory) + except Timeout: + continue # A scheduler or an orphaned-but-still-running worker owns it. + + +def cleanup_orphaned_aws_profiles() -> None: + root = _workspace_root() + with _lock(root / "registry.lock", timeout=-1): + _reap_orphans(root) + + +class AWSProfileWorkspace: + """Parent-owned workspace. No credential values or files exist at creation.""" + + def __init__(self) -> None: + root = _workspace_root() + with _lock(root / "registry.lock", timeout=-1): + _reap_orphans(root) + self.directory = Path(tempfile.mkdtemp(prefix="job-", dir=root)) + self._owner = _lock(self.directory / "parent.lock") + self._owner.acquire() + + def cleanup(self) -> None: + with _lock(self.directory.parent / "registry.lock", timeout=-1): + self._owner.release() + if self.directory.exists(): + shutil.rmtree(self.directory) + + +class _SourceCredentialFallback(ConfigProvider): + """Keep deployment defaults except source credentials/authentication mode. + + DLT resolves union alternatives field by field. A local profile_names or + org_role_name could otherwise select a different credential model, even with + this job's explicit profile_name. Destination/bootstrap defaults still work. + """ + + def __init__(self, provider: ConfigProvider) -> None: + self.provider = provider + + def get_value( + self, key: str, hint: type[Any], pipeline_name: str, *sections: str + ) -> tuple[Any, str]: + # Unsectioned providers cannot distinguish source credentials from other + # defaults. Exclude them entirely rather than enumerating AWS fields. + # The job provider itself is never wrapped. + if not self.provider.supports_sections: + return None, key + if (not sections or sections[0] in ("sources", "credentials")) and ( + key == "credentials" or "credentials" in sections + ): + return None, key + return self.provider.get_value(key, hint, pipeline_name, *sections) + + @property + def name(self) -> str: + return self.provider.name + + @property + def supports_secrets(self) -> bool: + return self.provider.supports_secrets + + @property + def supports_sections(self) -> bool: + return self.provider.supports_sections + + +def _write_ini(path: Path, sections: dict[str, dict[str, Any]]) -> None: + config = ConfigParser(interpolation=None) + try: + # Only the authentication bridge calls this writer. Arbitrary extension + # JSON remains in DLT, never in AWS INI files. + if any( + not isinstance(value, (str, int, float, bool)) + for settings in sections.values() + for value in settings.values() + ): + raise TypeError + config.read_dict(sections) + stream = StringIO() + config.write(stream) + serialized = stream.getvalue() + serialized.encode("utf-8") + # Require a lossless INI round trip (e.g. multiline section injection is + # not representable). This is serialization, not AWS field validation. + restored = ConfigParser(interpolation=None) + restored.read_string(serialized) + if {s: dict(restored[s]) for s in restored.sections()} != { + s: dict(config[s]) for s in config.sections() + }: + raise ValueError + except Exception: # noqa: BLE001 - serialization failures must omit supplied values + raise AWSAccessKeyError("Managed AWS profile could not be serialized") from None + # O_EXCL prevents following/reusing an existing file; permissions apply from + # creation, independent of the process umask. + with os.fdopen( + os.open(path, os.O_WRONLY | os.O_CREAT | os.O_EXCL, 0o600), + "w", + encoding="utf-8", + ) as stream: + stream.write(serialized) + + +@contextmanager +def _safe_sdk_logging() -> Iterator[None]: + # Spawned workers do not inherit the scheduler's Secrets Manager log sinks. + # SDK DEBUG responses can include the extension's temporary STS credentials. + loggers = [logging.getLogger("boto3"), logging.getLogger("botocore")] + loggers.extend( + value + for name, value in tuple(logging.Logger.manager.loggerDict.items()) + if name.startswith(("boto3.", "botocore.")) + and isinstance(value, logging.Logger) + ) + previous = { + value: (value.level, value.handlers[:], value.propagate, value.disabled) + for value in loggers + } + _configure_safe_sdk_logging() + try: + yield + finally: + for sdk_logger, (level, handlers, propagate, disabled) in previous.items(): + sdk_logger.setLevel(level) + sdk_logger.handlers = handlers + sdk_logger.propagate = propagate + sdk_logger.disabled = disabled + + +def _log_collection_failure(error: Exception) -> None: + # Collector/DLT wrappers can copy SDK messages into their own text. If an + # SDK cause exists, report only classified diagnostics, never either message: + # they can include temporary STS credentials not covered by pair redaction. + cause: BaseException | None = error + seen: set[int] = set() + while isinstance(cause, BaseException) and id(cause) not in seen: + seen.add(id(cause)) + if isinstance(cause, (BotoCoreError, ClientError)): + diagnostics = _classify_aws_error(cause).safe_diagnostics + assert diagnostics is not None # Always supplied by the SDK classifier. + message, context = diagnostics + logger.error( + "Managed AWS collection failed (%s): %s\n" + "Traceback (most recent call last):\n%s", + type(error).__name__, + message, + "".join(traceback.format_tb(error.__traceback__)), + extra=context, + ) + return + cause = ( + getattr(cause, "exception", None) or cause.__cause__ or cause.__context__ + ) + logger.exception("Managed AWS collection failed") + + +def validate_aws_profile( + profile: str, + *, + config_file: str | None = None, + credentials_file: str | None = None, +) -> None: + """Authenticate the initial collection identity before best-effort resources. + + Select the profile explicitly: a role profile must complete AssumeRole, never + retry with its source keys or the worker's default identity. Callers provide + safe SDK diagnostics and any generated-profile cleanup. This checks identity + only, not permissions for individual AWS services. + """ + sdk_session = None + if config_file or credentials_file: + # Match the extension's per-session file selection for existing-profile + # credential documents; never change process-wide files for these jobs. + sdk_session = BotocoreSession() + for variable, value in ( + ("config_file", config_file), + ("credentials_file", credentials_file), + ): + if value: + path = Path(value).expanduser().resolve() + if variable == "config_file" and not path.is_file(): + raise AWSAccessKeyError("Managed AWS config file is unavailable") + sdk_session.set_config_variable(variable, str(path)) + session = boto3.Session( + profile_name=profile, + region_name=dlt.secrets.get("sources.source.aws.credentials.region_name", str), + botocore_session=sdk_session, + ) + with closing( + session.client( + "sts", + endpoint_url=dlt.secrets.get( + "sources.source.aws.credentials.endpoint_url", str + ), + ) + ) as sts: + sts.get_caller_identity() + + +def validate_configured_aws_profiles() -> None: + """Temporary validation for managed jobs referencing existing named profiles. + + Organization setup is already validated before resource iteration by the AWS + extension. For independent profiles, preserve its skip-and-continue behavior, + but require at least one authenticated profile, regardless of resource rows. + Move this policy to the extension along with managed profile generation. + """ + prefix = "sources.source.aws.credentials" + if dlt.secrets.get(f"{prefix}.org_role_name", str) is not None: + return + profiles = dlt.secrets.get(f"{prefix}.profile_names", list) + multiple = profiles is not None + if profiles is None: + profile = dlt.secrets.get(f"{prefix}.profile_name", str) + if profile is None: + return # Leave credential-model selection to the extension/DLT. + profiles = [profile] + with _safe_sdk_logging(): + try: + authenticated = False + for profile in dict.fromkeys(profiles): + try: + validate_aws_profile( + profile, + config_file=dlt.secrets.get(f"{prefix}.config_file", str), + credentials_file=dlt.secrets.get( + f"{prefix}.credentials_file", str + ), + ) + except (BotoCoreError, ClientError) as error: + if not multiple: + raise + diagnostics = _classify_aws_error(error).safe_diagnostics + assert diagnostics is not None + message, context = diagnostics + logger.warning( + "Skipping managed AWS profile %r during authentication: %s", + profile, + message, + extra=context, + ) + continue + authenticated = True + if not authenticated: + raise ManagedConfigurationError( + "Managed AWS collection failed: no configured profile authenticated" + ) + except ManagedConfigurationError: + raise + except Exception as error: # noqa: BLE001 - sanitize SDK/config failures + _log_collection_failure(error) + raise ManagedConfigurationError( + "Managed AWS collection failed; check worker logs" + ) from None + + +@contextmanager +def managed_aws_profile( + params: dict[str, Any], credentials: AWSAccessKeyCredentials, directory: Path +) -> Iterator[str]: + """Install a direct-key or role profile; restore all worker state afterward.""" + document = aws_profile_params(params) + profile = f"openhound-job-{uuid4().hex}" + settings = document["sources"]["source"]["aws"]["credentials"] + # These fields belong to boto3's file provider, not the extension's DLT model. + role_settings = { + key: settings[key] for key in _ROLE_PROFILE_SETTINGS if key in settings + } + settings["profile_name"] = profile + source_profile = f"{profile}-source" if role_settings else profile + profile_config = ( + { + **role_settings, + "source_profile": source_profile, + "role_session_name": role_settings.get("role_session_name", profile), + } + if role_settings + else {} + ) + files = { + "AWS_SHARED_CREDENTIALS_FILE": directory / "credentials", + "AWS_CONFIG_FILE": directory / "config", + } + previous_environment = {key: os.environ.get(key) for key in files} + providers = Container()[PluggableRunContext].providers + previous_providers = providers.providers + worker_lock = _lock(directory / "worker.lock") + with _lock(directory.parent / "registry.lock", timeout=-1): + if not directory.is_dir() or directory.is_symlink(): + raise AWSAccessKeyError("Managed AWS credential workspace is unavailable") + worker_lock.acquire() + try: + with ( + redact_log_values( + (credentials.access_key_id, credentials.secret_access_key) + ), + _safe_sdk_logging(), + ): + try: + _write_ini( + files["AWS_SHARED_CREDENTIALS_FILE"], + { + source_profile: { + "aws_access_key_id": credentials.access_key_id, + "aws_secret_access_key": credentials.secret_access_key, + } + }, + ) + _write_ini( + files["AWS_CONFIG_FILE"], {f"profile {profile}": profile_config} + ) + os.environ.update({key: str(path) for key, path in files.items()}) + providers.providers = [ + _SourceCredentialFallback(p) for p in previous_providers + ] + with managed_job_configuration(document, {}): + yield profile + except ManagedConfigurationError: + raise + except Exception as error: # noqa: BLE001 - collector/SDK failures + _log_collection_failure(error) + # Keep remote tracebacks payload-free after logging locally + # while the worker's redaction and SDK controls are active. + raise ManagedConfigurationError( + "Managed AWS collection failed; check worker logs" + ) from None + finally: + providers.providers = previous_providers + for key, previous in previous_environment.items(): + if previous is None: + os.environ.pop(key, None) + else: + os.environ[key] = previous + try: + for path in files.values(): + path.unlink(missing_ok=True) + finally: + worker_lock.release() diff --git a/src/openhound/scheduler/configuration.py b/src/openhound/scheduler/configuration.py new file mode 100644 index 00000000..912195e4 --- /dev/null +++ b/src/openhound/scheduler/configuration.py @@ -0,0 +1,71 @@ +"""Job-scoped configuration using the same DLT resolver as local TOML files.""" + +from collections.abc import Iterator +from contextlib import contextmanager +from copy import deepcopy +from typing import Any + +from dlt.common.configuration.container import Container +from dlt.common.configuration.exceptions import ConfigurationException +from dlt.common.configuration.providers.doc import CustomLoaderDocProvider +from dlt.common.configuration.specs.pluggable_run_context import PluggableRunContext +from dlt.common.configuration.utils import TraceLogContext +from dlt.common.utils import update_dict_nested + + +class ManagedConfigurationError(ValueError): + """A managed configuration failure safe to report without its values.""" + + +def _validate_document(document: dict[str, Any], *, allow_empty: bool = False) -> None: + # Check only the document boundary. DLT and the collector own setting names, + # types, and values; custom TOML sections do not require framework changes. + if not isinstance(document, dict) or (not document and not allow_empty): + raise ManagedConfigurationError( + "Managed job configuration must be a non-empty nested JSON object" + ) + if any(not isinstance(value, dict) for value in document.values()): + # Never include user-supplied keys or values: either can contain secrets. + raise ManagedConfigurationError( + "Managed job configuration must use nested TOML sections " + "(for example sources.source.aws); flat configuration is not supported" + ) + + +@contextmanager +def managed_job_configuration( + params: dict[str, Any], secrets: dict[str, Any] +) -> Iterator[None]: + """Resolve full-hierarchy JSON for one worker job, then restore its providers. + + Retrieved secrets win per setting. Mappings merge recursively; lists are + individual settings and are replaced in full. Copy both documents so DLT or + collector mutations cannot affect the inputs retained by the scheduler. + """ + _validate_document(params) + _validate_document(secrets, allow_empty=True) + document = update_dict_nested(deepcopy(params), deepcopy(secrets)) + provider = CustomLoaderDocProvider( + "Managed job configuration", lambda: document, supports_secrets=True + ) + container = Container() + providers = container[PluggableRunContext].providers + previous = providers.providers + # Managed workers run one job at a time. Leave environment/TOML providers as + # fallback for deployment defaults, with this job's document taking priority. + providers.providers = [provider, *previous] + # Credential model fields can be plain strings, not DLT secret hints. Do not + # retain their values in DLT's configuration traces (including pipeline traces). + traces = TraceLogContext() + traces.logging_enabled = False + try: + with container.injectable_context(traces): + yield + except ConfigurationException: + # DLT coercion exceptions can include the value, even for credential fields. + raise ManagedConfigurationError( + "Managed job configuration could not be resolved; " + "check setting types and required fields" + ) from None + finally: + providers.providers = previous diff --git a/src/openhound/scheduler/dataflow.py b/src/openhound/scheduler/dataflow.py index 65bc1850..252abe61 100644 --- a/src/openhound/scheduler/dataflow.py +++ b/src/openhound/scheduler/dataflow.py @@ -36,23 +36,31 @@ def preprocess(extension: OpenHound) -> list[str]: return extracted_resources(load_info) -def convert(extension: OpenHound) -> list[str]: +def convert(extension: OpenHound, output_path: Path | None = None) -> list[str]: + options = {} + if output_path is not None: + options["output_path"] = output_path load_info = extension.converter( input_path=Path("/tmp/openhound") / extension.name, - method=Method.ingest, # type: ignore + method=Method.write if output_path is not None else Method.ingest, # type: ignore progress=Progress.log, # type: ignore + **options, ) return extracted_resources(load_info) -def pipeline(extension: OpenHound) -> dict[str, list[str]]: +def pipeline( + extension: OpenHound, output_path: Path | None = None +) -> dict[str, list[str]]: collect_result = collect(extension) preprocessor_result = [] if extension.preprocessor is not None: preprocessor_result = preprocess(extension) - convert_result = convert(extension) + convert_result = ( + convert(extension) if output_path is None else convert(extension, output_path) + ) return { "collect": collect_result, "preprocess": preprocessor_result, diff --git a/src/openhound/scheduler/service.py b/src/openhound/scheduler/service.py index e64a4b6e..604f80da 100644 --- a/src/openhound/scheduler/service.py +++ b/src/openhound/scheduler/service.py @@ -1,15 +1,36 @@ import logging +import multiprocessing +import shutil import signal +import tempfile import time +from collections.abc import Callable from concurrent.futures import Future, ProcessPoolExecutor from concurrent.futures.process import BrokenProcessPool -from dataclasses import dataclass +from contextlib import nullcontext +from dataclasses import dataclass, field +from datetime import UTC, datetime +from importlib import import_module from pathlib import Path +from typing import Any +from uuid import UUID import openhound import openhound.core.logging as openhound_logging -from openhound.core.clients.bloodhound_enterprise import BloodHoundEnterprise, JobStatus +from openhound.config import is_managed +from openhound.core.clients.aws_secrets_manager import ( + AWSSecretsManager, + SecretRetrievalError, +) +from openhound.core.clients.bhe_credentials import BHECredentials +from openhound.core.clients.bloodhound import BloodHoundHTTPError +from openhound.core.clients.bloodhound_enterprise import ( + BloodHoundEnterprise, + JobStatus, + ManagedJobOutcome, +) from openhound.core.clients.models.jobs import ( + CollectorJob, Job, ManagementOperation, ManagementOperationStatus, @@ -18,25 +39,95 @@ from openhound.core.manager import CollectorManager from openhound.core.support_bundle import create_support_bundle from openhound.scheduler import dataflow +from openhound.scheduler.aws_profile import ( + AWSAccessKeyCredentials, + AWSAccessKeyError, + AWSProfileWorkspace, + aws_profile_params, + cleanup_orphaned_aws_profiles, + managed_aws_profile, + prepare_aws_credentials, + validate_aws_profile, + validate_configured_aws_profiles, +) +from openhound.scheduler.configuration import ( + ManagedConfigurationError, + _validate_document, + managed_job_configuration, +) logger = logging.getLogger(__name__) POLL_INTERVAL = 30 # seconds; fixed poll/check-in cadence, must remain below BHE's 600s client-checkin timeout +HEARTBEAT_INTERVAL = ( + 300 # seconds; BHE expires managed jobs after 15 minutes without a heartbeat +) +# TODO: Remove this constant and the guard in _send_managed_heartbeat once BHE +# supports managed collector heartbeats. This is deliberately not configurable. +_MANAGED_HEARTBEATS_DISABLED = True + +# Explicit source registration for the collectors shipped with this PoC. +MANAGED_SOURCE_MODULES = { + "faker": "openhound_faker.source", + "aws": "openhound_aws.source", +} class ExtensionNotFoundError(Exception): """Raised when the configured collector extension cannot be found.""" - pass + +class ManagedRuntimeUnavailableError(RuntimeError): + """Managed scheduling requires a configured collection runtime.""" + + +class CredentialValidationError(ValueError): + """A managed job's required credential document could not be loaded.""" + + def __init__( + self, + message: str, + summary: str = "Credential document is invalid", + **context: Any, + ): + super().__init__(message) + self.summary = summary + self.context = context + + +class ManagedJobProcessingError(RuntimeError): + """A sanitized managed polling failure that can be retried next poll.""" + + +@dataclass(frozen=True) +class PreparedManagedJob: + """Validated job and credentials for the managed collection runtime. + + Keep credentials in the in-memory worker handoff; never persist or log them. + """ + + job: CollectorJob = field(repr=False) + credentials: dict[str, Any] | AWSAccessKeyCredentials = field(repr=False) @dataclass class Result: results: dict - job_id: int + job_id: int | str + output_path: Path | None = None + data_files: list[Path] = field(default_factory=list) + uploaded_files: int = 0 + upload_failure: str | None = None -def _subprocess_collect(collector_name: str, job_id: int) -> Result: +def _subprocess_collect( + collector_name: str, + job_id: int | str, + params: dict[str, Any] | None = None, + secrets: dict[str, Any] | AWSAccessKeyCredentials | None = None, + managed: bool = False, + aws_profile_directory: Path | None = None, +) -> Result: """A subprocess which runs the DLT pipeline for the specified collector. Loads the collector by name from Python entrypoints. @@ -44,6 +135,7 @@ def _subprocess_collect(collector_name: str, job_id: int) -> Result: Args: collector_name: The name of the extension to run. job_id: The Job ID as returned by BHE. + managed: Stage OpenGraph files for upload after collection completes. Returns: Result: The result of the collection, including the collected data and the job ID. @@ -52,10 +144,96 @@ def _subprocess_collect(collector_name: str, job_id: int) -> Result: ExtensionNotFoundError: If the collector cannot be found via entrypoints. """ signal.signal(signal.SIGINT, signal.SIG_IGN) + # Context and redaction surround imports, configuration, pipeline setup, and + # failure diagnostics. The queue handler applies them before serialization. + values = ( + (secrets.access_key_id, secrets.secret_access_key) + if isinstance(secrets, AWSAccessKeyCredentials) + else tuple(_credential_log_values(secrets)) + ) + with ( + openhound_logging.worker_log_context(job_id, collector_name), + openhound_logging.redact_log_values(values), + ): + try: + return _configured_worker_collect( + collector_name, job_id, params, secrets, managed, aws_profile_directory + ) + except ManagedConfigurationError: + # AWS failures were diagnosed inside the SDK logging controls; + # configuration-resolution failures deliberately omit their values. + logger.exception("Worker configuration or collection failed") + raise + except ExtensionNotFoundError: + logger.exception("Collection worker could not load extension") + raise + except Exception: + logger.exception("Collection worker failed") + if managed: + raise ManagedConfigurationError( + "Managed collection failed; check worker logs" + ) from None + raise + + +def _credential_log_values(value: Any): + if isinstance(value, str): + yield value + elif isinstance(value, dict): + for item in value.values(): + yield from _credential_log_values(item) + elif isinstance(value, list): + for item in value: + yield from _credential_log_values(item) + + +def _configured_worker_collect( + collector_name: str, + job_id: int | str, + params: dict[str, Any] | None, + secrets: dict[str, Any] | AWSAccessKeyCredentials | None, + managed: bool, + aws_profile_directory: Path | None, +) -> Result: + if isinstance(secrets, AWSAccessKeyCredentials): + if collector_name != "aws" or aws_profile_directory is None: + raise ManagedConfigurationError( + "Managed AWS profile worker context is required" + ) + with managed_aws_profile( + params or {}, secrets, aws_profile_directory + ) as profile: + validate_aws_profile(profile) + return _collect_in_worker(collector_name, job_id, managed, True) + configured = params is not None or secrets is not None + configuration = ( + managed_job_configuration(params or {}, secrets or {}) + if configured + else nullcontext() + ) + # Install configuration before importing collectors or constructing pipelines: + # both can resolve DLT settings. Always remove it, including on worker failure. + with configuration: + if managed and collector_name == "aws": + validate_configured_aws_profiles() + return _collect_in_worker(collector_name, job_id, managed, configured) + + +def _collect_in_worker( + collector_name: str, job_id: int | str, managed: bool, configured: bool +) -> Result: available_collectors = CollectorManager.from_entrypoint() for collector in available_collectors.collectors: if collector.name == collector_name: # pyright: ignore[reportAttributeAccessIssue] + if ( + configured + and collector.dlt_source is None + and collector.name in MANAGED_SOURCE_MODULES + ): + collector.dlt_source = import_module( + MANAGED_SOURCE_MODULES[collector.name] + ).source log_fields = { "collector_extension": collector.name, "collector_extension_version": ( @@ -75,13 +253,30 @@ def _subprocess_collect(collector_name: str, job_id: int) -> Result: job_id, extra=log_fields, ) - results = dataflow.pipeline(extension=collector) + output_path = ( + Path(tempfile.mkdtemp(prefix="openhound-job-")) if managed else None + ) + try: + results = ( + dataflow.pipeline(extension=collector) + if output_path is None + else dataflow.pipeline(extension=collector, output_path=output_path) + ) + except BaseException: + if output_path is not None: + shutil.rmtree(output_path) + raise logger.info( "Collection for job %s completed successfully.", job_id, extra=log_fields, ) - return Result(results=results, job_id=job_id) + return Result( + results=results, + job_id=job_id, + output_path=output_path, + data_files=sorted(output_path.glob("*.json")) if output_path else [], + ) logger.error(f"Collector '{collector_name}' not found in available collectors.") raise ExtensionNotFoundError(f"Collector '{collector_name}' not found.") @@ -101,25 +296,56 @@ def __init__( token_id: str, collector_name: str, log_base_path: Path | None = None, + interval: int = POLL_INTERVAL, + credential_refresh: Callable[[], BHECredentials] | None = None, + *, + managed: bool | None = None, + secrets_manager: AWSSecretsManager | None = None, + managed_job_runner: Callable[[PreparedManagedJob], None] | None = None, ): # BHE client settings self.bhe_uri = bhe_uri - self.collector_name = collector_name self.client = BloodHoundEnterprise( - bhe_uri=bhe_uri, token_key=token_key, token_id=token_id + bhe_uri=bhe_uri, + token_key=token_key, + token_id=token_id, + credential_refresh=credential_refresh, ) - # Interval how often to check for a job - self.interval = POLL_INTERVAL + + # Shared scheduler settings + self.collector_name = collector_name + self.interval = interval self.log_base_path = ( log_base_path or openhound_logging.logger_override.base_path ) - # Stores the ID of currently running BHE job - self.job_running: int | None = None + # Managed collection settings and recovery state + self.managed = is_managed() if managed is None else managed + self._aws_profile_workspace: AWSProfileWorkspace | None = None + if self.managed and self.collector_name == "aws": + cleanup_orphaned_aws_profiles() + self.secrets_manager = secrets_manager or AWSSecretsManager() + self.managed_job_runner = managed_job_runner or self._run_prepared_managed_job + self._pending_managed_claim: str | None = None + self._pending_managed_start: PreparedManagedJob | None = None + self._pending_managed_failure: str | None = None + self._pending_managed_failure_summary = "Credential document is invalid" + + # Stores the ID of the currently running BHE job. + # MVP runs one job at a time. + self.job_running: int | str | None = None + self.managed_job_id: str | None = None + self.next_heartbeat_at: float | None = None # Futures/results from the subprocess executor self.future: Future[Result] | None = None - self.executor = ProcessPoolExecutor(max_workers=1, max_tasks_per_child=1) + self._worker_context = multiprocessing.get_context("spawn") + self._worker_logging = openhound_logging.WorkerLogRouter(self._worker_context) + try: + self.executor = self._create_executor() + except BaseException: + self._worker_logging.close() + raise # Exit condition, changed to True when the process needs to stop self.exit = False @@ -130,18 +356,73 @@ def _exit_handler(self, sig: int, frame): logger.warning(f"Received signal {sig}, shutting down gracefully.") def _shutdown(self) -> None: - """Shut down the executor and wait for running tasks to finish. Executed when the service is shut down.""" + """Wait for the worker and finish any pending managed uploads.""" logger.info("Collection service stopping.") - self.executor.shutdown(wait=True, cancel_futures=True) + while self.future is not None and not self.future.done(): + self._poll_running_job() + if self.future is not None: + time.sleep(self.interval) + if self.managed_job_id is not None: + self._finish_managed_job_on_shutdown() + elif self.future is not None: + try: + self._handle_completed_job(self.future) + except Exception: + logger.exception( + "Unexpected error handling completed job during shutdown." + ) + self._clear_running_job() + try: + self.executor.shutdown(wait=True, cancel_futures=True) + finally: + self._worker_logging.close() logger.info("Collection service stopped.") + def _finish_managed_job_on_shutdown(self) -> None: + """Retry a completed managed job while maintaining its claim.""" + while self.future is not None: + self._poll_running_job() + if self.future is None: + return + try: + self._handle_completed_job(self.future) + except Exception: # noqa: BLE001 + # Retain managed work and redact sensitive errors. + logger.error( + "Error reporting completed managed job during shutdown; " + "retaining result for retry." + ) + if self.future is not None: + time.sleep(self.interval) + + def _clear_running_job(self) -> None: + self._cleanup_aws_profile_workspace() + self.future = None + self.job_running = None + self.managed_job_id = None + self.next_heartbeat_at = None + + def _cleanup_aws_profile_workspace(self) -> None: + if self._aws_profile_workspace is not None: + self._aws_profile_workspace.cleanup() + self._aws_profile_workspace = None + def _reset_executor(self) -> None: """Tear down and recreate the process pool after it has entered a broken state.""" try: - self.executor.shutdown(wait=False, cancel_futures=True) + self.executor.shutdown(wait=True, cancel_futures=True) except Exception: logger.exception("Error shutting down broken executor.") - self.executor = ProcessPoolExecutor(max_workers=1, max_tasks_per_child=1) + self.executor = self._create_executor() + + def _create_executor(self) -> ProcessPoolExecutor: + return ProcessPoolExecutor( + max_workers=1, + max_tasks_per_child=1, + mp_context=self._worker_context, + initializer=openhound_logging.configure_worker_logging, + initargs=(self._worker_logging.queue, logging.getLogger().level), + ) def check_jobs(self) -> Job | None: """Checks BloodHound enterprise for available jobs. These can either be new jobs or jobs currently started and not finished/stopped. @@ -170,6 +451,183 @@ def check_jobs(self) -> Job | None: return None + def check_managed_collector_jobs(self) -> CollectorJob | None: + """Return the first available job from the managed collector queue.""" + logger.info("Checking for new managed collector jobs in BloodHound Enterprise.") + available_jobs = self.client.available_collector_jobs(self.collector_name) + for job in available_jobs.data.jobs: + if job.job_key == self.collector_name: + return job + return None + + def _report_managed_credential_failure(self, job_id: str) -> None: + # Retain the unresolved ID if reporting exhausts its retries. On the next + # poll retry end, rather than claiming or retrieving another job. + self._pending_managed_failure = job_id + self.client.end_managed_job( + job_id, + ManagedJobOutcome.FAILED, + failure_message=self._pending_managed_failure_summary, + ) + self._pending_managed_failure = None + + def _load_managed_credentials( + self, job: CollectorJob + ) -> dict[str, Any] | AWSAccessKeyCredentials: + """Retrieve and validate credentials, exposing only a safe failure message.""" + logger.debug( + "Retrieving managed collector credentials.", + extra={ + "job_id": str(job.id), + "provider": "aws.secretsmanager", + "secret_count": 1, + }, + ) + try: + if job.job_key == "aws": + credentials = prepare_aws_credentials( + job, self.secrets_manager.get_secret + ) + else: + if job.secret_key_id is None or not job.secret_key_id.strip(): + raise CredentialValidationError( + "Managed job has no credential secret identifier" + ) + credentials = self.secrets_manager.get_secret(job.secret_key_id) + if isinstance(credentials, AWSAccessKeyCredentials): + return credentials + if isinstance(credentials, dict): + _validate_document(credentials, allow_empty=True) + return credentials + except AWSAccessKeyError as error: + raise CredentialValidationError( + str(error), error.summary, **error.context + ) from None + except CredentialValidationError: + raise + except ManagedConfigurationError as error: + raise CredentialValidationError(str(error)) from None + except SecretRetrievalError as error: + raise CredentialValidationError(str(error), error.failure_summary) from None + except Exception as error: # noqa: BLE001 + # Unexpected readers can expose payloads. + raise CredentialValidationError( + f"Credential secret reader failed ({type(error).__name__})", + "Credential secret retrieval failed", + ) from None + raise CredentialValidationError("Credential document must be a JSON object") + + def _validate_managed_claim(self, job: CollectorJob, job_id: str) -> None: + """Check the job identity and lease confirmed by BHE's claim response.""" + + # The authenticated claim endpoint confirms the holder. Reclaiming does + # not renew the lease, so an old response is not enough to permit work. + lease_expires_at = job.claim_expires_at + if ( + str(job.id) != job_id + or job.job_key != self.collector_name + or job.status != "claimed" + or job.claimed_by is None + or job.claimed_at is None + or lease_expires_at is None + or lease_expires_at.utcoffset() is None + or lease_expires_at <= datetime.now(UTC) + ): + logger.error( + "Managed claim response does not confirm a usable claim.", + extra={"job_id": job_id}, + ) + raise ManagedJobProcessingError("Managed claim could not be reconciled") + + def _claim_and_validate_managed_job(self, job_id: str) -> PreparedManagedJob | None: + """Claim a managed job and validate its credentials.""" + # A timeout can happen after BHE commits the claim. Keep the ID until we + # reconcile it; discovery no longer lists jobs held by this collector. + self._pending_managed_claim = job_id + try: + job = self.client.claim_managed_job(job_id) + except BloodHoundHTTPError as error: + if error.code in (404, 409): + # An existing claimed job held by this client would return 200. + self._pending_managed_claim = None + if error.code == 409: + logger.debug( + "Managed collector claim conflict.", + extra={"job_id": job_id}, + ) + return None + raise + + self._validate_managed_claim(job, job_id) + context = {"job_id": job_id} + logger.info("Job successfully picked up", extra=context) + try: + credentials = self._load_managed_credentials(job) + except CredentialValidationError as error: + logger.error( + "Cannot prepare managed job %s: %s. Reporting this job failed.", + job_id, + str(error).rstrip("."), + extra={ + **error.context, + **context, + **( + {} if job.job_key == "aws" else {"secret_id": job.secret_key_id} + ), + }, + ) + self._pending_managed_claim = None + self._pending_managed_failure_summary = error.summary + self._report_managed_credential_failure(str(job.id)) + return None + + # Secret retrieval may have consumed the remaining lease time. + self._validate_managed_claim(job, job_id) + logger.info("Credential retrieval valid", extra=context) + return PreparedManagedJob(job=job, credentials=credentials) + + def _clear_pending_managed_start(self) -> None: + """Discard preparation state once start is confirmed or rejected.""" + self._pending_managed_claim = None + self._pending_managed_start = None + + def _require_managed_runtime(self) -> Callable[[PreparedManagedJob], None]: + if self.managed_job_runner is None: + raise ManagedRuntimeUnavailableError( + "Managed collection runtime is required for managed scheduling" + ) + return self.managed_job_runner + + def _run_prepared_managed_job(self, prepared: PreparedManagedJob) -> None: + self.run_claimed_job(prepared.job.id, prepared.job.params, prepared.credentials) + + def _poll_managed_jobs(self) -> None: + runner = self._require_managed_runtime() + try: + if self._pending_managed_failure is not None: + self._report_managed_credential_failure(self._pending_managed_failure) + return + prepared = self._pending_managed_start + if prepared is None: + job_id = self._pending_managed_claim + if job_id is None: + available_job = self.check_managed_collector_jobs() + if available_job is None: + return + job_id = str(available_job.id) + prepared = self._claim_and_validate_managed_job(job_id) + # Keep validated inputs in memory across ambiguous start failures. + # Repeating BHE's idempotent start also reconciles a job that is + # already running remotely; re-claiming it would return 409. + self._pending_managed_start = prepared + if prepared is not None: + runner(prepared) + self._clear_pending_managed_start() + except Exception: # noqa: BLE001 - redact managed runtime errors + raise ManagedJobProcessingError( + "Managed collector job processing failed" + ) from None + def check_management(self) -> ManagementOperation | None: """Return the first pending support-bundle operation, if any.""" logger.info("Checking for management operations in BloodHound Enterprise.") @@ -181,6 +639,22 @@ def check_management(self) -> ManagementOperation | None: return operation return None + def _poll_management(self) -> bool: + """Return whether a support-bundle request was handled this poll.""" + try: + operation = self.check_management() + except Exception: + logger.exception("Error checking management operations.") + return False + + if operation is None: + return False + try: + self._send_support_bundle(operation) + except Exception: + logger.exception("Error executing management operation.") + return True + def _send_support_bundle(self, operation: ManagementOperation) -> None: """Claim, upload, and complete a support-bundle operation.""" bundle_path: Path | None = None @@ -212,102 +686,325 @@ def _send_support_bundle(self, operation: ManagementOperation) -> None: ) def _start_job(self, job: Job) -> None: - """Starts a BloodHound enterprise job by ID and runs the collection process in a subprocess. The results are then used to end the job in BHE with a complete status. - this function returns no value but updates the self.futures list with the future of the subprocess and sets the currently running job ID in self.job_running + """Start an unmanaged job using the shared worker lifecycle.""" + self._launch_job(job.id) - Args: - job (Job): The job to start as a Job object. - """ + def run_claimed_job( + self, + job_id: str | UUID, + params: dict[str, Any], + secrets: dict[str, Any] | AWSAccessKeyCredentials, + ) -> None: + """Start a validated managed job; the scheduler polls it until completion.""" + # Normalize managed UUIDs once at the runtime boundary. Unmanaged IDs + # remain integers throughout the shared lifecycle. + self._launch_job(str(job_id), params=params, secrets=secrets, managed=True) - logger.info(f"Starting job {job.id} with collector '{self.collector_name}'") - self.client.start_job(job.id) - self.job_running = job.id + def _launch_job( + self, + job_id: int | str, + *, + params: dict[str, Any] | None = None, + secrets: dict[str, Any] | AWSAccessKeyCredentials | None = None, + managed: bool = False, + ) -> None: + """Notify BHE and submit one collector worker.""" + if self.job_running is not None: + raise RuntimeError(f"Job {self.job_running} is already running") + + if isinstance(secrets, AWSAccessKeyCredentials): + params = aws_profile_params( + {} if params is None else params, collector_name=self.collector_name + ) + logger.info("Starting job %s with collector '%s'.", job_id, self.collector_name) + managed_job_id: str | None = None + if managed: + if not isinstance(job_id, str): + raise TypeError("Managed job IDs must be strings") + try: + self.client.start_collector_job(job_id) + except BloodHoundHTTPError as error: + if error.code in (404, 409): + self._clear_pending_managed_start() + raise + managed_job_id = job_id + self._clear_pending_managed_start() + else: + if not isinstance(job_id, int): + raise TypeError("Unmanaged job IDs must be integers") + self.client.start_job(job_id) + + if managed: + logger.info("Managed collector job %s started successfully.", job_id) + self.job_running = job_id + self.managed_job_id = managed_job_id + self.next_heartbeat_at = ( + time.monotonic() + HEARTBEAT_INTERVAL if managed else None + ) try: + if isinstance(secrets, AWSAccessKeyCredentials): + self._cleanup_aws_profile_workspace() + self._aws_profile_workspace = AWSProfileWorkspace() self.future = self.executor.submit( - _subprocess_collect, self.collector_name, job.id + _subprocess_collect, + self.collector_name, + job_id, + dict(params or {}) if managed else None, + ( + ( + secrets + if isinstance(secrets, AWSAccessKeyCredentials) + else dict(secrets or {}) + ) + if managed + else None + ), + *((True,) if managed else ()), + *( + (self._aws_profile_workspace.directory,) + if self._aws_profile_workspace is not None + else () + ), ) except BrokenProcessPool: - logger.exception( - f"Failed to submit job {job.id}: process pool is broken, resetting." - ) - self.future = None - self.job_running = None + logger.exception("Failed to submit job %s: process pool is broken.", job_id) self._reset_executor() - self.client.end_job( - JobStatus.FAILED, - f"Failed to start collector '{self.collector_name}': worker pool was broken", + try: + self._end_job( + JobStatus.FAILED, + f"Failed to start collector '{self.collector_name}': worker pool was broken", + ) + finally: + self._clear_running_job() + except Exception as error: + logger.error( + "Failed to submit job %s.", + job_id, + extra={"job_id": job_id, "error_type": type(error).__name__}, ) + try: + self._end_job( + JobStatus.FAILED, + f"Failed to start collector '{self.collector_name}'", + ) + finally: + self._clear_running_job() + raise - def _handle_completed_job(self, future: Future[Result]) -> None: - """Handles completion of a job and reports success or failure to BloodHound Enterprise + def _end_job(self, status: JobStatus, message: str) -> None: + if self.managed_job_id is not None: + self.client.end_managed_job( + self.managed_job_id, + ( + ManagedJobOutcome.SUCCEEDED + if status is JobStatus.COMPLETE + else ManagedJobOutcome.FAILED + ), + failure_message=message if status is JobStatus.FAILED else None, + ) + else: + self.client.end_job(status, message) + + def _send_managed_heartbeat(self, job_id: str) -> None: + if _MANAGED_HEARTBEATS_DISABLED: + return - Args: - future (Future[Result]): Future containing the result of the collection subprocess, including the job ID. - """ try: - result = future.result() - logger.info(f"Job {result.job_id} completed successfully, notifying BHE.") - self.client.end_job( - JobStatus.COMPLETE, - f"Collector '{self.collector_name}' completed successfully", + self.client.heartbeat_collector_job(job_id) + except Exception as error: # noqa: BLE001 + # Transient heartbeat failures must not stop collection. + if isinstance(error, BloodHoundHTTPError) and error.code == 404: + self._abort_managed_job(job_id) + return + logger.info( + "Managed collector job %s heartbeat request failed.", + job_id, + extra={ + "job_id": job_id, + "error_type": type(error).__name__, + "status_code": ( + error.code if isinstance(error, BloodHoundHTTPError) else None + ), + }, ) + else: + logger.info("Managed collector job %s heartbeat succeeded.", job_id) + def _abort_managed_job(self, job_id: str) -> None: + """Stop local work without reporting an outcome for a claim we no longer own.""" + logger.warning( + "Managed collector job %s claim lost; stopping collection.", job_id + ) + future = self.future + processes = list((self.executor._processes or {}).values()) + try: + for process in processes: + try: + process.kill() + except OSError: + logger.exception("Unable to kill worker process.") + self.executor.shutdown(wait=True, cancel_futures=True) + except Exception: + logger.exception("Error shutting down executor after claim loss.") + finally: + # Completed workers can leave staged output awaiting upload. Clean + # it before discarding the result, regardless of how claim loss arose. + if ( + future is not None + and future.done() + and not future.cancelled() + and future.exception() is None + ): + self._remove_job_output(future.result()) + self._clear_running_job() + self.executor = self._create_executor() + + def _handle_completed_job(self, future: Future[Result]) -> None: + """Report the completed worker's outcome through the active job API.""" + # Erase credentials before uploads/end retries. This also runs when a + # killed worker could not execute its context manager's finally block. + self._cleanup_aws_profile_workspace() + result: Result | None = None + try: + result = future.result() + logger.info(f"Job {result.job_id} completed successfully, notifying BHE.") except ExtensionNotFoundError: logger.error( f"Collector '{self.collector_name}' not found. Marking job as failed." ) - self.client.end_job( - JobStatus.FAILED, - f"Collector '{self.collector_name}' not found", + status = JobStatus.FAILED + message = f"Collector '{self.collector_name}' not found" + except ManagedConfigurationError as error: + status = JobStatus.FAILED + message = str(error) + logger.error( + "Managed job %s configuration is invalid: %s", + self.job_running, + message, + extra={ + "job_id": self.job_running, + "collector_extension": self.collector_name, + }, ) - except BrokenProcessPool: logger.exception( "Collection worker was terminated abruptly; resetting process pool." ) self._reset_executor() - self.client.end_job( - JobStatus.FAILED, - f"Collection worker for '{self.collector_name}' was terminated abruptly", + status = JobStatus.FAILED + message = ( + f"Collection worker for '{self.collector_name}' was terminated abruptly" ) - - except Exception: - logger.exception("Collection subprocess failed.") - self.client.end_job( - JobStatus.FAILED, - f"Unexpected error while running '{self.collector_name}' collector", + except Exception as error: # noqa: BLE001 - any worker error must end the job + logger.error( + "Collection subprocess failed.", + extra={"error_type": type(error).__name__}, + ) + status = JobStatus.FAILED + message = ( + f"Unexpected error while running '{self.collector_name}' collector" ) + else: + if self.managed_job_id is not None and result.output_path is not None: + if result.job_id != self.job_running: + raise ValueError( + "Managed collection result does not match the running job" + ) + while result.upload_failure is None and result.uploaded_files < len( + result.data_files + ): + data_file = result.data_files[result.uploaded_files] + try: + data = data_file.open("rb") + except FileNotFoundError: + result.upload_failure = ( + "Collector data upload failed: staged file is missing" + ) + break + try: + with data: + self.client.upload_collector_job_data( + self.managed_job_id, data + ) + except BloodHoundHTTPError as error: + if error.code in (404, 409): + self._abort_managed_job(self.managed_job_id) + return + if error.code not in (400, 401, 403, 415): + raise + # Retain the failure across polls if reporting it fails. + # Response bodies can contain sensitive collected data. + result.upload_failure = ( + f"Collector data upload rejected (HTTP {error.code})" + ) + break + result.uploaded_files += 1 + if result.upload_failure is not None: + status = JobStatus.FAILED + message = result.upload_failure + else: + status = JobStatus.COMPLETE + message = f"Collector '{self.collector_name}' completed successfully" + self._end_job(status, message) + if self.managed_job_id is not None and result is not None: + self._remove_job_output(result) + self._clear_running_job() - finally: - self.future = None - self.job_running = None + def _remove_job_output(self, result: Result) -> None: + if result.output_path is not None: + try: + shutil.rmtree(result.output_path) + except OSError: + logger.exception("Unable to remove output for job %s.", result.job_id) + + def _poll_running_job(self) -> None: + if self.managed_job_id is not None: + if ( + self.next_heartbeat_at is not None + and time.monotonic() >= self.next_heartbeat_at + ): + self._send_managed_heartbeat(self.managed_job_id) + if self.managed_job_id is not None: + self.next_heartbeat_at = time.monotonic() + HEARTBEAT_INTERVAL + else: + try: + _ = self.client.jobs_current + except Exception: + logger.exception("Error checking in-progress job.") def _poll(self) -> None: """Checks if jobs are completed and if a job should be run.""" - # If a job is currently running check if the future is done and handle completion try: if self.future is not None and self.future.done(): self._handle_completed_job(self.future) except Exception: - logger.exception("Unexpected error handling completed job.") - self.future = None - self.job_running = None + if self.managed: + logger.error( + "Error reporting completed managed job; retaining result for retry." + ) + else: + logger.exception( + "Error reporting completed job; retaining result for retry." + ) - # Management operations have priority over new collections while idle. - if self.job_running is None: - try: - operation = self.check_management() - except Exception: - logger.exception("Error checking management operations.") - operation = None + # Legacy support-bundle management is only available to unmanaged clients. + idle = ( + self.job_running is None + and self._pending_managed_claim is None + and self._pending_managed_start is None + and self._pending_managed_failure is None + ) + if not self.managed and idle and self._poll_management(): + return - if operation is not None: - try: - self._send_support_bundle(operation) - except Exception: - logger.exception("Error executing management operation.") - return + if self.managed: + if self.job_running is None: + self._poll_managed_jobs() + else: + self._poll_running_job() + return + if self.job_running is None: try: available_job = self.check_jobs() if available_job: @@ -315,10 +1012,7 @@ def _poll(self) -> None: except Exception: logger.exception("Error checking for or starting jobs.") else: - try: - self.client.jobs_current - except Exception: - logger.exception("Error checking in-progress job.") + self._poll_running_job() def start(self) -> None: """Start method to initiate the process of checking for jobs and running them. This method will run indefinitely until an exit signal is received""" @@ -327,15 +1021,29 @@ def start(self) -> None: logger.info( f"Service started, monitoring {self.bhe_uri} every {self.interval} seconds." ) - try: - self.client.update_client_metadata() - - except Exception as err: - logger.exception(f"Unable to update client metadata: {err}") + # Managed identities have no client metadata endpoint. + if not self.managed: + try: + self.client.update_client_metadata() + except Exception: + logger.exception("Unable to update client metadata.") try: while not self.exit: - self._poll() + if self.managed: + try: + self._poll() + except ManagedRuntimeUnavailableError: + logger.error( + "Managed collection runtime is unavailable; retrying next poll." + ) + except ManagedJobProcessingError: + logger.error( + "Managed scheduler poll failed; retrying next poll in %s seconds.", + self.interval, + ) + else: + self._poll() time.sleep(self.interval) finally: self._shutdown() diff --git a/src/scheduler.py b/src/scheduler.py index d18fa16e..5fc388c5 100644 --- a/src/scheduler.py +++ b/src/scheduler.py @@ -3,6 +3,8 @@ import dlt import typer +from openhound import config +from openhound.core.clients.bhe_credentials import AWSBHECredentials from openhound.scheduler.service import Service logger = logging.getLogger(__name__) @@ -11,14 +13,26 @@ def start(): - # Load BHE config and secrets bhe_uri = dlt.config["destination.bloodhoundenterprise.url"] collector_name = dlt.config["destination.bloodhoundenterprise.collector_name"] - # Load BHE secrets - token_key = dlt.secrets["destination.bloodhoundenterprise.token_key"] - token_id = dlt.secrets["destination.bloodhoundenterprise.token_id"] + if config.is_managed(): + managed_config = config.get_managed_bhe_config() + credentials = AWSBHECredentials(managed_config.secret_name) + current_credentials = credentials.poll_until_available( + managed_config.poll_interval + ) + credential_refresh = credentials.refresh + token_key = current_credentials.token_key + token_id = current_credentials.token_id + interval = managed_config.poll_interval + else: + # Load BHE secrets + token_key = dlt.secrets["destination.bloodhoundenterprise.token_key"] + token_id = dlt.secrets["destination.bloodhoundenterprise.token_id"] + interval = 30 + credential_refresh = None # Start the service logger.info(f"Initializing service for collector '{collector_name}'.") @@ -27,6 +41,8 @@ def start(): token_key=token_key, token_id=token_id, collector_name=collector_name, + interval=interval, + credential_refresh=credential_refresh, ) svc.start() diff --git a/tests/test_aws_profile_adapter.py b/tests/test_aws_profile_adapter.py new file mode 100644 index 00000000..aa144071 --- /dev/null +++ b/tests/test_aws_profile_adapter.py @@ -0,0 +1,1142 @@ +import logging +import multiprocessing +import os +import stat +import tempfile +import time +import traceback +import zipfile +from concurrent.futures import Future, ProcessPoolExecutor +from configparser import ConfigParser +from copy import deepcopy +from datetime import UTC, datetime, timedelta +from types import SimpleNamespace + +import boto3 +import dlt +import pytest +from botocore.exceptions import ClientError, ParamValidationError +from botocore.httpsession import URLLib3Session +from botocore.session import Session as BotocoreSession +from botocore.stub import Stubber +from dlt.common.configuration.container import Container +from dlt.common.configuration.providers.doc import CustomLoaderDocProvider +from dlt.common.configuration.providers.provider import ConfigProvider +from dlt.common.configuration.specs.pluggable_run_context import PluggableRunContext +from dlt.common.configuration.utils import TraceLogContext +from dlt.destinations import filesystem + +from openhound.core import logging as openhound_logging +from openhound.core.clients.bloodhound_enterprise import ManagedJobOutcome +from openhound.core.resources import safe_defer_wrapper, safe_resource_wrapper +from openhound.core.support_bundle import create_support_bundle +from openhound.scheduler import aws_profile, service +from openhound.scheduler.aws_profile import ( + AWSAccessKeyCredentials, + AWSProfileWorkspace, + cleanup_orphaned_aws_profiles, + managed_aws_profile, +) +from openhound.scheduler.configuration import ManagedConfigurationError + +ACCESS_KEY = "AKIAIOSFODNN7EXAMPLE" +SECRET_KEY = "wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY" +PAIR = AWSAccessKeyCredentials(ACCESS_KEY, SECRET_KEY) +SOURCE_ACCOUNT = "999999999999" +MANAGEMENT_ACCOUNT = "111111111111" +ROLE_ACCOUNT = "333333333333" +ROLE_ARN = f"arn:aws:iam::{ROLE_ACCOUNT}:role/security/CollectorRole" +ROLE_ACCESS_KEY = "ASIAIOSFODNN7EXAMPLE" +ROLE_SECRET_KEY = "T" * 40 +ROLE_TOKEN = "assumed-session-token" + + +@pytest.fixture(autouse=True) +def private_workspaces(tmp_path, monkeypatch): + monkeypatch.setattr( + aws_profile, + "tempfile", + SimpleNamespace(gettempdir=lambda: str(tmp_path), mkdtemp=tempfile.mkdtemp), + ) + monkeypatch.setenv("AWS_EC2_METADATA_DISABLED", "true") + monkeypatch.setenv( + "AWS_SHARED_CREDENTIALS_FILE", str(tmp_path / "scheduler-credentials") + ) + monkeypatch.setenv("AWS_CONFIG_FILE", str(tmp_path / "scheduler-config")) + monkeypatch.setattr( + URLLib3Session, + "send", + lambda *args, **kwargs: pytest.fail("Live AWS requests are forbidden"), + ) + + +@pytest.fixture +def workspace(): + workspace = AWSProfileWorkspace() + try: + yield workspace + finally: + workspace.cleanup() + + +@pytest.fixture +def profile_assumptions(monkeypatch): + """Stub the STS clients botocore creates internally for role profiles.""" + create_client = BotocoreSession.create_client + pending, stubbers = [], [] + + def expect(params, *, denied=False, secret=ROLE_SECRET_KEY): + pending.append((params, denied, secret)) + + def stub_sts(session, name, *args, **kwargs): + client = create_client(session, name, *args, **kwargs) + if name == "sts" and kwargs.get("aws_access_key_id") == ACCESS_KEY: + frozen = client._request_signer._credentials.get_frozen_credentials() + assert (frozen.access_key, frozen.secret_key, frozen.token) == ( + ACCESS_KEY, + SECRET_KEY, + None, + ) + assert pending, "Unexpected initial role assumption" + params, denied, secret = pending.pop(0) + stubber = Stubber(client) + if denied: + stubber.add_client_error( + "assume_role", + service_error_code="AccessDenied", + service_message="Role trust denied", + expected_params=params, + ) + else: + stubber.add_response( + "assume_role", + { + "Credentials": { + "AccessKeyId": ROLE_ACCESS_KEY, + "SecretAccessKey": secret, + "SessionToken": ROLE_TOKEN, + "Expiration": datetime.now(UTC) + timedelta(hours=1), + } + }, + params, + ) + stubber.activate() + stubbers.append(stubber) + return client + + monkeypatch.setattr(BotocoreSession, "create_client", stub_sts) + try: + yield expect + assert not pending, "Expected initial role assumption did not occur" + for stubber in stubbers: + stubber.assert_no_pending_responses() + finally: + for stubber in stubbers: + stubber.deactivate() + + +def _identity(account=SOURCE_ACCOUNT): + return { + "Account": account, + "Arn": f"arn:aws:iam::{account}:user/collector", + "UserId": "collector", + } + + +def _assert_source_credentials(session): + credentials = session.get_credentials() + assert credentials.method == "shared-credentials-file" + frozen = credentials.get_frozen_credentials() + assert (frozen.access_key, frozen.secret_key, frozen.token) == ( + ACCESS_KEY, + SECRET_KEY, + None, + ) + + +def _assert_role_credentials(session): + credentials = session.get_credentials() + assert credentials.method == "assume-role" + frozen = credentials.get_frozen_credentials() + assert (frozen.access_key, frozen.secret_key, frozen.token) == ( + ROLE_ACCESS_KEY, + ROLE_SECRET_KEY, + ROLE_TOKEN, + ) + + +@pytest.mark.parametrize("custom_session", [False, True], ids=["default", "custom"]) +def test_role_profile_resolves_through_real_sdk( + custom_session, + workspace, + profile_assumptions, + monkeypatch, +): + monkeypatch.setenv("AWS_ACCESS_KEY_ID", "scheduler-key") + monkeypatch.setenv("AWS_SECRET_ACCESS_KEY", "scheduler-secret") + monkeypatch.setenv("AWS_SESSION_TOKEN", "scheduler-token") + monkeypatch.setenv("AWS_PROFILE", "scheduler-profile") + monkeypatch.setenv("AWS_ROLE_ARN", "arn:aws:iam::123456789012:role/OtherRole") + options = {"role_arn": ROLE_ARN, "region_name": "eu-west-1"} + if custom_session: + options.update( + role_session_name="collection-session", external_id="customer/123" + ) + params = {"sources": {"source": {"aws": {"credentials": options}}}} + original_params, original_environment = deepcopy(params), dict(os.environ) + with managed_aws_profile(params, PAIR, workspace.directory) as profile: + session_name = "collection-session" if custom_session else profile + expected = {"RoleArn": ROLE_ARN, "RoleSessionName": session_name} + if custom_session: + expected["ExternalId"] = "customer/123" + profile_assumptions(expected) + settings = dlt.secrets.get("sources.source.aws.credentials", dict) + assert settings["profile_name"] == profile + assert settings["region_name"] == "eu-west-1" + session = boto3.Session( + profile_name=settings["profile_name"], region_name=settings["region_name"] + ) + _assert_role_credentials(session) + sts = session.client("sts") + with Stubber(sts) as stubber: + identity = { + "Account": ROLE_ACCOUNT, + "Arn": f"arn:aws:sts::{ROLE_ACCOUNT}:assumed-role/CollectorRole/{session_name}", + "UserId": f"role:{session_name}", + } + stubber.add_response("get_caller_identity", identity, {}) + assert sts.get_caller_identity() == identity + iam = session.client("iam") + assert iam._request_signer._credentials is session.get_credentials() + for key in ("role_arn", "role_session_name", "external_id"): + assert dlt.secrets.get( + f"sources.source.aws.credentials.{key}", str + ) == options.get(key) + config, credentials = ConfigParser(), ConfigParser() + config.read(workspace.directory / "config") + credentials.read(workspace.directory / "credentials") + role = config[f"profile {profile}"] + assert role["role_arn"] == ROLE_ARN + assert role["role_session_name"] == session_name + assert credentials.sections() == [role["source_profile"]] + assert credentials[role["source_profile"]]["aws_access_key_id"] == ACCESS_KEY + assert profile not in credentials + for name in ("config", "credentials"): + assert stat.S_IMODE((workspace.directory / name).stat().st_mode) == 0o600 + assert Container()[TraceLogContext].all_traces == [] + assert params == original_params + assert dict(os.environ) == original_environment + assert not (workspace.directory / "config").exists() + assert not (workspace.directory / "credentials").exists() + + +def test_role_profile_refreshes_temporary_credentials(workspace, profile_assumptions): + params = {"sources": {"source": {"aws": {"credentials": {"role_arn": ROLE_ARN}}}}} + with managed_aws_profile(params, PAIR, workspace.directory) as profile: + expected = {"RoleArn": ROLE_ARN, "RoleSessionName": profile} + profile_assumptions(expected) + session = boto3.Session(profile_name=profile, region_name="us-east-1") + _assert_role_credentials(session) + profile_assumptions(expected, secret="R" * 40) + provider = session._session.get_component("credential_provider").get_provider( + "assume-role" + ) + provider.cache.clear() + credentials = session.get_credentials() + credentials._expiry_time = datetime.now(UTC) - timedelta(seconds=1) + assert credentials.get_frozen_credentials().secret_key == "R" * 40 + + +def test_denied_profile_role_fails_without_source_fallback( + workspace, + profile_assumptions, + caplog, +): + original_environment = dict(os.environ) + params = {"sources": {"source": {"aws": {"credentials": {"role_arn": ROLE_ARN}}}}} + with ( + pytest.raises(ManagedConfigurationError, match="Managed AWS collection failed"), + managed_aws_profile(params, PAIR, workspace.directory) as profile, + ): + profile_assumptions( + {"RoleArn": ROLE_ARN, "RoleSessionName": profile}, denied=True + ) + session = boto3.Session(profile_name=profile, region_name="us-east-1") + session.get_credentials().get_frozen_credentials() + pytest.fail("Denied role must not fall back to source credentials") + assert "AssumeRole returned AccessDenied" in caplog.text + assert ACCESS_KEY not in caplog.text and SECRET_KEY not in caplog.text + assert dict(os.environ) == original_environment + assert not (workspace.directory / "config").exists() + assert not (workspace.directory / "credentials").exists() + + +@pytest.mark.parametrize( + "outcome", + [ + "assume-role-denied", + "invalid-source-key", + "success", + "recoverable", + "empty", + ], +) +def test_initial_identity_controls_worker_and_managed_outcome( + outcome, profile_assumptions, monkeypatch, tmp_path, caplog +): + """Real SDK profile resolution + DLT + worker; mock only AWS/BHE IO and executor.""" + failed = outcome in {"assume-role-denied", "invalid-source-key"} + role = outcome != "invalid-source-key" + options = { + "region_name": "eu-west-1", + "endpoint_url": "https://sts.example.test", + } + if role: + options.update(role_arn=ROLE_ARN, role_session_name="initial-collection") + profile_assumptions( + {"RoleArn": ROLE_ARN, "RoleSessionName": "initial-collection"}, + denied=outcome == "assume-role-denied", + ) + params = {"sources": {"source": {"aws": {"credentials": options}}}} + # Defaults must not become an alternative identity after a role denial. + monkeypatch.setenv("AWS_ACCESS_KEY_ID", "scheduler-key") + monkeypatch.setenv("AWS_SECRET_ACCESS_KEY", "scheduler-secret") + monkeypatch.setenv("AWS_SESSION_TOKEN", "scheduler-token") + original_environment = dict(os.environ) + providers = Container()[PluggableRunContext].providers + original_providers = providers.providers + create_client = BotocoreSession.create_client + identity_requests, stubbers, stages, ends, uploads = [], [], [], [], [] + + def stub_identity(session, name, *args, **kwargs): + client = create_client(session, name, *args, **kwargs) + if name == "sts" and kwargs.get("aws_access_key_id") is None: + # Stubber bypasses request signing. Resolve the real SDK credential + # chain first so the initial AssumeRole still runs (and can fail). + credentials = client._request_signer._credentials + frozen = credentials.get_frozen_credentials() + assert credentials.method == ( + "assume-role" if role else "shared-credentials-file" + ) + assert frozen.access_key == (ROLE_ACCESS_KEY if role else ACCESS_KEY) + assert client.meta.region_name == "eu-west-1" + assert client.meta.endpoint_url == "https://sts.example.test" + identity_requests.append(frozen.access_key) + stubber = Stubber(client) + if outcome == "invalid-source-key": + stubber.add_client_error( + "get_caller_identity", + service_error_code="InvalidClientTokenId", + expected_params={}, + ) + else: + stubber.add_response("get_caller_identity", _identity(ROLE_ACCOUNT), {}) + stubber.activate() + stubbers.append(stubber) + return client + + monkeypatch.setattr(BotocoreSession, "create_client", stub_identity) + + def denied_resource(): + raise ClientError( + {"Error": {"Code": "AccessDenied", "Message": "Optional bucket denied"}}, + "GetBucketPolicy", + ) + yield # Resource errors occur on iteration, as in the AWS extension. + + def denied_defer(): + raise ClientError( + {"Error": {"Code": "AccessDenied", "Message": "Optional object denied"}}, + "GetObject", + ) + + def collect(**kwargs): + stages.append("collect") + settings = dlt.secrets.get("sources.source.aws.credentials", dict) + assert settings["profile_name"].startswith("openhound-job-") + assert (directory / "credentials").exists() + resources = [dlt.resource([], name="accounts")] + if outcome != "empty": + resources = [dlt.resource([{"AccountId": ROLE_ACCOUNT}], name="accounts")] + if outcome == "recoverable": + resources.append( + dlt.resource(safe_resource_wrapper(denied_resource, "bucket_policy"))() + ) + + @dlt.resource + def objects(): + yield dlt.defer(safe_defer_wrapper(denied_defer))() + + resources.append(objects()) + pipeline = dlt.pipeline( + pipeline_name="initial_auth_regression", + pipelines_dir=str(tmp_path / "pipelines"), + destination=filesystem(bucket_url=str(tmp_path / "collected")), + dataset_name="aws", + ) + return pipeline.run(resources) + + def preprocess(extension): + stages.append("preprocess") + return [] + + def convert(extension, output_path): + stages.append("convert") + (output_path / "graph.json").write_text('{"graph": {}}') + return [] + + collector = SimpleNamespace( + name="aws", + dlt_source=object(), + package_version="test", + metadata=None, + collector=collect, + preprocessor=object(), + ) + monkeypatch.setattr( + service.CollectorManager, + "from_entrypoint", + lambda: SimpleNamespace(collectors=[collector]), + ) + monkeypatch.setattr(service.dataflow, "preprocess", preprocess) + monkeypatch.setattr(service.dataflow, "convert", convert) + output_directory = tmp_path / "staged-graph" + + def create_output(**kwargs): + output_directory.mkdir() + return str(output_directory) + + monkeypatch.setattr(service, "tempfile", SimpleNamespace(mkdtemp=create_output)) + scheduler = service.Service( + "https://bhe.example.test", "token", "id", "aws", managed=True + ) + monkeypatch.setattr(scheduler.client, "start_collector_job", lambda job_id: None) + monkeypatch.setattr( + scheduler.client, + "end_managed_job", + lambda *args, **kwargs: ends.append((args, kwargs)), + ) + monkeypatch.setattr( + scheduler.client, + "upload_collector_job_data", + lambda *args: uploads.append(args), + ) + + def submit(function, *args): + nonlocal directory + directory = scheduler._aws_profile_workspace.directory + future = Future() + try: + future.set_result(function(*args)) + except Exception as error: # noqa: BLE001 - emulate ProcessPoolExecutor + future.set_exception(error) + return future + + directory = None + monkeypatch.setattr(scheduler.executor, "submit", submit) + try: + scheduler.run_claimed_job("job-initial-auth", params, PAIR) + assert not (directory / "credentials").exists() + assert not (directory / "config").exists() + assert dict(os.environ) == original_environment + assert providers.providers is original_providers + if failed: + with pytest.raises( + ManagedConfigurationError, match="Managed AWS collection failed" + ): + scheduler.future.result() + assert stages == [] + assert "completed successfully" not in caplog.text + assert ( + "AssumeRole returned AccessDenied" + if role + else "GetCallerIdentity returned InvalidClientTokenId" + ) in caplog.text + else: + result = scheduler.future.result() + assert result.job_id == "job-initial-auth" + assert stages == ["collect", "preprocess", "convert"] + if outcome == "recoverable": + assert ( + "Error in resource 'bucket_policy' during iteration" in caplog.text + ) + assert "Error executing DLT defer" in caplog.text + scheduler._handle_completed_job(scheduler.future) + assert ends[0][0] == ( + "job-initial-auth", + ManagedJobOutcome.FAILED if failed else ManagedJobOutcome.SUCCEEDED, + ) + assert len(ends) == 1 + assert len(uploads) == (0 if failed else 1) + assert identity_requests == ( + [] + if outcome == "assume-role-denied" + else [ROLE_ACCESS_KEY if role else ACCESS_KEY] + ) + assert not directory.exists() + assert not output_directory.exists() + assert scheduler._aws_profile_workspace is None + assert scheduler.future is scheduler.job_running is None + assert ACCESS_KEY not in caplog.text and SECRET_KEY not in caplog.text + for stubber in stubbers: + stubber.assert_no_pending_responses() + finally: + for stubber in stubbers: + stubber.deactivate() + scheduler.executor.shutdown(wait=True, cancel_futures=True) + scheduler._worker_logging.close() + scheduler._cleanup_aws_profile_workspace() + + +def test_boto3_file_provider_uses_exact_pair_without_a_session_token( + workspace, monkeypatch +): + monkeypatch.setenv("AWS_ACCESS_KEY_ID", "runtime-access") + monkeypatch.setenv("AWS_SECRET_ACCESS_KEY", "runtime-secret") + monkeypatch.setenv("AWS_SESSION_TOKEN", "runtime-token") + with managed_aws_profile({}, PAIR, workspace.directory) as profile: + assert ( + dlt.secrets.get("sources.source.aws.credentials.profile_name", str) + == profile + ) + session = boto3.Session(profile_name=profile) + _assert_source_credentials(session) + sts = session.client("sts", region_name="us-east-1") + with Stubber(sts) as stubber: + stubber.add_response("get_caller_identity", _identity(), {}) + assert sts.get_caller_identity()["Account"] == SOURCE_ACCOUNT + + +def test_arbitrary_settings_survive_profile_adaptation_without_mutation(workspace): + options = { + "profile_name": "replaced-at-runtime", + "profile_names": ["other-profile"], + "auth_mode": "unknown-mode", + "source_account": "incomplete-organization", + "future": {"values": [None, {}, [], False, 0]}, + "role_arn": "not-an-arn", + "role_session_name": "x", + "external_id": "contains spaces", + } + params = { + "sources": { + "source": {"aws": {"credentials": options, "future": options["future"]}} + }, + "pipeline": {"credentials": {"unknown": None}}, + } + original = deepcopy(params) + pair = AWSAccessKeyCredentials("nonstandard-key", "short%secret") + with managed_aws_profile(params, pair, workspace.directory) as profile: + document = dlt.config.get("sources.source.aws", dict) + assert document == { + "credentials": {**options, "profile_name": profile}, + "future": options["future"], + } + assert dlt.config.get("pipeline.credentials", dict) == {"unknown": None} + config = ConfigParser(interpolation=None) + config.read(workspace.directory / "config") + role = config[f"profile {profile}"] + assert role["role_arn"] == "not-an-arn" + assert role["role_session_name"] == "x" + assert role["external_id"] == "contains spaces" + keys = ConfigParser(interpolation=None) + keys.read(workspace.directory / "credentials") + assert dict(keys[role["source_profile"]]) == { + "aws_access_key_id": "nonstandard-key", + "aws_secret_access_key": "short%secret", + } + assert params == original + + +@pytest.mark.parametrize( + "settings", + [ + {"role_arn": ROLE_ARN, "role_session_name": "collection-session"}, + {"role_arn": ROLE_ARN, "external_id": 123}, + {"role_arn": "user/not-a-role", "role_session_name": "x" * 65}, + ], +) +def test_role_fields_are_mechanically_serialized(workspace, settings): + params = {"sources": {"source": {"aws": {"credentials": settings}}}} + with managed_aws_profile(params, PAIR, workspace.directory) as profile: + config = ConfigParser(interpolation=None) + config.read(workspace.directory / "config") + for key, value in settings.items(): + assert config[f"profile {profile}"][key] == str(value) + assert config[f"profile {profile}"]["source_profile"] == f"{profile}-source" + + +@pytest.mark.parametrize("value", [None, {}, [], "unencodable-\ud800"]) +def test_unserializable_ini_values_fail_operationally_and_cleanup(workspace, value): + params = { + "sources": { + "source": { + "aws": {"credentials": {"role_arn": ROLE_ARN, "external_id": value}} + } + } + } + original = dict(os.environ) + with ( + pytest.raises( + ManagedConfigurationError, match="could not be serialized" + ) as error, + managed_aws_profile(params, PAIR, workspace.directory), + ): + pytest.fail("unserializable INI cannot be installed") + assert "None" not in str(error.value) + assert dict(os.environ) == original + assert not (workspace.directory / "credentials").exists() + + +@pytest.mark.parametrize("role_profile", [False, True], ids=["raw-user", "role"]) +def test_partial_profile_write_failure_restores_worker_state( + role_profile, workspace, monkeypatch, caplog +): + original_environment = dict(os.environ) + providers = Container()[PluggableRunContext].providers + original_providers = providers.providers + write_ini = aws_profile._write_ini + + def fail_config_write(path, sections): + if path.name == "config": + raise OSError(SECRET_KEY) + write_ini(path, sections) + + monkeypatch.setattr(aws_profile, "_write_ini", fail_config_write) + params = ( + {"sources": {"source": {"aws": {"credentials": {"role_arn": ROLE_ARN}}}}} + if role_profile + else {} + ) + with ( + pytest.raises(ManagedConfigurationError) as error, + managed_aws_profile(params, PAIR, workspace.directory), + ): + pytest.fail("partial profile cannot run collection") + assert SECRET_KEY not in "".join(traceback.format_exception(error.value)) + assert "OSError: [REDACTED]" in caplog.text + assert "fail_config_write" in caplog.text + assert SECRET_KEY not in caplog.text + assert not (workspace.directory / "credentials").exists() + assert dict(os.environ) == original_environment + assert providers.providers is original_providers + + +@pytest.mark.parametrize("wrapped", [False, True], ids=["direct", "wrapped"]) +@pytest.mark.parametrize("sdk_error", [False, True], ids=["client", "botocore"]) +def test_unlogged_sdk_failures_keep_safe_diagnostics( + workspace, caplog, wrapped, sdk_error +): + temporary_secret = "temporary-secret-not-in-raw-pair" + payload = f"{ACCESS_KEY} {SECRET_KEY} {temporary_secret}" + failure = ( + ParamValidationError(report=payload) + if sdk_error + else ClientError( + {"Error": {"Code": "AccessDenied", "Message": payload}}, + "GetCallerIdentity", + ) + ) + with ( + pytest.raises(ManagedConfigurationError) as error, + caplog.at_level(logging.INFO), + managed_aws_profile({}, PAIR, workspace.directory), + ): + try: + raise failure + except Exception as cause: + if wrapped: + raise RuntimeError(f"Collector failed: {payload}") from cause + raise + + record = next( + record + for record in caplog.records + if "Managed AWS collection failed" in record.getMessage() + ) + assert record.levelno == logging.ERROR + assert record.sdk_exception_type == type(failure).__name__ + assert "test_unlogged_sdk_failures_keep_safe_diagnostics" in record.getMessage() + if sdk_error: + assert "parameter validation failed" in record.getMessage() + else: + assert "GetCallerIdentity returned AccessDenied" in record.getMessage() + assert "resource permissions" in record.getMessage() + if wrapped: + assert "RuntimeError" in record.getMessage() + for value in (ACCESS_KEY, SECRET_KEY, temporary_secret): + assert value not in caplog.text + assert all(value not in repr(record.__dict__) for record in caplog.records) + assert value not in "".join(traceback.format_exception(error.value)) + assert str(error.value) == "Managed AWS collection failed; check worker logs" + assert not (workspace.directory / "credentials").exists() + + +def test_unlogged_import_failure_keeps_redacted_traceback(workspace, caplog): + with ( + pytest.raises(ManagedConfigurationError), + managed_aws_profile({}, PAIR, workspace.directory), + ): + raise ModuleNotFoundError(f"Missing extension for {ACCESS_KEY}: {SECRET_KEY}") + + assert "ModuleNotFoundError: Missing extension for [REDACTED]: [REDACTED]" in ( + caplog.text + ) + assert "Traceback (most recent call last)" in caplog.text + assert ACCESS_KEY not in caplog.text + assert SECRET_KEY not in caplog.text + + +def test_sdk_debug_responses_are_suppressed_and_logging_state_restored( + workspace, caplog +): + sdk_logger = logging.getLogger("botocore.managed_profile_test") + previous = ( + sdk_logger.level, + sdk_logger.handlers[:], + sdk_logger.disabled, + sdk_logger.propagate, + ) + sdk_logger.disabled = False + sdk_logger.propagate = True + sdk_logger.setLevel(logging.DEBUG) + try: + with ( + caplog.at_level(logging.DEBUG), + managed_aws_profile({}, PAIR, workspace.directory), + ): + sdk_logger.debug( + "AssumeRole SecretAccessKey=%s", "temporary-secret-not-in-raw-pair" + ) + logging.getLogger("botocore.new_managed_profile_test").debug( + "Response body %s", SECRET_KEY + ) + assert "temporary-secret-not-in-raw-pair" not in caplog.text + assert SECRET_KEY not in caplog.text + assert sdk_logger.disabled is False + assert sdk_logger.level == logging.DEBUG + finally: + sdk_logger.setLevel(previous[0]) + sdk_logger.handlers, sdk_logger.disabled, sdk_logger.propagate = previous[1:] + + +def test_profile_resolves_in_sdk_and_preserves_nested_configuration(workspace): + params = { + "sources": { + "source": { + "aws": { + "credentials": { + "region_name": "eu-west-1", + "endpoint_url": "https://aws.example.test", + }, + "max_attempts": 7, + "bucket_policy": { + "include": True, + "max_objects": 9, + "rule": [{"include": True, "match_bucket": "public/.*"}], + }, + } + } + } + } + original = deepcopy(params) + with managed_aws_profile(params, PAIR, workspace.directory) as profile: + settings = dlt.config.get("sources.source.aws", dict) + credentials = settings["credentials"] + assert credentials == { + **params["sources"]["source"]["aws"]["credentials"], + "profile_name": profile, + } + session = boto3.Session( + profile_name=credentials["profile_name"], + region_name=credentials["region_name"], + ) + _assert_source_credentials(session) + client = session.client("sts", endpoint_url=credentials["endpoint_url"]) + with Stubber(client) as stubber: + stubber.add_response("get_caller_identity", _identity(), {}) + assert client.get_caller_identity()["Account"] == SOURCE_ACCOUNT + assert client.meta.endpoint_url == "https://aws.example.test" + assert settings["max_attempts"] == 7 + assert ( + settings["bucket_policy"] + == params["sources"]["source"]["aws"]["bucket_policy"] + ) + for name in ("credentials", "config"): + path = workspace.directory / name + assert stat.S_IMODE(path.stat().st_mode) == 0o600 + assert stat.S_IMODE(workspace.directory.stat().st_mode) == 0o700 + assert ( + "aws_session_token" not in (workspace.directory / "credentials").read_text() + ) + assert Container()[TraceLogContext].all_traces == [] + assert params == original + assert not (workspace.directory / "credentials").exists() + assert not (workspace.directory / "config").exists() + + +def test_runtime_and_local_configuration_cannot_override_generated_identity( + workspace, monkeypatch +): + for key, value in { + "AWS_ACCESS_KEY_ID": "scheduler-access-key", + "AWS_SECRET_ACCESS_KEY": "scheduler-secret", + "AWS_SESSION_TOKEN": "scheduler-token", + "AWS_SECURITY_TOKEN": "legacy-token", + "AWS_PROFILE": "local-profile", + "AWS_DEFAULT_PROFILE": "local-profile", + "AWS_ROLE_ARN": "arn:aws:iam::123456789012:role/workload", + "AWS_WEB_IDENTITY_TOKEN_FILE": "/must/not/be/read", + "AWS_CONTAINER_CREDENTIALS_FULL_URI": "http://127.0.0.1:1/credentials", + "SOURCES__SOURCE__AWS__CREDENTIALS__PROFILE_NAME": "local-profile", + "SOURCES__SOURCE__AWS__CREDENTIALS__PROFILE_NAMES": '["local-profile"]', + "SOURCES__SOURCE__AWS__CREDENTIALS__ORG_ROLE_NAME": "LocalRole", + "CREDENTIALS__PROFILE_NAMES": '["local-profile"]', + "DESTINATION__BLOODHOUNDENTERPRISE__URL": "https://bhe.example.test", + }.items(): + monkeypatch.setenv(key, value) + providers = Container()[PluggableRunContext].providers + original_providers = providers.providers + local = CustomLoaderDocProvider( + "Local TOML equivalent", + lambda: { + "sources": { + "source": { + "aws": { + "credentials": { + "profile_names": ["local-profile"], + "org_role_name": "LocalRole", + } + } + } + } + }, + ) + + class UnsectionedDefaults(ConfigProvider): + def get_value(self, key, hint, pipeline_name, *sections): + return { + "profile_names": ["global-profile"], + "org_role_name": "GlobalRole", + }.get(key), key + + name = "Global source credential defaults" + supports_secrets = True + supports_sections = False + + providers.providers = [*original_providers, local, UnsectionedDefaults()] + original_environment = dict(os.environ) + try: + with managed_aws_profile({}, PAIR, workspace.directory) as profile: + assert ( + dlt.secrets.get("sources.source.aws.credentials.profile_name", str) + == profile + ) + assert ( + dlt.secrets.get( + "sources.source.aws.credentials.profile_names", list[str] + ) + is None + ) + assert ( + dlt.secrets.get("sources.source.aws.credentials.org_role_name", str) + is None + ) + assert dlt.secrets.get("credentials.profile_names", list[str]) is None + _assert_source_credentials(boto3.Session(profile_name=profile)) + assert ( + dlt.config.get("destination.bloodhoundenterprise.url", str) + == "https://bhe.example.test" + ) + assert dict(os.environ) == original_environment + finally: + providers.providers = original_providers + + +@pytest.mark.parametrize( + "role_profile", [False, True], ids=["raw-user", "initial-role"] +) +@pytest.mark.parametrize( + "dedicated_source", [False, True], ids=["management-profile", "source-account"] +) +def test_organization_settings_pass_through_with_generated_source_profile( + dedicated_source, + role_profile, + workspace, + profile_assumptions, +): + options = { + "org_role_name": "MemberCollectorRole", + "org_role_session_name": "job-session", + "include_management_account": False, + "region_name": "us-west-2", + } + if dedicated_source: + options.update( + source_account=SOURCE_ACCOUNT, + management_account_id=MANAGEMENT_ACCOUNT, + management_role_name="ManagementCollectorRole", + ) + if role_profile: + options.update( + role_arn=f"arn:aws:iam::{SOURCE_ACCOUNT}:role/InitialCollectionRole", + role_session_name="initial-session", + ) + profile_assumptions( + {"RoleArn": options["role_arn"], "RoleSessionName": "initial-session"} + ) + params = {"sources": {"source": {"aws": {"credentials": options}}}} + original = deepcopy(params) + with managed_aws_profile(params, PAIR, workspace.directory) as profile: + settings = dlt.secrets.get("sources.source.aws.credentials", dict) + assert settings == {**options, "profile_name": profile} + session = boto3.Session( + profile_name=settings["profile_name"], region_name=settings["region_name"] + ) + if role_profile: + _assert_role_credentials(session) + else: + _assert_source_credentials(session) + assert params == original + + +@pytest.mark.parametrize( + "role_profile", [False, True], ids=["raw-user", "role-then-raw"] +) +@pytest.mark.parametrize("failed", [False, True], ids=["success", "failure"]) +def test_successive_jobs_restore_environment_providers_and_remove_files( + failed, role_profile, profile_assumptions, monkeypatch +): + monkeypatch.setenv("AWS_ACCESS_KEY_ID", "scheduler-access") + monkeypatch.setenv("AWS_SECRET_ACCESS_KEY", "scheduler-secret") + scheduler_session = boto3.Session() + assert scheduler_session.get_credentials().method == "env" + original_environment = dict(os.environ) + providers = Container()[PluggableRunContext].providers + original_providers = providers.providers + original_traces = Container()[TraceLogContext] + paths, profiles = [], [] + for index in range(2): + workspace = AWSProfileWorkspace() + paths.append(workspace.directory) + options = ( + { + "sources": { + "source": {"aws": {"credentials": {"region_name": "eu-west-1"}}} + } + } + if index == 0 + else {} + ) + if role_profile and index == 0: + options["sources"]["source"]["aws"]["credentials"]["role_arn"] = ROLE_ARN + try: + with managed_aws_profile(options, PAIR, workspace.directory) as profile: + profiles.append(profile) + settings = dlt.secrets.get("sources.source.aws.credentials", dict) + assert settings["profile_name"] == profile + assert settings.get("region_name") == ( + "eu-west-1" if index == 0 else None + ) + session = boto3.Session( + profile_name=settings["profile_name"], + region_name=settings.get("region_name"), + ) + if role_profile and index == 0: + profile_assumptions( + {"RoleArn": ROLE_ARN, "RoleSessionName": profile} + ) + _assert_role_credentials(session) + else: + _assert_source_credentials(session) + if failed and index == 0: + raise RuntimeError(SECRET_KEY) + except ManagedConfigurationError as error: + assert failed and index == 0 + assert SECRET_KEY not in "".join(traceback.format_exception(error)) + finally: + assert not (workspace.directory / "credentials").exists() + workspace.cleanup() + assert dict(os.environ) == original_environment + assert providers.providers is original_providers + assert Container()[TraceLogContext] is original_traces + assert boto3.Session().get_credentials().access_key == "scheduler-access" + assert scheduler_session.get_credentials().access_key == "scheduler-access" + assert len(set(paths)) == len(set(profiles)) == 2 + assert all(not path.exists() for path in paths) + + +def test_worker_logs_errors_and_support_bundle_are_redacted( + workspace, tmp_path, caplog +): + log_path = tmp_path / "logs" + log_path.mkdir() + handler = logging.FileHandler(log_path / "openhound.log") + handler.setFormatter(openhound_logging.OpenHoundJSONFormatter()) + root = logging.getLogger() + root.addHandler(handler) + try: + with ( + pytest.raises(ManagedConfigurationError) as error, + caplog.at_level(logging.DEBUG), + managed_aws_profile({}, PAIR, workspace.directory), + ): + logger = logging.getLogger("openhound_aws.auth") + logger.info( + "Credential failure: %s", + SECRET_KEY, + extra={"response": {"credential": SECRET_KEY}}, + ) + try: + raise RuntimeError(f"AccessDenied for {ACCESS_KEY}: {SECRET_KEY}") + except RuntimeError: + logger.exception("SDK failure") + # Exercise a handler configured after the adapter was entered. + late_handler = openhound_logging._protect_handler( + logging.FileHandler(log_path / "ext_aws.log") + ) + root.addHandler(late_handler) + try: + logger.error("Late handler %s", SECRET_KEY) + finally: + root.removeHandler(late_handler) + late_handler.close() + raise RuntimeError(SECRET_KEY) + assert SECRET_KEY not in "".join(traceback.format_exception(error.value)) + assert SECRET_KEY not in caplog.text + assert ACCESS_KEY not in caplog.text + assert all(SECRET_KEY not in repr(record.__dict__) for record in caplog.records) + assert SECRET_KEY not in repr(PAIR) + assert "RuntimeError: AccessDenied" in (log_path / "openhound.log").read_text() + bundle = create_support_bundle("aws", log_path) + try: + with zipfile.ZipFile(bundle) as archive: + assert set(archive.namelist()) == {"openhound.log", "ext_aws.log"} + assert all( + SECRET_KEY.encode() not in archive.read(name) + for name in archive.namelist() + ) + assert b"RuntimeError: AccessDenied" in archive.read("openhound.log") + finally: + bundle.unlink() + bundle.parent.rmdir() + finally: + root.removeHandler(handler) + handler.close() + + +def _blocked_collection(collector, job_id, params, credentials, managed, directory): + with managed_aws_profile(params, credentials, directory): + time.sleep(60) + + +def _wait_for_credentials(directory): + deadline = time.monotonic() + 15 + while not (directory / "config").exists() and time.monotonic() < deadline: + time.sleep(0.01) + assert (directory / "credentials").is_file() + assert (directory / "config").is_file() + + +@pytest.mark.parametrize( + "claim_lost", [False, True], ids=["abrupt-death", "scheduler-kill"] +) +def test_parent_removes_credentials_after_forced_worker_termination( + claim_lost, monkeypatch +): + scheduler = service.Service( + "https://bhe.example.test", "token", "id", "aws", managed=True + ) + monkeypatch.setattr(scheduler.client, "start_collector_job", lambda job_id: None) + ends = [] + monkeypatch.setattr( + scheduler.client, + "end_managed_job", + lambda *args, **kwargs: ends.append((args, kwargs)), + ) + executor = scheduler.executor + submit = executor.submit + monkeypatch.setattr( + executor, "submit", lambda function, *args: submit(_blocked_collection, *args) + ) + try: + scheduler.run_claimed_job("job", {}, PAIR) + directory = scheduler._aws_profile_workspace.directory + _wait_for_credentials(directory) + assert SECRET_KEY in (directory / "credentials").read_text() + workers = list(executor._processes.values()) + if claim_lost: + scheduler._abort_managed_job("job") + assert ends == [] + else: + for worker in workers: + worker.kill() + worker.join(timeout=10) + scheduler.future.exception(timeout=10) + scheduler._handle_completed_job(scheduler.future) + assert ( + ends[0][1]["failure_message"] + == "Collection worker for 'aws' was terminated abruptly" + ) + assert all(not worker.is_alive() for worker in workers) + assert not directory.exists() + assert scheduler._aws_profile_workspace is None + finally: + for worker in list((executor._processes or {}).values()): + if worker.is_alive(): + worker.kill() + executor.shutdown(wait=True, cancel_futures=True) + scheduler.executor.shutdown(wait=True, cancel_futures=True) + scheduler._worker_logging.close() + scheduler._cleanup_aws_profile_workspace() + + +def _orphan_scheduler(temp_root, connection): + aws_profile.tempfile.tempdir = temp_root + workspace = AWSProfileWorkspace() + with managed_aws_profile({}, PAIR, workspace.directory): + connection.send(workspace.directory) + time.sleep(60) + + +def test_startup_reclaims_orphans_after_scheduler_death(tmp_path): + ctx = multiprocessing.get_context("spawn") + parent, child = ctx.Pipe() + process = ctx.Process(target=_orphan_scheduler, args=(str(tmp_path), child)) + process.start() + try: + assert parent.poll(15) + directory = parent.recv() + cleanup_orphaned_aws_profiles() + assert directory.exists() # Both process-held locks are still live. + process.kill() + process.join(timeout=10) + assert not process.is_alive() + cleanup_orphaned_aws_profiles() + assert not directory.exists() + finally: + if process.is_alive(): + process.kill() + process.join(timeout=10) + parent.close() + child.close() + + +def test_orphan_reaper_preserves_live_worker_after_scheduler_death(workspace): + executor = ProcessPoolExecutor(max_workers=1, max_tasks_per_child=1) + try: + future = executor.submit( + _blocked_collection, "aws", "job", {}, PAIR, True, workspace.directory + ) + _wait_for_credentials(workspace.directory) + workspace._owner.release() # Simulate the scheduler's OS lock being released. + cleanup_orphaned_aws_profiles() + assert (workspace.directory / "credentials").exists() + for worker in list(executor._processes.values()): + worker.kill() + worker.join(timeout=10) + future.exception(timeout=10) + cleanup_orphaned_aws_profiles() + assert not workspace.directory.exists() + finally: + executor.shutdown(wait=True, cancel_futures=True) diff --git a/tests/test_aws_secrets_manager.py b/tests/test_aws_secrets_manager.py new file mode 100644 index 00000000..85f5d29e --- /dev/null +++ b/tests/test_aws_secrets_manager.py @@ -0,0 +1,937 @@ +import base64 +import errno +import json +import logging +import traceback +from collections.abc import Mapping, Sequence +from typing import Any, NoReturn, cast + +import boto3 +import pytest +from botocore.awsrequest import AWSResponse +from botocore.exceptions import ( + BotoCoreError, + ClientError, + ConfigNotFound, + CredentialRetrievalError, + EndpointConnectionError, + HTTPClientError, + InvalidConfigError, + MetadataRetrievalError, + NoCredentialsError, + NoRegionError, + ParamValidationError, + PartialCredentialsError, + ProfileNotFound, +) + +from openhound.core.clients.aws_secrets_manager import ( + AWSConfigurationError, + AWSSecretsManager, + SecretBatchError, + SecretNotFoundError, + SecretPermissionError, + SecretRequestError, +) + +LOGGER_NAME = "openhound.core.clients.aws_secrets_manager" +SECRET_ID = "openhound/bhe" +SECRET_VALUE = "secret-value-that-must-not-be-logged" +PROVIDER_MESSAGE = ( + "role arn:aws:iam::123456789012:role/openhound cannot read the secret" +) + + +@pytest.fixture(autouse=True) +def restore_sdk_logging(): + def sdk_loggers(): + return { + name: sdk_logger + for name, sdk_logger in logging.Logger.manager.loggerDict.copy().items() + if name in ("boto3", "botocore") or name.startswith(("boto3.", "botocore.")) + if isinstance(sdk_logger, logging.Logger) + } + + original = { + name: ( + sdk_logger.level, + sdk_logger.disabled, + sdk_logger.propagate, + sdk_logger.handlers[:], + ) + for name, sdk_logger in sdk_loggers().items() + } + try: + yield + finally: + for name, sdk_logger in sdk_loggers().items(): + level, disabled, propagate, handlers = original.get( + name, (logging.NOTSET, False, True, []) + ) + sdk_logger.setLevel(level) + sdk_logger.disabled = disabled + sdk_logger.propagate = propagate + sdk_logger.handlers[:] = handlers + + +class FakeSecretsManagerClient: + def __init__( + self, + *, + get_response: Mapping[str, Any] | None = None, + get_error: BaseException | None = None, + batch_responses: Sequence[Mapping[str, Any] | BaseException] | None = None, + ) -> None: + self.get_response = get_response + self.get_error = get_error + self.batch_responses = list(batch_responses or []) + self.get_requests: list[str] = [] + self.batch_requests: list[list[str]] = [] + + def get_secret_value(self, *, SecretId: str) -> Mapping[str, Any]: + self.get_requests.append(SecretId) + if self.get_error is not None: + raise self.get_error + return self.get_response or {} + + def batch_get_secret_value(self, *, SecretIdList: list[str]) -> Mapping[str, Any]: + self.batch_requests.append(list(SecretIdList)) + if not self.batch_responses: + return {"SecretValues": [], "Errors": []} + response = self.batch_responses.pop(0) + if isinstance(response, BaseException): + raise response + return response + + +def aws_error( + code: str, + message: str, + operation: str = "GetSecretValue", +) -> ClientError: + return ClientError( + { + "Error": {"Code": code, "Message": message}, + "ResponseMetadata": { + "HTTPHeaders": {}, + "HTTPStatusCode": 400, + "HostId": "", + "RequestId": "", + "RetryAttempts": 0, + }, + }, + operation, + ) + + +def assert_safe_failure(error: BaseException, forbidden: set[str]) -> None: + assert all(value not in str(error) for value in forbidden) + + +def assert_safe_logs(caplog: pytest.LogCaptureFixture, forbidden: set[str]) -> None: + assert all(value not in caplog.text for value in forbidden) + assert all( + all(value not in record.getMessage() for value in forbidden) + and all(value not in repr(record.__dict__) for value in forbidden) + for record in caplog.records + ) + + +@pytest.mark.parametrize( + ("secret_text", "expected"), + [ + (SECRET_VALUE, SECRET_VALUE), + ('{"username":"user","enabled":true}', {"username": "user", "enabled": True}), + ('{"malformed":', '{"malformed":'), + ('[{"not": "an object result"}]', '[{"not": "an object result"}]'), + ], +) +def test_get_secret_uses_get_secret_value_for_text_and_json_objects( + secret_text: str, expected: str | dict[str, Any] +) -> None: + client = FakeSecretsManagerClient(get_response={"SecretString": secret_text}) + + assert AWSSecretsManager(client).get_secret(SECRET_ID) == expected + assert client.get_requests == [SECRET_ID] + assert client.batch_requests == [] + + +def test_get_secrets_returns_text_and_json_objects_from_batch_values() -> None: + json_secret = "json-secret" + client = FakeSecretsManagerClient( + batch_responses=[ + { + "SecretValues": [ + {"ARN": json_secret, "SecretString": '{"value": 2}'}, + {"Name": "text-secret", "SecretString": "value-1"}, + ], + "Errors": [], + } + ] + ) + + result = AWSSecretsManager(client).get_secrets(["text-secret", json_secret]) + + assert result == {"text-secret": "value-1", json_secret: {"value": 2}} + assert list(result) == ["text-secret", json_secret] + assert client.batch_requests == [["text-secret", json_secret]] + assert client.get_requests == [] + + +def test_get_secrets_splits_requests_into_chunks_of_20_and_merges_values() -> None: + secret_ids = [f"secret-{index}" for index in range(41)] + responses = [ + { + "SecretValues": [ + {"Name": secret_id, "SecretString": f"value-{secret_id}"} + for secret_id in reversed(secret_ids[start : start + 20]) + ] + } + for start in range(0, len(secret_ids), 20) + ] + client = FakeSecretsManagerClient(batch_responses=responses) + + result = AWSSecretsManager(client).get_secrets(secret_ids) + + assert result == {secret_id: f"value-{secret_id}" for secret_id in secret_ids} + assert list(result) == secret_ids + assert [len(request) for request in client.batch_requests] == [20, 20, 1] + assert client.batch_requests == [ + secret_ids[:20], + secret_ids[20:40], + secret_ids[40:], + ] + + +def test_get_secrets_continues_and_aggregates_typed_failures_safely( + caplog: pytest.LogCaptureFixture, +) -> None: + secret_ids = [f"secret-{index}" for index in range(21)] + first_failure = secret_ids[2] + second_failure = secret_ids[-1] + client = FakeSecretsManagerClient( + batch_responses=[ + { + "SecretValues": [ + {"Name": secret_id, "SecretString": SECRET_VALUE} + for secret_id in secret_ids[:20] + if secret_id != first_failure + ], + "Errors": [ + { + "SecretId": first_failure, + "ErrorCode": "ResourceNotFoundException", + "Message": PROVIDER_MESSAGE, + } + ], + }, + aws_error( + "AccessDeniedException", + PROVIDER_MESSAGE, + operation="BatchGetSecretValue", + ), + ] + ) + forbidden = {SECRET_VALUE} + + with ( + caplog.at_level(logging.DEBUG, logger=LOGGER_NAME), + pytest.raises(SecretBatchError) as raised, + ): + AWSSecretsManager(client).get_secrets(secret_ids) + + assert client.batch_requests == [secret_ids[:20], secret_ids[20:]] + assert set(raised.value.failures) == {first_failure, second_failure} + assert isinstance(raised.value.failures[first_failure], SecretNotFoundError) + assert isinstance(raised.value.failures[second_failure], SecretPermissionError) + assert_safe_failure(raised.value, forbidden) + for failure in raised.value.failures.values(): + assert_safe_failure(failure, forbidden) + + attempts = [ + cast(Any, record) + for record in caplog.records + if getattr(record, "outcome", None) == "attempt" + ] + assert [record.batch_number for record in attempts] == [1, 2] + error_records = [ + cast(Any, record) + for record in caplog.records + if getattr(record, "outcome", None) == "failure" + ] + assert len(error_records) == 2 + assert first_failure in caplog.text + assert "ResourceNotFoundException" in caplog.text + assert PROVIDER_MESSAGE in caplog.text + assert all( + record.failure_category == "SecretBatchError" for record in error_records + ) + assert_safe_logs(caplog, forbidden) + + +@pytest.mark.parametrize( + ("error", "expected_type"), + [ + ( + aws_error("ResourceNotFoundException", PROVIDER_MESSAGE), + SecretNotFoundError, + ), + ( + aws_error("AccessDeniedException", PROVIDER_MESSAGE), + SecretPermissionError, + ), + ( + EndpointConnectionError(endpoint_url="https://secretsmanager.example"), + SecretRequestError, + ), + ], +) +def test_get_secret_classifies_aws_errors( + error: BotoCoreError | ClientError, + expected_type: type[Exception], +) -> None: + client = FakeSecretsManagerClient(get_error=error) + + with pytest.raises(expected_type) as raised: + AWSSecretsManager(client).get_secret(SECRET_ID) + + assert raised.value.context["sdk_exception_type"] == type(error).__name__ + if isinstance(error, ClientError): + assert PROVIDER_MESSAGE in str(raised.value) + else: + assert "https://secretsmanager.example" in str(raised.value) + + +@pytest.mark.parametrize( + "sdk_error", + [ + NoCredentialsError(), + PartialCredentialsError(provider="aws", cred_var="secret-key"), + NoRegionError(), + InvalidConfigError(error_msg=PROVIDER_MESSAGE), + ConfigNotFound(path=PROVIDER_MESSAGE), + ProfileNotFound(profile=PROVIDER_MESSAGE), + CredentialRetrievalError( + provider="credential-provider", error_msg=PROVIDER_MESSAGE + ), + ], +) +def test_client_creation_configuration_failures_are_classified( + monkeypatch: pytest.MonkeyPatch, + sdk_error: BotoCoreError, +) -> None: + def fail_client(*_args: object, **_kwargs: object) -> NoReturn: + raise sdk_error + + monkeypatch.setattr( + "openhound.core.clients.aws_secrets_manager.boto3.client", fail_client + ) + + with pytest.raises(AWSConfigurationError) as raised: + AWSSecretsManager().get_secret(SECRET_ID) + assert type(sdk_error).__name__ in str(raised.value) + assert raised.value.context["sdk_exception_type"] == type(sdk_error).__name__ + if isinstance(sdk_error, NoRegionError): + assert "AWS_DEFAULT_REGION" in str(raised.value) + assert "selected AWS profile" in str(raised.value) + elif not isinstance(sdk_error, CredentialRetrievalError): + assert str(sdk_error) in str(raised.value) + + +def test_successful_retrieval_logs_safe_attempt_and_result_records( + caplog: pytest.LogCaptureFixture, +) -> None: + client = FakeSecretsManagerClient(get_response={"SecretString": SECRET_VALUE}) + + with caplog.at_level(logging.DEBUG, logger=LOGGER_NAME): + AWSSecretsManager(client).get_secret(SECRET_ID) + + records = [cast(Any, record) for record in caplog.records] + assert [(record.levelno, record.outcome) for record in records] == [ + (logging.DEBUG, "attempt"), + (logging.DEBUG, "success"), + ] + assert all(record.provider == "aws.secretsmanager" for record in records) + assert all(record.operation == "get_secret_value" for record in records) + assert all(record.secret_count == 1 for record in records) + assert_safe_logs(caplog, {SECRET_ID, SECRET_VALUE}) + + caplog.clear() + batch_client = FakeSecretsManagerClient( + batch_responses=[ + { + "SecretValues": [ + {"Name": "first", "SecretString": SECRET_VALUE}, + {"Name": "second", "SecretString": "second-value"}, + ] + } + ] + ) + with caplog.at_level(logging.DEBUG, logger=LOGGER_NAME): + AWSSecretsManager(batch_client).get_secrets(["first", "second"]) + + batch_records = [cast(Any, record) for record in caplog.records] + assert [(record.levelno, record.outcome) for record in batch_records] == [ + (logging.DEBUG, "attempt"), + (logging.DEBUG, "success"), + ] + assert all(record.operation == "batch_get_secret_value" for record in batch_records) + assert all(record.secret_count == 2 for record in batch_records) + assert all(record.batch_number == 1 for record in batch_records) + assert_safe_logs( + caplog, + {"first", "second", SECRET_VALUE, "second-value"}, + ) + + +def test_get_secrets_rejects_a_string_instead_of_ids() -> None: + with pytest.raises(TypeError, match="iterable of secret identifiers"): + AWSSecretsManager(FakeSecretsManagerClient()).get_secrets(SECRET_ID) + + +def test_bhe_binary_json_secret_is_decoded_in_memory(caplog) -> None: + client = FakeSecretsManagerClient( + get_response={ + "SecretBinary": b'{"credential":"secret-value-that-must-not-be-logged"}' + } + ) + with caplog.at_level(logging.DEBUG, logger=LOGGER_NAME): + result = AWSSecretsManager(client).get_secret(SECRET_ID) + assert result == {"credential": SECRET_VALUE} + assert_safe_logs(caplog, {SECRET_ID, SECRET_VALUE}) + + +def test_invalid_binary_secret_is_safely_classified(caplog) -> None: + client = FakeSecretsManagerClient( + get_response={"SecretBinary": b"\xffsecret-value-that-must-not-be-logged"} + ) + with ( + caplog.at_level(logging.DEBUG, logger=LOGGER_NAME), + pytest.raises(SecretRequestError) as raised, + ): + AWSSecretsManager(client).get_secret(SECRET_ID) + assert_safe_failure(raised.value, {SECRET_ID, SECRET_VALUE}) + assert_safe_logs(caplog, {SECRET_ID, SECRET_VALUE}) + + +def test_boto_client_uses_bounded_timeouts_and_retries(monkeypatch) -> None: + captured = {} + client = FakeSecretsManagerClient(get_response={"SecretString": SECRET_VALUE}) + + def create_client(service, **kwargs): + captured.update(service=service, **kwargs) + return client + + monkeypatch.setattr( + "openhound.core.clients.aws_secrets_manager.boto3.client", create_client + ) + assert AWSSecretsManager().get_secret(SECRET_ID) == SECRET_VALUE + config = captured["config"] + assert captured["service"] == "secretsmanager" + assert config.connect_timeout == 10 + assert config.read_timeout == 20 + assert config.retries == {"mode": "standard", "total_max_attempts": 4} + + +@pytest.mark.parametrize( + "operation,response_kind,injected", + [ + pytest.param("single", "string", False, id="single-string-created"), + pytest.param("single", "binary", True, id="single-binary-injected"), + pytest.param("single", "provider_error", False, id="single-error-created"), + pytest.param("batch", "string", True, id="batch-string-injected"), + pytest.param("batch", "binary", False, id="batch-binary-created"), + pytest.param("batch", "provider_error", True, id="batch-error-injected"), + ], +) +def test_real_sdk_http_pipeline_does_not_expose_secrets( + monkeypatch, caplog, operation, response_kind, injected +) -> None: + document = json.dumps({"credential": SECRET_VALUE}) + encoded_document = base64.b64encode(document.encode()).decode() + secret = ( + {"SecretBinary": encoded_document} + if response_kind == "binary" + else {"SecretString": document} + ) + if response_kind == "provider_error": + response_body = { + "__type": "AccessDeniedException", + "message": f"{PROVIDER_MESSAGE}: {SECRET_ID}", + } + status = 400 + else: + response_body = ( + {"SecretValues": [{"Name": SECRET_ID, **secret}], "Errors": []} + if operation == "batch" + else secret + ) + status = 200 + raw_body = json.dumps(response_body).encode() + forbidden = { + SECRET_VALUE, + encoded_document, + raw_body.decode(), + } + requests = [] + clients = [] + + class RawResponse: + def stream(self, *_args, **_kwargs): + yield raw_body + + def create_client(service, **kwargs): + # This preexisting child has its own handler and explicit DEBUG level. + logging.getLogger("botocore.endpoint").debug( + "SDK construction %s", SECRET_ID, extra={"payload": SECRET_VALUE} + ) + session = boto3.Session( + aws_access_key_id="fake-access-key", + aws_secret_access_key="fake-secret-key", + region_name="us-east-1", + ) + client = session.client(service, **kwargs) + clients.append(client) + monkeypatch.setattr(client._endpoint.http_session, "send", send) + return client + + def send(request): + requests.append(request) + # These children are first created after suppression; explicit DEBUG + # must not bypass the namespace sink, even for ERROR records. + for namespace in ("boto3.resources", "botocore.endpoint"): + child = logging.getLogger( + f"{namespace}.openhound_pipeline_{operation}_{response_kind}_{injected}" + ) + child.setLevel(logging.DEBUG) + child.error( + "SDK response %s", PROVIDER_MESSAGE, extra={"payload": raw_body} + ) + return AWSResponse( + request.url, + status, + {"content-type": "application/x-amz-json-1.1"}, + RawResponse(), + ) + + monkeypatch.setattr( + "openhound.core.clients.aws_secrets_manager.boto3.client", create_client + ) + if injected: + # An injected real client is constructed by its caller. Only the + # wrapper's subsequent retrieval is under test in this case. + client = boto3.Session( + aws_access_key_id="fake-access-key", + aws_secret_access_key="fake-secret-key", + region_name="us-east-1", + ).client("secretsmanager") + clients.append(client) + monkeypatch.setattr(client._endpoint.http_session, "send", send) + else: + client = None + + for name in ( + "boto3", + "botocore", + "boto3.resources", + "botocore.endpoint", + "botocore.parsers", + ): + sdk_logger = logging.getLogger(name) + sdk_logger.disabled = False + sdk_logger.setLevel(logging.DEBUG) + sdk_logger.addHandler(caplog.handler) + sdk_logger.propagate = False + caplog.clear() + + try: + with caplog.at_level(logging.DEBUG): + manager = AWSSecretsManager(client) + if response_kind == "provider_error": + error_type = ( + SecretBatchError if operation == "batch" else SecretPermissionError + ) + with pytest.raises(error_type) as raised: + if operation == "batch": + manager.get_secrets([SECRET_ID]) + else: + manager.get_secret(SECRET_ID) + assert_safe_failure(raised.value, forbidden) + formatted = "".join(traceback.format_exception(raised.value)) + assert all(value not in formatted for value in forbidden) + if operation == "batch": + assert isinstance( + raised.value.failures[SECRET_ID], SecretPermissionError + ) + assert_safe_failure(raised.value.failures[SECRET_ID], forbidden) + elif operation == "batch": + assert manager.get_secrets([SECRET_ID]) == { + SECRET_ID: {"credential": SECRET_VALUE} + } + else: + assert manager.get_secret(SECRET_ID) == {"credential": SECRET_VALUE} + finally: + for sdk_client in clients: + sdk_client.close() + + assert len(requests) == 1 + assert SECRET_ID in requests[0].body.decode() + assert_safe_logs(caplog, forbidden) + records = [record for record in caplog.records if record.name == LOGGER_NAME] + assert [getattr(record, "outcome", None) for record in records] == [ + "attempt", + "failure" if response_kind == "provider_error" else "success", + ] + assert all(record.exc_info is None for record in records) + assert not any( + record.name.startswith(("boto3", "botocore")) for record in caplog.records + ) + + +@pytest.mark.parametrize("during_construction", [False, True]) +def test_sts_failure_retains_actual_operation_and_metadata( + monkeypatch, caplog, during_construction +): + error = aws_error( + "AccessDenied", + "Role trust condition did not match", + "AssumeRoleWithWebIdentity", + ) + error.response["ResponseMetadata"].update( + RequestId="request-123", HTTPStatusCode=403, RetryAttempts=2 + ) + + def fail(*args, **kwargs): + raise error + + monkeypatch.setattr("openhound.core.clients.aws_secrets_manager.boto3.client", fail) + manager = AWSSecretsManager( + None if during_construction else FakeSecretsManagerClient(get_error=error) + ) + with caplog.at_level(logging.DEBUG), pytest.raises(AWSConfigurationError) as raised: + manager.get_secret(SECRET_ID) + text = str(raised.value) + for expected in ( + "STS AssumeRoleWithWebIdentity", + "AccessDenied", + "Role trust condition", + "role trust", + "request-123", + "HTTP 403", + "SDK retries 2", + ): + assert expected in text + assert "Secrets Manager denied" not in text + assert raised.value.context["aws_operation"] == "AssumeRoleWithWebIdentity" + assert raised.value.context["RetryAttempts"] == 2 + assert all(record.exc_info is None for record in caplog.records) + + +@pytest.mark.parametrize("batch", [False, True]) +def test_unrelated_runtime_errors_are_classified_without_provider_output(caplog, batch): + error = RuntimeError(SECRET_VALUE) + client = FakeSecretsManagerClient(get_error=error, batch_responses=[error]) + manager = AWSSecretsManager(client) + with caplog.at_level(logging.DEBUG): + if batch: + with pytest.raises(SecretBatchError) as raised: + manager.get_secrets([SECRET_ID]) + failure = raised.value.failures[SECRET_ID] + else: + with pytest.raises(SecretRequestError) as raised: + manager.get_secret(SECRET_ID) + failure = raised.value + + assert isinstance(failure, SecretRequestError) + assert str(failure) == "AWS secret retrieval failed (RuntimeError)" + assert failure.context["sdk_exception_type"] == "RuntimeError" + assert SECRET_VALUE not in "".join(traceback.format_exception(raised.value)) + assert_safe_logs(caplog, {SECRET_VALUE}) + + +@pytest.mark.parametrize( + "code,expected_type", + [ + ("ExpiredTokenException", AWSConfigurationError), + ("UnrecognizedClientException", AWSConfigurationError), + ("ThrottlingException", SecretRequestError), + ("InternalServiceError", SecretRequestError), + ], +) +def test_service_failures_preserve_code_message_region_and_secret(code, expected_type): + from types import SimpleNamespace + + client = FakeSecretsManagerClient(get_error=aws_error(code, PROVIDER_MESSAGE + ".")) + client.meta = SimpleNamespace(region_name="us-west-2") + with pytest.raises(expected_type) as raised: + AWSSecretsManager(client).get_secret(SECRET_ID) + for expected in (code, PROVIDER_MESSAGE, "GetSecretValue", SECRET_ID, "us-west-2"): + assert expected in str(raised.value) + assert raised.value.context["region"] == "us-west-2" + assert "all attempts" not in str(raised.value) + assert ".." not in str(raised.value) + + +@pytest.mark.parametrize("code", ["SignatureDoesNotMatch", "InvalidSignatureException"]) +def test_signature_diagnostics_exclude_authentication_material_from_normal_logs( + caplog, code +): + from openhound.core.clients.bhe_credentials import AWSBHECredentials + + error = aws_error( + code, + "The request signature we calculated does not match the signature you provided. " + "The Canonical String for this request should have been\n" + f"GET\n/\n\nhost:secretsmanager.example\nx-amz-security-token:{SECRET_VALUE}\n" + f"Authorization: AWS4-HMAC-SHA256 Credential={SECRET_VALUE}", + ) + error.response["ResponseMetadata"].update(RequestId="request-123", RetryAttempts=2) + manager = AWSSecretsManager(FakeSecretsManagerClient(get_error=error)) + with caplog.at_level(logging.ERROR), pytest.raises(AWSConfigurationError) as raised: + AWSBHECredentials(SECRET_ID, manager).refresh() + + assert len(caplog.records) == 1 + for expected in ( + code, + "GetSecretValue", + "request signature", + "request-123", + "SDK retries 2", + ): + assert expected in caplog.text + assert "Canonical String" not in caplog.text + assert SECRET_VALUE not in "".join(traceback.format_exception(raised.value)) + assert_safe_logs(caplog, {SECRET_VALUE}) + + +@pytest.mark.parametrize( + "failure,expected", + [ + ("invalid-json", "unable to parse JSON"), + ("timeout", "ReadTimeoutError"), + ("http-error", "HTTP 403"), + ], +) +def test_container_provider_failures_keep_safe_cause(caplog, failure, expected): + from types import SimpleNamespace + + from botocore.credentials import ContainerProvider + from botocore.exceptions import ReadTimeoutError + from botocore.utils import ContainerMetadataFetcher + + from openhound.core.clients.bhe_credentials import AWSBHECredentials + + endpoint = "http://localhost/credentials" + + def send(request): + if failure == "timeout": + raise ReadTimeoutError(endpoint_url=f"{endpoint}?token={SECRET_VALUE}") + return SimpleNamespace( + status_code=403 if failure == "http-error" else 200, + content=SECRET_VALUE.encode(), + ) + + fetcher = ContainerMetadataFetcher( + session=SimpleNamespace(send=send), sleep=lambda _: None + ) + fetcher.RETRY_ATTEMPTS = 1 + provider = ContainerProvider( + environ={"AWS_CONTAINER_CREDENTIALS_FULL_URI": endpoint}, fetcher=fetcher + ) + + def retrieve(**kwargs): + return provider.load() + + manager = AWSSecretsManager(SimpleNamespace(get_secret_value=retrieve)) + with caplog.at_level(logging.ERROR), pytest.raises(AWSConfigurationError) as raised: + AWSBHECredentials(SECRET_ID, manager).refresh() + + assert len(caplog.records) == 1 + assert expected in caplog.text + assert "container-role" in caplog.text + if failure == "timeout": + assert raised.value.context["transport_exception_type"] == "ReadTimeoutError" + assert SECRET_VALUE not in "".join(traceback.format_exception(raised.value)) + assert_safe_logs(caplog, {SECRET_VALUE}) + + +@pytest.mark.parametrize( + "parameter,constraint", + [ + ("RoleArn", "valid min length: 20"), + ("WebIdentityToken", "valid types: "), + ], +) +def test_validation_keeps_parameter_and_rule_without_values( + caplog, parameter, constraint +): + from botocore.session import get_session + from botocore.validate import validate_parameters + + from openhound.core.clients.bhe_credentials import AWSBHECredentials + + shape = ( + get_session() + .get_service_model("sts") + .operation_model("AssumeRoleWithWebIdentity") + .input_shape + ) + params: dict[str, Any] = { + "RoleArn": "arn:aws:iam::123456789012:role/openhound", + "RoleSessionName": "openhound", + "WebIdentityToken": SECRET_VALUE, + } + params[parameter] = "x" if parameter == "RoleArn" else {"token": SECRET_VALUE} + with pytest.raises(ParamValidationError) as sdk_error: + validate_parameters(params, shape) + + manager = AWSSecretsManager(FakeSecretsManagerClient(get_error=sdk_error.value)) + caplog.clear() + with caplog.at_level(logging.ERROR), pytest.raises(SecretRequestError) as raised: + AWSBHECredentials(SECRET_ID, manager).refresh() + + assert len(caplog.records) == 1 + assert f"parameter {parameter}" in caplog.text + assert constraint in caplog.text + assert "value:" not in str(raised.value) + assert SECRET_VALUE not in "".join(traceback.format_exception(raised.value)) + assert_safe_logs(caplog, {SECRET_VALUE}) + + +@pytest.mark.parametrize( + "error,explanation", + [ + ( + CredentialRetrievalError(provider="custom-process", error_msg=SECRET_VALUE), + "custom-process", + ), + ( + CredentialRetrievalError( + provider="container-role", + error_msg=f"Received non 200 response 403 from container metadata: {SECRET_VALUE}", + ), + "container-role", + ), + ( + ParamValidationError(report=f"Invalid WebIdentityToken: {SECRET_VALUE}"), + "parameter validation", + ), + ( + MetadataRetrievalError( + error_msg=f"Received non 200 response 401 from container metadata: {SECRET_VALUE}" + ), + "HTTP 401", + ), + (HTTPClientError(error=ValueError(SECRET_VALUE)), "ValueError"), + ], +) +def test_unsafe_provider_and_validation_details_are_omitted(error, explanation, caplog): + with ( + caplog.at_level(logging.DEBUG), + pytest.raises((AWSConfigurationError, SecretRequestError)) as raised, + ): + AWSSecretsManager(FakeSecretsManagerClient(get_error=error)).get_secret( + SECRET_ID + ) + assert explanation in str(raised.value) + assert "omitted" not in str(raised.value) + if ( + isinstance(error, CredentialRetrievalError) + and error.kwargs["provider"] == "container-role" + ): + assert "HTTP 403" in str(raised.value) + assert raised.value.context["credential_endpoint_http_status"] == 403 + assert SECRET_VALUE not in "".join(traceback.format_exception(raised.value)) + assert_safe_logs(caplog, {SECRET_VALUE}) + + +@pytest.mark.parametrize( + "flag,retries,explanation", + [ + ("MaxAttemptsReached", 3, "SDK maximum attempts reached"), + ("RetryQuotaReached", 1, "SDK retry quota exhausted"), + ], +) +def test_sdk_retry_exhaustion_is_reported_only_when_supplied( + flag, retries, explanation +): + error = aws_error("ThrottlingException", "Rate exceeded") + error.response["ResponseMetadata"].update({"RetryAttempts": retries, flag: True}) + manager = AWSSecretsManager(FakeSecretsManagerClient(get_error=error)) + with pytest.raises(SecretRequestError) as raised: + manager.get_secret(SECRET_ID) + assert explanation in str(raised.value) + assert f"SDK retries {retries}" in str(raised.value) + assert raised.value.context[flag] is True + + del error.response["ResponseMetadata"][flag] + with pytest.raises(SecretRequestError) as raised: + manager.get_secret(SECRET_ID) + assert explanation not in str(raised.value) + assert flag not in raised.value.context + + +@pytest.mark.parametrize( + "error_type,error_number,during_construction", + [ + (FileNotFoundError, errno.ENOENT, False), + (PermissionError, errno.EACCES, True), + ], +) +def test_aws_file_access_failures_retain_path_and_reason_in_bootstrap_logs( + monkeypatch, caplog, error_type, error_number, during_construction +): + import os + from types import SimpleNamespace + + from botocore.utils import FileWebIdentityTokenLoader + + from openhound.core.clients.bhe_credentials import AWSBHECredentials + + token_path = "/var/run/secrets/eks.amazonaws.com/serviceaccount/token" + + def fail_open(path): + raise error_type(error_number, os.strerror(error_number), path) from ValueError( + SECRET_VALUE + ) + + token_loader = FileWebIdentityTokenLoader(token_path, _open=fail_open) + + def fail(*args, **kwargs): + return token_loader() + + monkeypatch.setattr("openhound.core.clients.aws_secrets_manager.boto3.client", fail) + manager = AWSSecretsManager( + None if during_construction else SimpleNamespace(get_secret_value=fail) + ) + with pytest.raises(AWSConfigurationError) as raised: + manager.get_secret(SECRET_ID) + assert raised.value.context["file_path"] == token_path + assert raised.value.context["os_errno"] == error_number + assert SECRET_VALUE not in "".join(traceback.format_exception(raised.value)) + assert_safe_logs(caplog, {SECRET_VALUE}) + caplog.clear() + + def stop_polling(interval): + assert interval == 17 + raise RuntimeError("stop bootstrap polling") + + monkeypatch.setattr( + "openhound.core.clients.bhe_credentials.time.sleep", stop_polling + ) + with ( + caplog.at_level(logging.WARNING), + pytest.raises(RuntimeError, match="stop bootstrap polling"), + ): + AWSBHECredentials(SECRET_ID, manager).poll_until_available(17) + assert len(caplog.records) == 1 + for expected in ( + error_type.__name__, + token_path, + os.strerror(error_number), + "Retrying bootstrap retrieval in 17 seconds", + ): + assert expected in caplog.text + assert "mounts and permissions" in caplog.text + assert_safe_logs(caplog, {SECRET_VALUE}) + assert all(record.exc_info is None for record in caplog.records) diff --git a/tests/test_bhe_credentials.py b/tests/test_bhe_credentials.py new file mode 100644 index 00000000..a82d59c7 --- /dev/null +++ b/tests/test_bhe_credentials.py @@ -0,0 +1,160 @@ +import pytest + +from openhound.core.clients.bhe_credentials import ( + AWSBHECredentials, + BHECredentials, + InvalidBHECredentialsSecret, +) + + +class FakeSecretsManager: + def __init__(self, secret): + self.secret = secret + self.requested = [] + + def get_secret(self, secret_id): + self.requested.append(secret_id) + return self.secret + + +def test_aws_bhe_credentials_loads_token_pair_from_json_secret(): + manager = FakeSecretsManager( + {"token_id": "12345678-1234-1234-1234-123456789abc", "token_key": "key"} + ) + + result = AWSBHECredentials("production/bhe", manager).refresh() + + assert result == BHECredentials( + token_id="12345678-1234-1234-1234-123456789abc", token_key="key" + ) + assert manager.requested == ["production/bhe"] + + +@pytest.mark.parametrize( + "secret", + [ + "not-json", + {"token_id": "not-a-uuid", "token_key": "key"}, + {"token_id": " ", "token_key": "key"}, + {"token_id": 123, "token_key": "key"}, + {}, + {"token_id": "id"}, + {"token_key": "key"}, + {"token_id": "", "token_key": "key"}, + ], +) +def test_aws_bhe_credentials_rejects_invalid_secret_values(secret): + with pytest.raises(InvalidBHECredentialsSecret): + AWSBHECredentials("production/bhe", FakeSecretsManager(secret)).refresh() + + +def test_expired_sdk_refresh_logs_safe_cause(caplog): + import logging + from datetime import UTC, datetime, timedelta + + import boto3 + import botocore.session + from botocore.credentials import RefreshableCredentials + + from openhound.core.clients.aws_secrets_manager import ( + AWSConfigurationError, + AWSSecretsManager, + ) + + metadata = { + "access_key": "private-access-key", + "secret_key": "private-secret-key", + "token": "private-session-token", + "expiry_time": (datetime.now(UTC) - timedelta(hours=1)).isoformat(), + } + session = botocore.session.get_session() + session._credentials = RefreshableCredentials.create_from_metadata( + metadata=metadata, refresh_using=lambda: metadata, method="container-role" + ) + client = boto3.Session(botocore_session=session).client( + "secretsmanager", region_name="us-east-1" + ) + provider = AWSBHECredentials("production/bhe", AWSSecretsManager(client)) + caplog.clear() + + with caplog.at_level(logging.INFO), pytest.raises(AWSConfigurationError) as raised: + provider.refresh() + + assert raised.value.failure_summary == "AWS credentials unavailable" + assert len(caplog.records) == 1 + assert "refreshed credentials are still expired" in caplog.text + assert "production/bhe" in caplog.text + assert "private-" not in repr([record.__dict__ for record in caplog.records]) + assert all(record.exc_info is None for record in caplog.records) + + +def test_unexpected_refresh_failure_logs_type_without_payload(caplog): + import logging + from types import SimpleNamespace + + error = ValueError("private-provider-output") + + def retrieve(secret_id): + raise error + + provider = AWSBHECredentials("production/bhe", SimpleNamespace(get_secret=retrieve)) + with caplog.at_level(logging.INFO), pytest.raises(ValueError) as raised: + provider.refresh() + + assert raised.value is error + assert len(caplog.records) == 1 + assert "Unexpected credential reader failure (ValueError)" in caplog.text + assert "production/bhe" in caplog.text + assert "private-provider-output" not in repr(caplog.records[0].__dict__) + assert caplog.records[0].exc_info is None + + +@pytest.mark.parametrize("failure", ["region", "provider", "document"]) +def test_bootstrap_logs_each_failure_with_cause_and_delay(monkeypatch, caplog, failure): + import logging + from types import SimpleNamespace + + from botocore.exceptions import CredentialRetrievalError, NoRegionError + + from openhound.core.clients.aws_secrets_manager import AWSSecretsManager + + calls = [] + + def retrieve(**kwargs): + calls.append(kwargs) + if len(calls) <= 2: + if failure == "document": + return {"SecretString": '{"token_key":"private-token-key"}'} + if failure == "provider": + raise CredentialRetrievalError( + provider="custom-process", error_msg="private-token-key" + ) + raise NoRegionError() + return { + "SecretString": '{"token_id":"12345678-1234-1234-1234-123456789abc","token_key":"private-token-key"}' + } + + provider = AWSBHECredentials( + "production/bhe", AWSSecretsManager(SimpleNamespace(get_secret_value=retrieve)) + ) + delays = [] + monkeypatch.setattr( + "openhound.core.clients.bhe_credentials.time.sleep", delays.append + ) + with caplog.at_level(logging.WARNING): + assert provider.poll_until_available(17).token_key == "private-token-key" + assert delays == [17, 17] + assert len(caplog.records) == 2 + expected = { + "region": "AWS_DEFAULT_REGION", + "provider": "custom-process", + "document": "non-empty token_id", + }[failure] + assert all(expected in record.getMessage() for record in caplog.records) + assert all( + "Retrying bootstrap retrieval in 17 seconds" in record.getMessage() + for record in caplog.records + ) + assert "production/bhe" in caplog.text + assert "private-token-key" not in repr([r.__dict__ for r in caplog.records]) + assert all(record.exc_info is None for record in caplog.records) diff --git a/tests/test_bhe_job_scheduling.py b/tests/test_bhe_job_scheduling.py index 7e3746f3..97a77a55 100644 --- a/tests/test_bhe_job_scheduling.py +++ b/tests/test_bhe_job_scheduling.py @@ -3,28 +3,57 @@ import hashlib import json import logging -from concurrent.futures import Future +import os +import subprocess +import sys +import time +import traceback +from concurrent.futures import Future, ProcessPoolExecutor from concurrent.futures.process import BrokenProcessPool +from copy import deepcopy +from datetime import UTC, datetime, timedelta from pathlib import Path from types import SimpleNamespace from urllib.parse import urlsplit +from uuid import UUID +import boto3 +import dlt import pytest +import requests +from botocore.exceptions import ClientError +from botocore.httpsession import URLLib3Session +from botocore.session import Session as BotocoreSession +from botocore.stub import Stubber from fastapi import FastAPI, Request, Response from fastapi.testclient import TestClient from openhound.core.clients import bloodhound, bloodhound_enterprise +from openhound.core.clients.aws_secrets_manager import ( + AWSSecretsManager, + SecretPermissionError, +) +from openhound.core.clients.bhe_credentials import AWSBHECredentials, BHECredentials from openhound.core.clients.bloodhound import BloodHoundHTTPError -from openhound.core.clients.bloodhound_enterprise import JobStatus +from openhound.core.clients.bloodhound_enterprise import ( + JobStatus, + ManagedJobOutcome, +) from openhound.core.clients.models.jobs import ( + CollectorJob, + CollectorJobsAvailable, ManagementOperation, ManagementOperationStatus, ManagementOperationType, ) from openhound.core.models.graph import Graph from openhound.scheduler import service as scheduler_service +from openhound.scheduler.aws_profile import AWSAccessKeyCredentials, AWSAccessKeyError +from openhound.scheduler.configuration import ManagedConfigurationError from openhound.scheduler.service import ( ExtensionNotFoundError, + ManagedJobProcessingError, + ManagedRuntimeUnavailableError, Result, Service, _subprocess_collect, @@ -32,6 +61,12 @@ TEST_DATA_DIR = Path(__file__).parent / "test_data" / "api" / "jobs" MANAGEMENT_DATA_DIR = Path(__file__).parent / "test_data" / "api" / "management" +CLAIMED_PARAMS = {"sources": {"source": {"random": {"node_count": 4}}}} +COLLECTOR_SECRETS = { + "sources": { + "source": {"random": {"credentials": {"token": "private-credential-value"}}} + } +} def load_json(filename: str) -> dict: @@ -51,10 +86,32 @@ def mock_bloodhound_api(): app.state.job_started = False app.state.job_ended = False app.state.end_payload = None + app.state.collector_job_end_requests = [] + app.state.collector_job_end_statuses = [] + app.state.collector_job_claim_requests = [] + app.state.collector_job_claim_statuses = [] + app.state.collector_job_stateful_claims = False + app.state.collector_job_active_id = None + app.state.collector_job_lost_claim_responses = 0 + app.state.collector_job_claim_response = load_json( + "collector_jobs_available_with_job.json" + )["data"]["jobs"][0] + app.state.managed_events = [] + app.state.collector_job_start_requests = [] + app.state.collector_job_start_statuses = [] + app.state.collector_job_heartbeat_requests = [] app.state.start_payload = None + app.state.jobs_available_requests = 0 + app.state.jobs_current_requests = 0 + app.state.collector_job_queue_requests = [] + app.state.collector_job_queue_response = load_json( + "collector_jobs_available_empty.json" + ) + app.state.collector_job_queue_error_status = None app.state.client_update_payload = None app.state.ingested_edges = 0 app.state.management_operations = [] + app.state.management_available_requests = 0 app.state.operation_started = False app.state.operation_ended = False app.state.operation_start_payload = None @@ -65,29 +122,122 @@ def mock_bloodhound_api(): app.state.uploaded_parts = [] app.state.artifact_completed = False app.state.ingested_nodes = 0 + app.state.managed_ingest_requests = [] + app.state.managed_ingested_nodes = 0 + app.state.managed_ingested_edges = 0 @app.get("/api/v2/jobs/available") async def jobs_available(): + app.state.jobs_available_requests += 1 if not app.state.job_started: return load_json("jobs_available_with_job.json") return load_json("jobs_available_empty.json") @app.get("/api/v2/jobs/current") async def jobs_current(): + app.state.jobs_current_requests += 1 return Response(status_code=404) + @app.get("/api/v2/collie-job-queue/available") + async def collector_job_queue(request: Request): + query_params = dict(request.query_params) + query_params["sort_by"] = request.query_params.getlist("sort_by") + app.state.collector_job_queue_requests.append(query_params) + if app.state.collector_job_queue_error_status is not None: + return Response(status_code=app.state.collector_job_queue_error_status) + if app.state.collector_job_stateful_claims: + jobs = [ + job + for job in app.state.collector_job_queue_response["data"]["jobs"] + if job["id"] != app.state.collector_job_active_id + ] + return { + "count": len(jobs), + "skip": 0, + "limit": 1, + "data": {"jobs": jobs[:1]}, + } + return app.state.collector_job_queue_response + @app.post("/api/v2/jobs/start") async def start_job(body: dict): app.state.job_started = True app.state.start_payload = body return load_json("job_start.json") + @app.post("/api/v2/collie-job-queue/{job_id}/claim") + async def claim_managed_job(job_id: str, request: Request): + app.state.managed_events.append("claim") + app.state.collector_job_claim_requests.append( + {"job_id": job_id, "body": await request.body()} + ) + if app.state.collector_job_claim_statuses: + status = app.state.collector_job_claim_statuses.pop(0) + if status != 200: + return Response("raw-provider-secret", status_code=status) + if app.state.collector_job_stateful_claims: + if app.state.collector_job_active_id not in (None, job_id): + return Response("raw-provider-secret", status_code=409) + if ( + app.state.collector_job_active_id == job_id + and app.state.collector_job_claim_response["status"] == "running" + ): + return Response("raw-provider-secret", status_code=409) + app.state.collector_job_active_id = job_id + if app.state.collector_job_lost_claim_responses: + app.state.collector_job_lost_claim_responses -= 1 + # BHE has committed the assignment, but its response is lost. + raise requests.Timeout("raw-provider-secret") + return {"data": {"job": app.state.collector_job_claim_response}} + @app.post("/api/v2/jobs/end") async def end_job(body: dict): app.state.job_ended = True app.state.end_payload = body return load_json("job_end.json") + @app.post("/api/v2/collie-job-queue/{job_id}/end") + async def end_managed_job(job_id: str, body: dict): + app.state.managed_events.append("end") + app.state.collector_job_end_requests.append({"job_id": job_id, "body": body}) + status = ( + app.state.collector_job_end_statuses.pop(0) + if app.state.collector_job_end_statuses + else 200 + ) + return Response( + "raw-provider-secret" if status != 200 else "", status_code=status + ) + + @app.post("/api/v2/collie-job-queue/{job_id}/start") + async def start_collector_job(job_id: str, request: Request): + app.state.collector_job_start_requests.append((job_id, await request.body())) + if app.state.collector_job_start_statuses: + status = app.state.collector_job_start_statuses.pop(0) + if status != 200: + return Response("raw-provider-secret", status_code=status) + if app.state.collector_job_stateful_claims: + if app.state.collector_job_active_id != job_id: + return Response("raw-provider-secret", status_code=404) + job = app.state.collector_job_claim_response + if job["status"] != "running": + if job["status"] != "claimed" or datetime.fromisoformat( + job["claim_expires_at"] + ) <= datetime.now(UTC): + return Response("raw-provider-secret", status_code=409) + job["status"] = "running" + job["claim_expires_at"] = ( + datetime.now(UTC) + timedelta(minutes=15) + ).isoformat() + return Response(status_code=200) + + @app.post("/api/v2/collie-job-queue/{job_id}/heartbeat") + async def heartbeat_collector_job(job_id: str, request: Request): + app.state.collector_job_heartbeat_requests.append( + (job_id, await request.body()) + ) + return Response(status_code=200) + @app.post("/api/v2/ingest") async def ingest(request: Request): body = await request.body() @@ -97,6 +247,16 @@ async def ingest(request: Request): app.state.ingested_edges += len(validate_graph.graph.edges) return {"status": "success"} + @app.post("/api/v2/collie-job-queue/{job_id}/ingest") + async def ingest_managed_job(job_id: str, request: Request): + assert not app.state.collector_job_end_requests + assert request.headers["content-type"] == "application/json" + graph = Graph.model_validate_json(await request.body()) + app.state.managed_ingest_requests.append(job_id) + app.state.managed_ingested_nodes += len(graph.graph.nodes) + app.state.managed_ingested_edges += len(graph.graph.edges) + return Response(status_code=202) + @app.put("/api/v2/clients/update") async def update_client(body: dict): app.state.client_update_payload = body @@ -104,6 +264,7 @@ async def update_client(body: dict): @app.get("/api/v2/clients/management/available") async def management_available(): + app.state.management_available_requests += 1 return {"data": app.state.management_operations} @app.post("/api/v2/clients/management/start") @@ -201,6 +362,7 @@ def mock_service(mock_bloodhound_api, monkeypatch): class DummyExecutor: def __init__(self, *args, **kwargs): self.submitted = [] + self._processes = {} def submit(self, *args, **kwargs): future = Future() @@ -211,7 +373,13 @@ def shutdown(self, *args, **kwargs): return None def mock_request(method, url, **kwargs): - path = urlsplit(url).path + body = kwargs.get("data") + if body is not None and hasattr(body, "read"): + kwargs["data"] = body.read() + parsed_url = urlsplit(url) + path = parsed_url.path + if parsed_url.query: + path = f"{path}?{parsed_url.query}" if method.upper() == "GET": return mock_bloodhound_api.get(path) if method.upper() == "POST": @@ -223,12 +391,18 @@ def mock_request(method, url, **kwargs): monkeypatch.setattr("requests.request", mock_request) monkeypatch.setattr(scheduler_service, "ProcessPoolExecutor", DummyExecutor) + monkeypatch.setattr( + scheduler_service.openhound_logging, + "WorkerLogRouter", + lambda context: SimpleNamespace(queue=None, close=lambda: None), + ) return Service( bhe_uri="http://localhost:8000", token_key="test-key", token_id="test-id", collector_name="openhound-faker", + managed=False, ) @@ -251,6 +425,36 @@ def test_client_update_sends_metadata(mock_service, mock_bloodhound_api, monkeyp } +@pytest.mark.parametrize("managed", [False, True]) +def test_scheduler_start_routes_requests_by_mode( + mock_service, mock_bloodhound_api, monkeypatch, managed +): + mock_service.managed = managed + state = mock_bloodhound_api.app.state + monkeypatch.setattr(scheduler_service.signal, "signal", lambda *args: None) + + def next_poll(seconds): + assert seconds == mock_service.interval + mock_service.exit = True + if mock_service.future is not None: + mock_service.future.set_result(Result({}, mock_service.job_running)) + + monkeypatch.setattr(scheduler_service.time, "sleep", next_poll) + + mock_service.start() + + if managed: + assert state.client_update_payload is None + assert state.management_available_requests == 0 + assert len(state.collector_job_queue_requests) == 1 + assert state.jobs_available_requests == state.jobs_current_requests == 0 + else: + assert state.client_update_payload is not None + assert state.management_available_requests == 1 + assert state.collector_job_queue_requests == [] + assert state.jobs_available_requests == 1 + + def test_request_forwards_connect_and_read_timeouts(monkeypatch): captured = {} @@ -270,6 +474,67 @@ def mock_request(**kwargs): assert captured["timeout"] == (1.5, 3.5) +def test_managed_client_refreshes_credentials_and_retries_once_after_401(monkeypatch): + responses = iter( + [ + SimpleNamespace(status_code=401, text="expired"), + SimpleNamespace(status_code=204, text=""), + ] + ) + requests = [] + + def mock_request(**kwargs): + requests.append(kwargs) + return next(responses) + + refreshes = [] + + def refresh(): + refreshes.append(True) + return BHECredentials(token_id="new-id", token_key="new-key") + + monkeypatch.setattr(bloodhound.requests, "request", mock_request) + client = bloodhound_enterprise.BloodHoundEnterprise( + token_key="old-key", + token_id="old-id", + credential_refresh=refresh, + ) + + client.request(method="GET", path="/test", timeout=(10, 20)) + + assert refreshes == [True] + assert len(requests) == 2 + assert requests[0]["headers"]["Authorization"] == "bhesignature old-id" + assert requests[1]["headers"]["Authorization"] == "bhesignature new-id" + assert [request["timeout"] for request in requests] == [(10, 20), (10, 20)] + assert client.token_id == "new-id" + assert client.token_key == "new-key" + + +def test_managed_client_does_not_retry_a_second_401(monkeypatch): + def mock_request(**kwargs): + return SimpleNamespace(status_code=401, text="still expired") + + refreshes = [] + + def refresh(): + refreshes.append(True) + return BHECredentials(token_id="new-id", token_key="new-key") + + monkeypatch.setattr(bloodhound.requests, "request", mock_request) + client = bloodhound_enterprise.BloodHoundEnterprise( + token_key="old-key", + token_id="old-id", + credential_refresh=refresh, + ) + + with pytest.raises(BloodHoundHTTPError) as error: + client.request(method="GET", path="/test") + + assert error.value.code == 401 + assert refreshes == [True] + + def test_upload_artifact_part_uses_bounded_timeouts(mock_service, monkeypatch): captured = {} @@ -286,6 +551,25 @@ def capture_request(*args, **kwargs): ) +def test_available_collector_jobs_uses_bounded_timeouts(mock_service, monkeypatch): + captured = {} + + def capture_request(*args, **kwargs): + captured.update(kwargs) + return SimpleNamespace( + json=lambda: load_json("collector_jobs_available_empty.json") + ) + + monkeypatch.setattr(mock_service.client, "request", capture_request) + + mock_service.client.available_collector_jobs("openhound-faker") + + assert captured["timeout"] == ( + bloodhound_enterprise.MANAGED_COLLECTOR_JOB_AVAILABLE_CONNECT_TIMEOUT_SECONDS, + bloodhound_enterprise.MANAGED_COLLECTOR_JOB_AVAILABLE_READ_TIMEOUT_SECONDS, + ) + + def test_client_update_uses_unknown_when_hostname_lookup_fails( mock_service, mock_bloodhound_api, monkeypatch ): @@ -324,12 +608,240 @@ def raise_error(hostname: str): } +@pytest.mark.parametrize( + ("outcome", "failure_message", "metadata", "expected_body"), + [ + ( + ManagedJobOutcome.SUCCEEDED, + None, + None, + {"outcome": "succeeded"}, + ), + ( + ManagedJobOutcome.FAILED, + "Collection failed", + {"phase": "collect"}, + { + "outcome": "failed", + "failure_message": "Collection failed", + "metadata": {"phase": "collect"}, + }, + ), + ], +) +def test_end_managed_job_sends_managed_outcome( + mock_service, + mock_bloodhound_api, + outcome, + failure_message, + metadata, + expected_body, +): + mock_service.client.end_managed_job( + "collector-job-123", + outcome, + failure_message=failure_message, + metadata=metadata, + ) + + assert mock_bloodhound_api.app.state.collector_job_end_requests == [ + {"job_id": "collector-job-123", "body": expected_body} + ] + + +def test_start_and_heartbeat_collector_job_use_bodyless_requests( + mock_service, mock_bloodhound_api +): + job_id = "e0d43dbd-2cdf-4d68-9a81-2f95cd58a92f" + mock_service.client.start_collector_job(job_id) + mock_service.client.heartbeat_collector_job(job_id) + + assert mock_bloodhound_api.app.state.collector_job_start_requests == [(job_id, b"")] + assert mock_bloodhound_api.app.state.collector_job_heartbeat_requests == [ + (job_id, b"") + ] + + +def test_start_and_heartbeat_collector_job_use_bounded_timeouts( + mock_service, monkeypatch +): + requests_sent = [] + monkeypatch.setattr( + mock_service.client, + "request", + lambda *args, **kwargs: requests_sent.append(kwargs), + ) + + mock_service.client.start_collector_job("collector-job-123") + mock_service.client.heartbeat_collector_job("collector-job-123") + + assert [request["timeout"] for request in requests_sent] == [ + ( + bloodhound_enterprise.MANAGED_JOB_LIFECYCLE_CONNECT_TIMEOUT_SECONDS, + bloodhound_enterprise.MANAGED_JOB_LIFECYCLE_READ_TIMEOUT_SECONDS, + ) + ] * 2 + + +def test_end_managed_job_uses_bounded_timeouts(mock_service, monkeypatch): + captured = {} + + def capture_request(*args, **kwargs): + captured.update(kwargs) + + monkeypatch.setattr(mock_service.client, "request", capture_request) + + mock_service.client.end_managed_job( + "collector-job-123", ManagedJobOutcome.SUCCEEDED + ) + + assert captured["timeout"] == ( + bloodhound_enterprise.MANAGED_JOB_END_CONNECT_TIMEOUT_SECONDS, + bloodhound_enterprise.MANAGED_JOB_END_READ_TIMEOUT_SECONDS, + ) + + +@pytest.mark.parametrize( + "transient_error", + [ + BloodHoundHTTPError("server error", 500), + requests.ConnectionError("connection failed"), + requests.Timeout("request timed out"), + ], + ids=["http-500", "connection-error", "timeout"], +) +def test_end_managed_job_retries_transient_errors( + mock_service, monkeypatch, caplog, transient_error +): + attempts = 0 + delays = [] + + def flaky_request(*args, **kwargs): + nonlocal attempts + attempts += 1 + if attempts == 1: + raise transient_error + return object() + + monkeypatch.setattr(mock_service.client, "request", flaky_request) + monkeypatch.setattr(bloodhound_enterprise.time, "sleep", delays.append) + + with caplog.at_level( + logging.DEBUG, logger="openhound.core.clients.bloodhound_enterprise" + ): + mock_service.client.end_managed_job( + "collector-job-123", ManagedJobOutcome.SUCCEEDED + ) + + assert attempts == 2 + assert delays == [bloodhound_enterprise.MANAGED_JOB_END_RETRY_DELAY_SECONDS] + assert [ + record.levelno + for record in caplog.records + if record.getMessage().startswith( + ( + "Attempting to end managed collector job", + "Managed collector job collector-job-123 end attempt failed", + "Managed collector job collector-job-123 ended successfully", + ) + ) + ] == [logging.INFO, logging.INFO, logging.INFO, logging.INFO] + assert any( + record.levelno == logging.DEBUG + and getattr(record, "endpoint", None) + == "/api/v2/collie-job-queue/collector-job-123/end" + for record in caplog.records + ) + + +def test_end_managed_job_raises_after_transient_retries_are_exhausted( + mock_service, monkeypatch +): + attempts = 0 + delays = [] + + def unavailable(*args, **kwargs): + nonlocal attempts + attempts += 1 + raise BloodHoundHTTPError("temporarily unavailable", 503) + + monkeypatch.setattr(mock_service.client, "request", unavailable) + monkeypatch.setattr(bloodhound_enterprise.time, "sleep", delays.append) + + with pytest.raises(BloodHoundHTTPError) as error: + mock_service.client.end_managed_job( + "collector-job-123", ManagedJobOutcome.FAILED + ) + + assert error.value.code == 503 + assert attempts == 4 + assert delays == [bloodhound_enterprise.MANAGED_JOB_END_RETRY_DELAY_SECONDS] * ( + bloodhound_enterprise.MANAGED_JOB_END_MAX_ATTEMPTS - 1 + ) + + +def test_end_managed_job_does_not_retry_non_transient_client_error( + mock_service, monkeypatch +): + attempts = 0 + delays = [] + + def bad_request(*args, **kwargs): + nonlocal attempts + attempts += 1 + raise BloodHoundHTTPError("bad request", 400) + + monkeypatch.setattr(mock_service.client, "request", bad_request) + monkeypatch.setattr(bloodhound_enterprise.time, "sleep", delays.append) + + with pytest.raises(BloodHoundHTTPError) as error: + mock_service.client.end_managed_job( + "collector-job-123", ManagedJobOutcome.FAILED + ) + + assert error.value.code == 400 + assert attempts == 1 + assert delays == [] + + +@pytest.mark.parametrize( + "request_error", + [ + requests.exceptions.InvalidURL("invalid URL"), + requests.exceptions.TooManyRedirects("too many redirects"), + ], + ids=["invalid-url", "too-many-redirects"], +) +def test_end_managed_job_does_not_retry_non_transient_request_error( + mock_service, monkeypatch, request_error +): + attempts = 0 + delays = [] + + def invalid_request(*args, **kwargs): + nonlocal attempts + attempts += 1 + raise request_error + + monkeypatch.setattr(mock_service.client, "request", invalid_request) + monkeypatch.setattr(bloodhound_enterprise.time, "sleep", delays.append) + + with pytest.raises(type(request_error)): + mock_service.client.end_managed_job( + "collector-job-123", ManagedJobOutcome.FAILED + ) + + assert attempts == 1 + assert delays == [] + + def test_jobs_starts_new_job(mock_service, mock_bloodhound_api): """Runs _check_jobs and checks if the new job is started when available.""" job = mock_service.check_jobs() assert job is not None assert job.id == 123 + assert isinstance(job.id, int) assert mock_bloodhound_api.app.state.job_started is False @@ -340,17 +852,192 @@ def test_jobs_no_jobs_available(mock_service, mock_bloodhound_api): assert mock_service.check_jobs() is None +def test_managed_collector_job_check_requests_first_queue_job( + mock_service, mock_bloodhound_api +): + mock_bloodhound_api.app.state.collector_job_queue_response = load_json( + "collector_jobs_available_with_job.json" + ) + + job = mock_service.check_managed_collector_jobs() + + assert job is not None + assert job.id == UUID("11111111-1111-1111-1111-111111111111") + assert mock_bloodhound_api.app.state.collector_job_queue_requests == [ + { + "limit": "1", + "job_key": "eq:openhound-faker", + "sort_by": ["-priority", "created_at"], + } + ] + assert mock_bloodhound_api.app.state.jobs_available_requests == 0 + + +def test_collector_job_queue_response_parses_contract_fields( + mock_service, mock_bloodhound_api +): + payload = load_json("collector_jobs_available_with_jobs.json") + mock_bloodhound_api.app.state.collector_job_queue_response = payload + + response = mock_service.client.available_collector_jobs("openhound-faker") + job = response.data.jobs[0] + + assert response.count == 2 + assert response.skip == 0 + assert response.limit == 2 + assert job == CollectorJob.model_validate(payload["data"]["jobs"][0]) + assert job.id == UUID("11111111-1111-1111-1111-111111111111") + assert job.job_schedule_id is None + assert job.job_profile_id is None + assert job.job_type_id == 7 + assert isinstance(job.job_type_id, int) + assert job.job_key == "openhound-faker" + assert job.version == "v2" + assert job.params == { + "sources": { + "source": { + "random": { + "domains": ["example.com", "example.org"], + "include_deleted": False, + "nested": {"depth": 2}, + } + } + } + } + assert job.scope_client_id is None + assert job.external_secret_key_id is None + assert job.secret_key_id is None + assert job.priority == 10 + assert job.status == "ready" + assert job.run_at.isoformat() == "2026-02-20T10:00:00+00:00" + assert job.unclaimed_deadline_at.isoformat() == "2026-02-20T10:05:00+00:00" + assert job.attempts == 0 + assert job.max_attempts == 3 + assert job.last_failure is None + assert job.claimed_by is None + assert job.claimed_at is None + assert job.claim_expires_at is None + assert job.created_at.isoformat() == "2026-02-20T09:00:00+00:00" + assert job.updated_at.isoformat() == "2026-02-20T09:30:00+00:00" + + associated_job = response.data.jobs[1] + assert associated_job.external_secret_key_id == "external-key" + assert associated_job.id == UUID("22222222-2222-2222-2222-222222222222") + assert associated_job.job_schedule_id == UUID( + "44444444-4444-4444-4444-444444444444" + ) + assert associated_job.job_profile_id == UUID("55555555-5555-5555-5555-555555555555") + assert associated_job.job_type_id == 8 + assert isinstance(associated_job.job_type_id, int) + + +@pytest.mark.parametrize( + "field", + ["id", "job_schedule_id", "job_profile_id", "scope_client_id", "claimed_by"], +) +def test_collector_job_queue_response_rejects_invalid_uuids(field): + payload = load_json("collector_jobs_available_with_jobs.json") + + payload["data"]["jobs"][1][field] = "not-a-uuid" + + with pytest.raises(ValueError): + CollectorJobsAvailable.model_validate(payload) + + +@pytest.mark.parametrize( + "field", ["job_schedule_id", "job_profile_id", "scope_client_id", "claimed_by"] +) +def test_collector_job_queue_response_requires_nullable_uuid_fields(field): + payload = load_json("collector_jobs_available_with_jobs.json") + + del payload["data"]["jobs"][0][field] + + with pytest.raises(ValueError): + CollectorJobsAvailable.model_validate(payload) + + +@pytest.mark.parametrize("operation", ["available", "claim"]) +@pytest.mark.parametrize("field", ["job_schedule_id", "job_profile_id"]) +@pytest.mark.parametrize( + "value", ["not-a-uuid", 42, "42"], ids=["malformed", "integer", "numeric-string"] +) +def test_managed_queue_clients_reject_invalid_associations( + mock_service, mock_bloodhound_api, operation, field, value +): + state = mock_bloodhound_api.app.state + payload = load_json("collector_jobs_available_with_job.json") + job = payload["data"]["jobs"][0] + job[field] = value + state.collector_job_queue_response = payload + state.collector_job_claim_response = job + + if operation == "available": + with pytest.raises( + RuntimeError, match="^Invalid managed collector queue response$" + ): + mock_service.client.available_collector_jobs("openhound-faker") + assert len(state.collector_job_queue_requests) == 1 + else: + with pytest.raises( + RuntimeError, match="^Invalid managed collector claim response$" + ): + mock_service.client.claim_managed_job(job["id"]) + assert len(state.collector_job_claim_requests) == 1 + + +def test_managed_collector_job_check_returns_no_job_for_empty_queue( + mock_service, mock_bloodhound_api +): + assert mock_service.check_managed_collector_jobs() is None + assert len(mock_bloodhound_api.app.state.collector_job_queue_requests) == 1 + + +def test_queue_http_error_is_logged_and_re_raised( + mock_service, mock_bloodhound_api, caplog +): + mock_bloodhound_api.app.state.collector_job_queue_error_status = 503 + + with caplog.at_level(logging.ERROR), pytest.raises(BloodHoundHTTPError): + mock_service.check_managed_collector_jobs() + + queue_errors = [ + record + for record in caplog.records + if record.name == "openhound.core.clients.bloodhound_enterprise" + and record.levelno == logging.ERROR + and record.getMessage() == "Managed collector job queue request failed." + ] + assert len(queue_errors) == 1 + + def test_poll_starts_new_job(mock_service, mock_bloodhound_api, monkeypatch): """Similar to test_jobs_starts_new_job but using the poll method""" submitted = Future() - monkeypatch.setattr(mock_service.executor, "submit", lambda *args: submitted) + secret_requests = [] + monkeypatch.setattr( + mock_service.secrets_manager, "get_secret", secret_requests.append + ) + submitted_args = [] + + def submit(*args): + submitted_args.append(args) + return submitted + + monkeypatch.setattr(mock_service.executor, "submit", submit) mock_service._poll() assert mock_service.job_running == 123 + assert isinstance(mock_service.job_running, int) assert mock_service.future is submitted + assert submitted_args[0][1:] == ("openhound-faker", 123, None, None) + assert isinstance(submitted_args[0][2], int) assert mock_bloodhound_api.app.state.job_started is True assert mock_bloodhound_api.app.state.start_payload == {"id": 123} + assert mock_bloodhound_api.app.state.collector_job_queue_requests == [] + assert mock_bloodhound_api.app.state.collector_job_claim_requests == [] + assert mock_bloodhound_api.app.state.collector_job_end_requests == [] + assert secret_requests == [] def test_job_already_running(mock_service, monkeypatch): @@ -511,8 +1198,16 @@ def raise_error(self): mock_service._poll() -def test_scheduler_ingest_opengraph(mock_service, mock_bloodhound_api, monkeypatch): +def test_scheduler_ingest_opengraph( + mock_service, mock_bloodhound_api, monkeypatch, tmp_path +): """Run the DLT pipeline with the openhound-faker collector + check the amount of ingested nodes + edges""" + monkeypatch.setenv("DLT_DATA_DIR", str(tmp_path / ".dlt")) + monkeypatch.setattr( + scheduler_service.dataflow, + "Path", + lambda value: tmp_path if value == "/tmp/openhound" else Path(value), + ) monkeypatch.setenv( "DESTINATION__BLOODHOUNDENTERPRISE__URL", "http://localhost:8000" ) @@ -524,6 +1219,18 @@ def test_scheduler_ingest_opengraph(mock_service, mock_bloodhound_api, monkeypat assert result.job_id == 123 assert mock_bloodhound_api.app.state.ingested_nodes == 1000 assert mock_bloodhound_api.app.state.ingested_edges == 10000 + assert result.output_path is None + assert result.data_files == [] + + future = Future() + future.set_result(result) + mock_service.job_running = 123 + mock_service.future = future + mock_service._shutdown() + + assert mock_bloodhound_api.app.state.job_ended is True + assert mock_bloodhound_api.app.state.managed_ingest_requests == [] + assert mock_bloodhound_api.app.state.collector_job_end_requests == [] def _support_bundle_operation() -> dict: @@ -564,6 +1271,7 @@ def test_poll_prioritizes_management_over_a_new_job( assert len(sent) == 1 assert mock_bloodhound_api.app.state.job_started is False + assert mock_bloodhound_api.app.state.collector_job_queue_requests == [] def test_poll_starts_job_when_no_management_work( @@ -577,16 +1285,28 @@ def test_poll_starts_job_when_no_management_work( assert mock_bloodhound_api.app.state.job_started is True +@pytest.mark.parametrize("managed", [False, True]) def test_poll_still_checks_jobs_when_management_endpoint_fails( - mock_service, mock_bloodhound_api, monkeypatch + managed_service, mock_bloodhound_api, monkeypatch, caplog, managed ): - submitted = Future() - monkeypatch.setattr(mock_service, "check_management", lambda: 1 / 0) - monkeypatch.setattr(mock_service.executor, "submit", lambda *args: submitted) + managed_service.managed = managed - mock_service._poll() + def fail_management(): + raise BloodHoundHTTPError("management unavailable", 503) - assert mock_bloodhound_api.app.state.job_started is True + monkeypatch.setattr(managed_service, "check_management", fail_management) + with caplog.at_level(logging.ERROR): + managed_service._poll() + + state = mock_bloodhound_api.app.state + if managed: + assert state.managed_events == ["claim", "retrieve", "start_handoff"] + else: + assert state.job_started is True + assert ("management unavailable" in caplog.text) is (not managed) + assert any(record.exc_info is not None for record in caplog.records) is ( + not managed + ) def test_poll_does_not_start_a_job_when_management_work_fails( @@ -602,6 +1322,7 @@ def fail(operation): mock_service._poll() assert mock_bloodhound_api.app.state.job_started is False + assert mock_bloodhound_api.app.state.collector_job_queue_requests == [] def test_send_support_bundle_claims_uploads_completes_and_cleans_up( @@ -692,10 +1413,12 @@ def fail(*args, **kwargs): "status": ManagementOperationStatus.FAILED.value, } assert "Support bundle operation" in caplog.text + assert f"{failure_point} failed" in caplog.text + assert any(record.exc_info is not None for record in caplog.records) def test_send_support_bundle_retries_transient_part_upload_failure( - mock_service, monkeypatch, tmp_path + mock_service, monkeypatch, tmp_path, caplog ): log = tmp_path / "openhound.log" log.write_text("support log") @@ -710,7 +1433,7 @@ def flaky_request(method, path, **kwargs): if "/parts/" in path: attempts += 1 if attempts < 3: - raise BloodHoundHTTPError("temporary failure", 503) + raise BloodHoundHTTPError("temporary upload failure", 503) return original_request(method, path, **kwargs) monkeypatch.setattr(mock_service.client, "request", flaky_request) @@ -720,10 +1443,39 @@ def flaky_request(method, path, **kwargs): assert attempts == 3 assert delays == [2, 2] + assert "temporary upload failure" in caplog.text + assert any( + record.levelno == logging.WARNING and record.exc_info is not None + for record in caplog.records + ) + + +def test_support_bundle_retry_preserves_broad_request_exception_policy( + mock_service, monkeypatch +): + attempts = 0 + delays = [] + + def request(): + nonlocal attempts + attempts += 1 + if attempts == 1: + raise requests.exceptions.InvalidURL("invalid URL") + return "ok" + + monkeypatch.setattr(bloodhound_enterprise.time, "sleep", delays.append) + + result = mock_service.client._retry_support_bundle_request( + "test support bundle request", request + ) + + assert result == "ok" + assert attempts == 2 + assert delays == [bloodhound_enterprise.SUPPORT_BUNDLE_RETRY_DELAY_SECONDS] def test_send_support_bundle_fails_after_transient_retries_are_exhausted( - mock_service, mock_bloodhound_api, monkeypatch, tmp_path + mock_service, mock_bloodhound_api, monkeypatch, tmp_path, caplog ): log = tmp_path / "openhound.log" log.write_text("support log") @@ -752,6 +1504,8 @@ def unavailable_part_upload(method, path, **kwargs): "operation_id": operation.id, "status": ManagementOperationStatus.FAILED.value, } + assert "temporarily unavailable" in caplog.text + assert any(record.exc_info is not None for record in caplog.records) def test_collection_logs_include_extension_and_openhound_versions(monkeypatch, caplog): @@ -776,7 +1530,9 @@ def test_collection_logs_include_extension_and_openhound_versions(monkeypatch, c collection_records = [ record for record in caplog.records - if record.getMessage().startswith(("Subprocess running collection", "Collection for job")) + if record.getMessage().startswith( + ("Subprocess running collection", "Collection for job") + ) ] assert len(collection_records) == 2 for record in collection_records: @@ -784,3 +1540,3324 @@ def test_collection_logs_include_extension_and_openhound_versions(monkeypatch, c assert record.collector_extension_version == "2.3.4" assert record.openhound_version == "4.5.6" assert record.job_id == 42 + + +@pytest.mark.parametrize("collection_error", [None, RuntimeError("collection failed")]) +@pytest.mark.parametrize( + "report_error", + [BloodHoundHTTPError("unauthorized", 401), RuntimeError("refresh failed")], +) +def test_completed_job_retries_reporting_without_losing_outcome( + mock_service, monkeypatch, collection_error, report_error +): + future = Future() + if collection_error is None: + future.set_result(Result(results={"collect": ["a"]}, job_id=123)) + expected = JobStatus.COMPLETE + else: + future.set_exception(collection_error) + expected = JobStatus.FAILED + mock_service.future = future + mock_service.job_running = 123 + reports = [] + + def end_job(status, message): + reports.append((status, message)) + if len(reports) == 1: + raise report_error + + monkeypatch.setattr(mock_service.client, "end_job", end_job) + monkeypatch.setattr(mock_service, "check_jobs", lambda: None) + + mock_service._poll() + + assert mock_service.future is future + assert mock_service.job_running == 123 + assert [status for status, _ in reports] == [expected] + + mock_service._poll() + + assert mock_service.future is None + assert mock_service.job_running is None + assert reports == [reports[0], reports[0]] + + +def test_managed_refresh_rejects_invalid_id_and_recovers(mock_service, monkeypatch): + from openhound.core.clients.bhe_credentials import ( + AWSBHECredentials, + InvalidBHECredentialsSecret, + ) + + secret = {"token_id": "invalid-id", "token_key": "new-key"} + provider = AWSBHECredentials("bhe", SimpleNamespace(get_secret=lambda _: secret)) + client = mock_service.client + client._credential_refresh = provider.refresh + old_pair = (client.token_id, client.token_key) + attempts = [] + valid_id = "12345678-1234-1234-1234-123456789abc" + + def request(self, **kwargs): + attempts.append((self.token_id, self.token_key)) + if self.token_id != valid_id: + raise BloodHoundHTTPError("unauthorized", 401) + return "ok" + + monkeypatch.setattr(bloodhound.BloodHound, "request", request) + + with pytest.raises(InvalidBHECredentialsSecret): + client.request("GET", "/test") + + assert (client.token_id, client.token_key) == old_pair + assert attempts == [old_pair] + + secret["token_id"] = valid_id + assert client.request("GET", "/test") == "ok" + assert attempts == [old_pair, old_pair, (valid_id, "new-key")] + + +@pytest.fixture +def managed_service(mock_service, mock_bloodhound_api): + state = mock_bloodhound_api.app.state + state.collector_job_queue_response = load_json( + "collector_jobs_available_with_job.json" + ) + # Discovery is intentionally stale. Only the claimed snapshot is authoritative. + state.collector_job_queue_response["data"]["jobs"][0]["secret_key_id"] = ( + "stale-reference" + ) + claimed_at = datetime.now(UTC) + state.collector_job_claim_response.update( + secret_key_id="private-secret-reference", + status="claimed", + params=CLAIMED_PARAMS, + claimed_by="33333333-3333-3333-3333-333333333333", + claimed_at=claimed_at.isoformat(), + claim_expires_at=(claimed_at + timedelta(minutes=15)).isoformat(), + ) + mock_service.managed = True + mock_service.handoffs = [] + mock_service.secret_requests = [] + + def get_secret_value(*, SecretId): + state.managed_events.append("retrieve") + mock_service.secret_requests.append(SecretId) + return {"SecretString": json.dumps(COLLECTOR_SECRETS)} + + mock_service.secrets_manager = AWSSecretsManager( + SimpleNamespace(get_secret_value=get_secret_value) + ) + + def start_handoff(prepared): + state.managed_events.append("start_handoff") + mock_service.handoffs.append(prepared) + + mock_service.managed_job_runner = start_handoff + return mock_service + + +@pytest.mark.parametrize("has_schedule", [False, True], ids=["no-schedule", "schedule"]) +@pytest.mark.parametrize("has_profile", [False, True], ids=["no-profile", "profile"]) +def test_claim_request_is_bodyless_and_parses_data_job( + managed_service, mock_bloodhound_api, monkeypatch, has_schedule, has_profile +): + state = mock_bloodhound_api.app.state + associated_job = load_json("collector_jobs_available_with_jobs.json")["data"][ + "jobs" + ][1] + schedule_id = associated_job["job_schedule_id"] if has_schedule else None + profile_id = associated_job["job_profile_id"] if has_profile else None + state.collector_job_claim_response.update( + job_schedule_id=schedule_id, job_profile_id=profile_id + ) + captured = [] + original = managed_service.client.request + + def request(**kwargs): + captured.append(kwargs) + return original(**kwargs) + + monkeypatch.setattr(managed_service.client, "request", request) + job_id = mock_bloodhound_api.app.state.collector_job_claim_response["id"] + claimed = managed_service.client.claim_managed_job(job_id) + + assert claimed == CollectorJob.model_validate( + mock_bloodhound_api.app.state.collector_job_claim_response + ) + assert claimed.id == UUID(job_id) + assert claimed.job_schedule_id == (UUID(schedule_id) if schedule_id else None) + assert claimed.job_profile_id == (UUID(profile_id) if profile_id else None) + assert claimed.job_type_id == 7 + assert isinstance(claimed.job_type_id, int) + assert claimed.version == "v2" + assert claimed.params == CLAIMED_PARAMS + assert captured == [ + { + "method": "POST", + "path": f"/api/v2/collie-job-queue/{job_id}/claim", + "timeout": ( + bloodhound_enterprise.MANAGED_JOB_CLAIM_CONNECT_TIMEOUT_SECONDS, + bloodhound_enterprise.MANAGED_JOB_CLAIM_READ_TIMEOUT_SECONDS, + ), + } + ] + assert mock_bloodhound_api.app.state.collector_job_claim_requests == [ + {"job_id": job_id, "body": b""} + ] + + +def test_managed_poll_claims_retrieves_and_hands_off_authoritative_job( + managed_service, mock_bloodhound_api, caplog +): + with caplog.at_level(logging.DEBUG): + managed_service._poll() + + state = mock_bloodhound_api.app.state + assert state.managed_events == ["claim", "retrieve", "start_handoff"] + assert managed_service.secret_requests == ["private-secret-reference"] + (prepared,) = managed_service.handoffs + assert prepared.job.params == CLAIMED_PARAMS + assert prepared.job.status == "claimed" + assert prepared.credentials == COLLECTOR_SECRETS + assert state.job_started is False + assert state.jobs_available_requests == state.jobs_current_requests == 0 + assert state.collector_job_end_requests == [] + for message in ("Job successfully picked up", "Credential retrieval valid"): + assert any( + r.levelno == logging.INFO and r.getMessage() == message + for r in caplog.records + ) + for forbidden in ( + "private-secret-reference", + "private-credential-value", + "test-key", + "test-id", + ): + assert forbidden not in caplog.text + assert forbidden not in repr(prepared) + assert all(forbidden not in repr(r.__dict__) for r in caplog.records) + + +def test_managed_job_retrieves_full_hierarchy_secret_before_dlt_resolution( + managed_service, mock_bloodhound_api, monkeypatch, caplog +): + service = managed_service + bhe_credentials = (service.client.token_id, service.client.token_key) + state = mock_bloodhound_api.app.state + data_dir = Path(__file__).parent / "test_data" + params = json.loads((data_dir / "managed_aws_configuration.json").read_text()) + secret_document = (data_dir / "managed_aws_secrets.json").read_text() + service.collector_name = "aws" + state.collector_job_queue_response["data"]["jobs"][0]["job_key"] = "aws" + state.collector_job_claim_response.update(job_key="aws", params=params) + requests_seen = [] + + def get_secret_value(*, SecretId): + requests_seen.append(SecretId) + return {"SecretString": secret_document} + + service.secrets_manager = AWSSecretsManager( + SimpleNamespace(get_secret_value=get_secret_value) + ) + service.managed_job_runner = service._run_prepared_managed_job + collector = SimpleNamespace( + name="aws", dlt_source=SimpleNamespace(), package_version="test", metadata=None + ) + monkeypatch.setattr( + scheduler_service.CollectorManager, + "from_entrypoint", + lambda: SimpleNamespace(collectors=[collector]), + ) + resolved = {} + + def collect(extension, output_path): + resolved.update( + profiles=dlt.secrets.get( + "sources.source.aws.credentials.profile_names", list + ), + region=dlt.secrets.get("sources.source.aws.credentials.region_name", str), + max_objects=dlt.config.get( + "sources.source.aws.bucket_policy.max_objects", int + ), + ) + return {"collect": []} + + monkeypatch.setattr(scheduler_service.dataflow, "pipeline", collect) + # This test verifies configuration handoff, not existing AWS profiles. + monkeypatch.setattr( + scheduler_service, "validate_configured_aws_profiles", lambda: None + ) + + def submit(function, *args): + future = Future() + future.set_result(function(*args)) + return future + + monkeypatch.setattr(service.executor, "submit", submit) + with caplog.at_level(logging.DEBUG): + service._poll() + service._handle_completed_job(service.future) + + assert requests_seen == ["private-secret-reference"] + assert resolved == { + "profiles": ["retrieved-profile-a", "retrieved-profile-b"], + "region": "us-west-2", + "max_objects": 50, + } + assert (service.client.token_id, service.client.token_key) == bhe_credentials + assert state.collector_job_end_requests[0]["body"] == {"outcome": "succeeded"} + for credential in ("retrieved-profile-a", "retrieved-profile-b"): + assert credential not in caplog.text + + +AWS_ACCESS_KEY = "AKIAIOSFODNN7EXAMPLE" +AWS_SECRET_KEY = "wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY" + + +@pytest.mark.parametrize( + "case", ["single-denied", "single-success", "mixed", "all-denied"] +) +def test_existing_aws_profiles_report_authentication_outcome_through_queue( + managed_service, mock_bloodhound_api, case, monkeypatch, tmp_path, caplog +): + """Exercise discovery/claim/start/end against the local API and real SDK files.""" + scheduler = managed_service + scheduler.collector_name = "aws" + scheduler.managed_job_runner = scheduler._run_prepared_managed_job + state = mock_bloodhound_api.app.state + state.collector_job_queue_response["data"]["jobs"][0]["job_key"] = "aws" + config, credentials = tmp_path / "config", tmp_path / "credentials" + config.write_text( + "[profile good]\nregion=eu-west-1\n[profile denied]\nregion=eu-west-1\n" + ) + contents = "".join( + f"[{profile}]\naws_access_key_id={AWS_ACCESS_KEY}\naws_secret_access_key={AWS_SECRET_KEY}\n" + for profile in ("good", "denied") + ) + credentials.write_text(contents) + options = {"config_file": str(config), "credentials_file": str(credentials)} + if case in ("mixed", "all-denied"): + options["profile_names"] = ["denied", "good"] if case == "mixed" else ["denied"] + else: + options["profile_name"] = "good" if case == "single-success" else "denied" + document = {"sources": {"source": {"aws": {"credentials": options}}}} + scheduler.secrets_manager = AWSSecretsManager( + SimpleNamespace( + get_secret_value=lambda **kwargs: {"SecretString": json.dumps(document)} + ) + ) + state.collector_job_claim_response.update( + job_key="aws", + external_secret_key_id=None, + params={ + "sources": { + "source": {"aws": {"credentials": {"region_name": "eu-west-1"}}} + } + }, + ) + monkeypatch.setenv("AWS_ACCESS_KEY_ID", "scheduler-key") + monkeypatch.setenv("AWS_SECRET_ACCESS_KEY", "scheduler-secret") + monkeypatch.setenv("AWS_EC2_METADATA_DISABLED", "true") + monkeypatch.setattr( + URLLib3Session, "send", lambda *a, **kw: pytest.fail("Live AWS request") + ) + environment = dict(os.environ) + create_client = BotocoreSession.create_client + identities, stubbers, stages = [], [], [] + + def stub_identity(session, name, *args, **kwargs): + client = create_client(session, name, *args, **kwargs) + assert name == "sts" + profile = session.get_config_variable("profile") + assert profile in ("good", "denied") + assert client._request_signer._credentials.method == "shared-credentials-file" + assert client._request_signer._credentials.access_key == AWS_ACCESS_KEY + identities.append(profile) + stubber = Stubber(client) + if profile == "denied": + stubber.add_client_error( + "get_caller_identity", + service_error_code="InvalidClientTokenId", + expected_params={}, + ) + else: + stubber.add_response( + "get_caller_identity", + { + "Account": "123456789012", + "Arn": "arn:aws:iam::123456789012:user/collector", + "UserId": "collector", + }, + {}, + ) + stubber.activate() + stubbers.append(stubber) + return client + + monkeypatch.setattr(BotocoreSession, "create_client", stub_identity) + + def collect(**kwargs): + stages.append("collect") + return SimpleNamespace(load_packages=[]) + + collector = SimpleNamespace( + name="aws", + dlt_source=object(), + package_version="test", + metadata=None, + collector=collect, + preprocessor=object(), + ) + monkeypatch.setattr( + scheduler_service.CollectorManager, + "from_entrypoint", + lambda: SimpleNamespace(collectors=[collector]), + ) + monkeypatch.setattr( + scheduler_service.dataflow, + "preprocess", + lambda extension: stages.append("preprocess") or [], + ) + + def convert(extension, output_path): + stages.append("convert") + (output_path / "graph.json").write_text(Graph(graph={}).model_dump_json()) + return [] + + monkeypatch.setattr(scheduler_service.dataflow, "convert", convert) + + def submit(function, *args): + future = Future() + try: + future.set_result(function(*args)) + except Exception as error: # noqa: BLE001 - emulate ProcessPoolExecutor + future.set_exception(error) + return future + + monkeypatch.setattr(scheduler.executor, "submit", submit) + failed = case in ("single-denied", "all-denied") + try: + scheduler._poll() + if failed: + with pytest.raises(ManagedConfigurationError): + scheduler.future.result() + assert stages == [] + else: + assert scheduler.future.result().job_id == scheduler.job_running + assert stages == ["collect", "preprocess", "convert"] + state.collector_job_queue_response = load_json( + "collector_jobs_available_empty.json" + ) + scheduler._poll() + assert len(state.collector_job_start_requests) == 1 + assert len(state.collector_job_end_requests) == 1 + assert state.collector_job_end_requests[0]["body"]["outcome"] == ( + "failed" if failed else "succeeded" + ) + assert len(state.managed_ingest_requests) == (0 if failed else 1) + assert identities == ( + ["denied", "good"] if case == "mixed" else ["denied" if failed else "good"] + ) + assert credentials.read_text() == contents + assert scheduler._aws_profile_workspace is None + assert scheduler.future is scheduler.job_running is None + assert dict(os.environ) == environment + assert AWS_ACCESS_KEY not in caplog.text and AWS_SECRET_KEY not in caplog.text + for stubber in stubbers: + stubber.assert_no_pending_responses() + finally: + for stubber in stubbers: + stubber.deactivate() + + +@pytest.fixture +def aws_managed_service(managed_service, mock_bloodhound_api): + state = mock_bloodhound_api.app.state + managed_service.collector_name = "aws" + state.collector_job_queue_response["data"]["jobs"][0].update( + job_key="aws", external_secret_key_id="stale-discovery-key" + ) + state.collector_job_claim_response.update( + job_key="aws", + external_secret_key_id=AWS_ACCESS_KEY, + params={ + "sources": { + "source": {"aws": {"credentials": {"region_name": "us-west-2"}}} + } + }, + ) + managed_service.secrets_manager = AWSSecretsManager( + SimpleNamespace( + get_secret_value=lambda **kwargs: {"SecretString": AWS_SECRET_KEY} + ) + ) + managed_service.managed_job_runner = managed_service._run_prepared_managed_job + return managed_service + + +@pytest.mark.parametrize( + "failed", [False, True], ids=["successful-collection", "failed-collection"] +) +@pytest.mark.parametrize( + "binary_secret", [False, True], ids=["secret-string", "bhe-utf8-binary"] +) +def test_active_aws_poll_uses_claimed_key_and_raw_secret_in_worker_adapter( + aws_managed_service, + mock_bloodhound_api, + failed, + binary_secret, + monkeypatch, + tmp_path, + caplog, +): + scheduler = aws_managed_service + state = mock_bloodhound_api.app.state + # Secrets Manager continues using the scheduler's identity. Only the worker + # receives the customer's pair and changes its SDK file/config providers. + client = boto3.Session( + aws_access_key_id="scheduler-key", + aws_secret_access_key="scheduler-secret", + region_name="us-east-1", + ).client("secretsmanager") + scheduler.secrets_manager = AWSSecretsManager(client) + handoffs, seen_files = [], [] + runner = scheduler.managed_job_runner + + def handoff(prepared): + handoffs.append(prepared) + assert prepared.job.external_secret_key_id == AWS_ACCESS_KEY + assert isinstance(prepared.credentials, AWSAccessKeyCredentials) + assert prepared.credentials.secret_access_key == AWS_SECRET_KEY + runner(prepared) + + scheduler.managed_job_runner = handoff + + create_client = BotocoreSession.create_client + identity_stubbers = [] + + def stub_identity(session, name, *args, **kwargs): + sts = create_client(session, name, *args, **kwargs) + if name == "sts": + assert sts._request_signer._credentials.access_key == AWS_ACCESS_KEY + stubber = Stubber(sts) + stubber.add_response( + "get_caller_identity", + { + "Account": "123456789012", + "Arn": "arn:aws:iam::123456789012:user/collector", + "UserId": "collector", + }, + {}, + ) + stubber.activate() + identity_stubbers.append(stubber) + return sts + + monkeypatch.setattr(BotocoreSession, "create_client", stub_identity) + + def load_collectors(): + assert state.collector_job_start_requests == [ + (state.collector_job_claim_response["id"], b"") + ] + settings = dlt.secrets.get("sources.source.aws.credentials", dict) + assert settings["profile_name"].startswith("openhound-job-") + assert settings["region_name"] == "us-west-2" + credentials = boto3.Session( + profile_name=settings["profile_name"], region_name=settings["region_name"] + ).get_credentials() + assert credentials.method == "shared-credentials-file" + frozen = credentials.get_frozen_credentials() + assert (frozen.access_key, frozen.secret_key, frozen.token) == ( + AWS_ACCESS_KEY, + AWS_SECRET_KEY, + None, + ) + seen_files.append(Path(os.environ["AWS_SHARED_CREDENTIALS_FILE"])) + return SimpleNamespace( + collectors=[ + SimpleNamespace( + name="aws", + dlt_source=object(), + package_version="test", + metadata=None, + ) + ] + ) + + monkeypatch.setattr( + scheduler_service.CollectorManager, "from_entrypoint", load_collectors + ) + + def collect(extension, output_path): + assert not seen_files[0].is_relative_to(output_path) + if failed: + logging.getLogger("openhound_aws.auth").error( + "Provider failure: %s", AWS_SECRET_KEY + ) + raise RuntimeError(AWS_SECRET_KEY) + (output_path / "graph.json").write_text(Graph(graph={}).model_dump_json()) + return {"collect": ["account"], "convert": ["graph"]} + + monkeypatch.setattr(scheduler_service.dataflow, "pipeline", collect) + + def submit(function, *args): + future = Future() + try: + future.set_result(function(*args)) + except Exception as error: # noqa: BLE001 - emulate ProcessPoolExecutor + future.set_exception(error) + return future + + monkeypatch.setattr(scheduler.executor, "submit", submit) + environment = dict(os.environ) + with Stubber(client) as stubber, caplog.at_level(logging.DEBUG): + stubber.add_response( + "get_secret_value", + ( + {"SecretBinary": AWS_SECRET_KEY.encode()} + if binary_secret + else {"SecretString": AWS_SECRET_KEY} + ), + {"SecretId": "private-secret-reference"}, + ) + scheduler._poll() + directory = scheduler._aws_profile_workspace.directory + assert not seen_files[0].exists() # Worker cleanup precedes parent reporting. + assert directory.exists() + if not failed: + result = scheduler.future.result() + assert AWS_SECRET_KEY not in repr(result) + assert all( + AWS_SECRET_KEY not in path.read_text() for path in result.data_files + ) + scheduler._handle_completed_job(scheduler.future) + stubber.assert_no_pending_responses() + assert len(handoffs) == 1 + assert not directory.exists() + assert scheduler._aws_profile_workspace is None + assert dict(os.environ) == environment + assert client._request_signer._credentials.access_key == "scheduler-key" + assert state.collector_job_end_requests[0]["body"]["outcome"] == ( + "failed" if failed else "succeeded" + ) + assert len(state.managed_ingest_requests) == (0 if failed else 1) + assert AWS_SECRET_KEY not in repr(state.collector_job_end_requests) + assert AWS_SECRET_KEY not in caplog.text + assert AWS_SECRET_KEY not in repr(handoffs) + assert all(AWS_SECRET_KEY not in repr(record.__dict__) for record in caplog.records) + assert len(identity_stubbers) == 1 + for stubber in identity_stubbers: + stubber.assert_no_pending_responses() + stubber.deactivate() + + +@pytest.mark.parametrize( + "field,value", + [ + ("external_secret_key_id", None), + ("secret_key_id", None), + ("secret_key_id", ""), + ("secret_key_id", " "), + ], +) +def test_invalid_aws_metadata_fails_before_start_without_runtime_fallback( + aws_managed_service, mock_bloodhound_api, field, value, monkeypatch, caplog +): + monkeypatch.setenv("AWS_PROFILE", "local-profile") + monkeypatch.setenv("AWS_ACCESS_KEY_ID", "scheduler-key") + monkeypatch.setenv("AWS_SECRET_ACCESS_KEY", "scheduler-secret") + state = mock_bloodhound_api.app.state + state.collector_job_claim_response[field] = value + requests_seen = [] + get_secret = aws_managed_service.secrets_manager.get_secret + + def retrieve(reference): + requests_seen.append(reference) + return get_secret(reference) + + monkeypatch.setattr(aws_managed_service.secrets_manager, "get_secret", retrieve) + with caplog.at_level(logging.DEBUG): + aws_managed_service._poll() + # A null key ID permits a legacy document, whose shape requires retrieval. + assert requests_seen == ( + ["private-secret-reference"] + if field == "external_secret_key_id" and value is None + else [] + ) + assert state.collector_job_start_requests == [] + assert aws_managed_service.executor.submitted == [] + assert aws_managed_service._aws_profile_workspace is None + assert state.collector_job_end_requests[0]["body"]["outcome"] == "failed" + assert AWS_SECRET_KEY not in caplog.text + assert all(record.exc_info is None for record in caplog.records) + + +@pytest.mark.parametrize( + "response", + [ + {}, + {"SecretString": json.dumps({"aws_secret_access_key": AWS_SECRET_KEY})}, + {"SecretString": json.dumps(COLLECTOR_SECRETS)}, + {"SecretBinary": b"\xff\xfe"}, + ], +) +def test_invalid_raw_aws_payloads_are_not_trimmed_wrapped_or_treated_as_documents( + aws_managed_service, mock_bloodhound_api, response, caplog +): + aws_managed_service.secrets_manager = AWSSecretsManager( + SimpleNamespace(get_secret_value=lambda **kwargs: response) + ) + with caplog.at_level(logging.DEBUG): + aws_managed_service._poll() + state = mock_bloodhound_api.app.state + assert state.collector_job_start_requests == [] + assert aws_managed_service.executor.submitted == [] + assert state.collector_job_end_requests[0]["body"]["outcome"] == "failed" + assert AWS_SECRET_KEY not in repr(state.collector_job_end_requests) + assert AWS_SECRET_KEY not in caplog.text + + +@pytest.mark.parametrize( + "options", + [ + {"profile_name": "local"}, + {"profile_names": ["local"]}, + {"profile_names": []}, + {"auth_mode": "workload"}, + {"authentication_mode": "profile"}, + {"aws_session_token": "unsupported-token"}, + {"aws_access_key_id": AWS_ACCESS_KEY}, + {"source_account": "999999999999"}, + {"org_role_name": ""}, + {"org_role_name": 123}, + {"role_arn": "arn:aws:iam::123456789012:user/Collector"}, + {"role_arn": "arn:aws:iam::123456789:role/Collector"}, + {"role_arn": "arn:aws:iam::123456789012:role/Collector\n"}, + { + "role_arn": "arn:aws:iam::123456789012:role/Collector", + "role_session_name": "x", + }, + { + "role_arn": "arn:aws:iam::123456789012:role/Collector", + "role_session_name": "x" * 65, + }, + { + "role_arn": "arn:aws:iam::123456789012:role/Collector", + "role_session_name": "session\n[profile other]", + }, + { + "role_arn": "arn:aws:iam::123456789012:role/Collector", + "role_session_name": None, + }, + { + "role_arn": "arn:aws:iam::123456789012:role/Collector", + "external_id": "contains spaces", + }, + { + "role_arn": "arn:aws:iam::123456789012:role/Collector", + "external_id": "x" * 1225, + }, + { + "role_arn": "arn:aws:iam::123456789012:role/Collector", + "external_id": 123, + }, + { + "role_arn": "arn:aws:iam::123456789012:role/Collector", + "source_account": "123456789012", + }, + ], +) +def test_raw_aws_settings_reach_worker_handoff_without_aws_validation( + aws_managed_service, mock_bloodhound_api, options, monkeypatch +): + state = mock_bloodhound_api.app.state + state.collector_job_claim_response["params"] = { + "sources": {"source": {"aws": {"credentials": options}}} + } + monkeypatch.setattr( + aws_managed_service.secrets_manager, + "get_secret", + lambda _: AWS_SECRET_KEY, + ) + aws_managed_service._poll() + assert len(state.collector_job_start_requests) == 1 + submitted = aws_managed_service.executor.submitted[0][0] + assert submitted[3] == state.collector_job_claim_response["params"] + assert submitted[3]["sources"]["source"]["aws"]["credentials"] == options + assert submitted[4].secret_access_key == AWS_SECRET_KEY + assert state.collector_job_end_requests == [] + aws_managed_service._cleanup_aws_profile_workspace() + + +@pytest.mark.parametrize( + "options", + [ + {"role_session_name": "collection-session"}, + {"external_id": "customer-123"}, + {"role_session_name": "collection-session", "external_id": "customer-123"}, + {"role_arn": None}, + {"role_arn": ""}, + {"role_arn": " \t"}, + {"role_arn": False}, + {"role_arn": 123}, + {"role_arn": []}, + {"role_arn": {}}, + {"role_arn": "", "role_session_name": "collection-session"}, + {"role_arn": None, "external_id": "customer-123"}, + ], +) +def test_incomplete_aws_role_settings_fail_before_secret_retrieval_and_start( + aws_managed_service, mock_bloodhound_api, options, monkeypatch, caplog +): + state = mock_bloodhound_api.app.state + state.collector_job_claim_response["params"] = { + "sources": {"source": {"aws": {"credentials": options}}} + } + monkeypatch.setattr( + aws_managed_service.secrets_manager, + "get_secret", + lambda _: pytest.fail("Incomplete role settings must not retrieve credentials"), + ) + aws_managed_service._poll() + assert "Managed AWS role settings require a non-empty role_arn" in caplog.text + assert state.collector_job_start_requests == [] + assert aws_managed_service.executor.submitted == [] + assert aws_managed_service._aws_profile_workspace is None + assert state.collector_job_end_requests[0]["body"]["outcome"] == "failed" + + +def test_missing_new_bhe_job_field_rejects_claim_before_start( + aws_managed_service, mock_bloodhound_api +): + state = mock_bloodhound_api.app.state + del state.collector_job_claim_response["external_secret_key_id"] + with pytest.raises(ManagedJobProcessingError): + aws_managed_service._poll() + assert state.collector_job_start_requests == [] + assert aws_managed_service.executor.submitted == [] + + +def test_raw_aws_key_secret_not_reinterpreted_for_other_collectors( + managed_service, mock_bloodhound_api +): + managed_service.secrets_manager = AWSSecretsManager( + SimpleNamespace( + get_secret_value=lambda **kwargs: {"SecretString": AWS_SECRET_KEY} + ) + ) + managed_service._poll() + assert managed_service.handoffs == [] + assert ( + mock_bloodhound_api.app.state.collector_job_end_requests[0]["body"]["outcome"] + == "failed" + ) + + +def test_aws_missing_secret_reports_sanitized_failure_before_start( + aws_managed_service, mock_bloodhound_api, caplog +): + def missing_secret(**kwargs): + raise ClientError( + { + "Error": { + "Code": "ResourceNotFoundException", + "Message": "Secret does not exist", + } + }, + "GetSecretValue", + ) + + aws_managed_service.secrets_manager = AWSSecretsManager( + SimpleNamespace(get_secret_value=missing_secret) + ) + with caplog.at_level(logging.DEBUG): + aws_managed_service._poll() + state = mock_bloodhound_api.app.state + assert state.collector_job_start_requests == [] + assert aws_managed_service.executor.submitted == [] + assert state.collector_job_end_requests[0]["body"] == { + "outcome": "failed", + "failure_message": "Credential secret not found", + } + assert AWS_SECRET_KEY not in caplog.text + + +@pytest.mark.parametrize("error_type", [RuntimeError, SecretPermissionError]) +def test_aws_preparation_does_not_expose_secret_reader_payloads( + aws_managed_service, mock_bloodhound_api, monkeypatch, caplog, error_type +): + def fail_reader(reference): + error = error_type(f"AccessDenied for {AWS_ACCESS_KEY}: {AWS_SECRET_KEY}") + if isinstance(error, SecretPermissionError): + error.context.update( + aws_error_code=AWS_SECRET_KEY, RequestId=AWS_ACCESS_KEY + ) + raise error + + monkeypatch.setattr(aws_managed_service.secrets_manager, "get_secret", fail_reader) + with caplog.at_level(logging.DEBUG): + aws_managed_service._poll() + state = mock_bloodhound_api.app.state + assert state.collector_job_start_requests == [] + assert aws_managed_service.executor.submitted == [] + assert state.collector_job_end_requests[0]["body"] == { + "outcome": "failed", + "failure_message": ( + "AWS secret access denied" + if error_type is SecretPermissionError + else "Credential secret retrieval failed" + ), + } + assert error_type.__name__ in caplog.text + for value in (AWS_ACCESS_KEY, AWS_SECRET_KEY): + assert value not in caplog.text + assert value not in repr(state.collector_job_end_requests) + assert all(value not in repr(record.__dict__) for record in caplog.records) + + +@pytest.mark.parametrize("level", [logging.INFO, logging.DEBUG]) +@pytest.mark.parametrize("raw_key", [False, True], ids=["nested-document", "raw-key"]) +@pytest.mark.parametrize( + "failure,expected,summary", + [ + ("region", "AWS_DEFAULT_REGION", "AWS credentials unavailable"), + ("credentials", "NoCredentialsError", "AWS credentials unavailable"), + ( + "refresh", + "refreshed credentials are still expired", + "AWS credentials unavailable", + ), + ( + "ExpiredTokenException", + "ExpiredTokenException", + "AWS credentials unavailable", + ), + ("AccessDeniedException", "secret permissions", "AWS secret access denied"), + ("kms", "kms:Decrypt", "AWS secret access denied"), + ("DecryptionFailure", "KMS key permissions", "AWS secret retrieval failed"), + ( + "ResourceNotFoundException", + "secret identifier", + "Credential secret not found", + ), + ("sts", "role trust conditions", "AWS credentials unavailable"), + ("timeout", "ReadTimeoutError", "AWS secret retrieval failed"), + ("network", "EndpointConnectionError", "AWS secret retrieval failed"), + ("provider", "credential provider", "AWS credentials unavailable"), + ("metadata", "MetadataRetrievalError", "AWS credentials unavailable"), + ("parameter", "parameter validation", "AWS secret retrieval failed"), + ("tainted", "AWS operation returned Unknown", "AWS secret retrieval failed"), + ], +) +def test_aws_preparation_keeps_safe_runtime_diagnostics( + aws_managed_service, + mock_bloodhound_api, + monkeypatch, + caplog, + level, + raw_key, + failure, + expected, + summary, +): + from botocore.exceptions import ( + CredentialRetrievalError, + EndpointConnectionError, + MetadataRetrievalError, + NoCredentialsError, + NoRegionError, + ParamValidationError, + ReadTimeoutError, + ) + + payload = f"{AWS_ACCESS_KEY} {AWS_SECRET_KEY} private-runtime-session-token" + errors = { + "region": NoRegionError(), + "credentials": NoCredentialsError(), + "refresh": RuntimeError( + "Credentials were refreshed, but the refreshed credentials are still expired." + ), + "timeout": ReadTimeoutError( + endpoint_url=f"https://example.test/?token={payload}" + ), + "network": EndpointConnectionError( + endpoint_url=f"https://example.test/?token={payload}" + ), + "provider": CredentialRetrievalError( + provider="custom-process", error_msg=payload + ), + "metadata": MetadataRetrievalError( + error_msg=f"Received non 200 response 403 from container metadata: {payload}" + ), + "parameter": ParamValidationError( + report=f"Invalid WebIdentityToken: {payload}" + ), + } + request_id = "aaaaaaaa-bbbb-cccc-dddd-eeeeeeeeeeee" + operation = "AssumeRoleWithWebIdentity" if failure == "sts" else "GetSecretValue" + code = "AccessDeniedException" if failure in ("kms", "sts") else failure + if failure == "tainted": + operation, code, request_id = AWS_ACCESS_KEY, AWS_SECRET_KEY, AWS_SECRET_KEY + error = errors.get(failure) or ClientError( + { + "Error": { + "Code": code, + "Message": payload + (" kms:Decrypt" if failure == "kms" else ""), + }, + "ResponseMetadata": { + "RequestId": request_id, + "HTTPStatusCode": AWS_SECRET_KEY if failure == "tainted" else 403, + "RetryAttempts": AWS_SECRET_KEY if failure == "tainted" else 2, + }, + }, + operation, + ) + + def fail(*args, **kwargs): + raise error + + # Cover both SDK construction (including workload STS) and retrieval failures. + monkeypatch.setattr("openhound.core.clients.aws_secrets_manager.boto3.client", fail) + aws_managed_service.secrets_manager = AWSSecretsManager( + None + if failure in ("region", "sts") + else SimpleNamespace( + get_secret_value=fail, + meta=SimpleNamespace( + region_name=AWS_SECRET_KEY if failure == "tainted" else "us-west-2" + ), + ) + ) + state = mock_bloodhound_api.app.state + if not raw_key: + state.collector_job_claim_response["external_secret_key_id"] = None + with caplog.at_level(level): + aws_managed_service._poll() + + record = next( + r for r in caplog.records if "Cannot prepare managed job" in r.getMessage() + ) + assert record.levelno == logging.ERROR + assert "OpenHound runtime secret retrieval failed" in record.getMessage() + assert expected in record.getMessage() + assert record.sdk_exception_type == type(error).__name__ + if failure == "tainted": + assert record.aws_operation == "AWS operation" + assert record.aws_error_code == "Unknown" + assert all( + key not in record.__dict__ + for key in ("region", "RequestId", "HTTPStatusCode", "RetryAttempts") + ) + elif isinstance(error, ClientError): + assert record.aws_operation == operation + assert record.aws_error_code == code + assert record.RequestId == request_id + assert record.HTTPStatusCode == 403 + assert record.RetryAttempts == 2 + if failure not in ("region", "sts", "tainted"): + assert record.region == "us-west-2" + if failure == "provider": + assert record.credential_provider == "custom-process" + if failure == "metadata": + assert record.credential_endpoint_http_status == 403 + assert state.collector_job_end_requests[0]["body"] == { + "outcome": "failed", + "failure_message": summary, + } + assert state.collector_job_start_requests == [] + assert aws_managed_service.executor.submitted == [] + for forbidden in (AWS_ACCESS_KEY, AWS_SECRET_KEY, "private-runtime-session-token"): + assert forbidden not in caplog.text + assert all(forbidden not in repr(r.__dict__) for r in caplog.records) + assert forbidden not in repr(state.collector_job_end_requests) + assert all(r.exc_info is None for r in caplog.records) + + +@pytest.mark.parametrize("launch_path", ["run_claimed_job", "_launch_job"]) +@pytest.mark.parametrize("wrong_collector", [False, True]) +def test_aws_launch_paths_validate_before_start( + mock_service, mock_bloodhound_api, launch_path, wrong_collector +): + mock_service.collector_name = "other" if wrong_collector else "aws" + params = ( + {} + if wrong_collector + else {"sources": {"source": {"aws": {"credentials": "not-an-object"}}}} + ) + kwargs = {"managed": True} if launch_path == "_launch_job" else {} + with pytest.raises(AWSAccessKeyError): + getattr(mock_service, launch_path)( + "job", + params=params, + secrets=AWSAccessKeyCredentials(AWS_ACCESS_KEY, AWS_SECRET_KEY), + **kwargs, + ) + assert mock_bloodhound_api.app.state.collector_job_start_requests == [] + assert mock_service.executor.submitted == [] + assert mock_service._aws_profile_workspace is None + assert mock_service.future is mock_service.job_running is None + + +@pytest.mark.parametrize( + "params", + [ + {"credentials": {"profile_names": ["other"]}}, + {"sources": {"credentials": {"org_role_name": "OtherRole"}}}, + {"sources": {"source": {"credentials": {"profile_name": "other"}}}}, + {"sources": {"aws": {"credentials": {"profile_names": ["other"]}}}}, + {"aws_collector": {"credentials": {"profile_names": ["other"]}}}, + { + "aws_collector": { + "sources": { + "source": {"aws": {"credentials": {"profile_name": "other"}}} + } + } + }, + ], +) +def test_raw_aws_preserves_alternative_nested_authentication_layouts( + aws_managed_service, mock_bloodhound_api, params +): + state = mock_bloodhound_api.app.state + state.collector_job_claim_response["params"] = params + aws_managed_service._poll() + assert len(state.collector_job_start_requests) == 1 + submitted = aws_managed_service.executor.submitted[0][0] + expected = deepcopy(params) + expected.setdefault("sources", {}).setdefault("source", {}).setdefault( + "aws", {} + ).setdefault("credentials", {}) + assert submitted[3] == expected + assert state.collector_job_end_requests == [] + aws_managed_service._cleanup_aws_profile_workspace() + + +@pytest.mark.parametrize( + "key", ["", "nonstandard-key", "ASIAIOSFODNN7EXAMPLE", "key\n"] +) +@pytest.mark.parametrize( + "secret", ["", "short", "%" * 41, '"quoted-key"', '["key"]', " key\n"] +) +def test_nonstandard_raw_keys_reach_handoff_unchanged( + aws_managed_service, mock_bloodhound_api, key, secret +): + state = mock_bloodhound_api.app.state + state.collector_job_claim_response["external_secret_key_id"] = key + aws_managed_service.secrets_manager = AWSSecretsManager( + SimpleNamespace(get_secret_value=lambda **kwargs: {"SecretString": secret}) + ) + aws_managed_service._poll() + submitted = aws_managed_service.executor.submitted[0][0] + assert submitted[4].access_key_id == key + assert submitted[4].secret_access_key == secret + assert len(state.collector_job_start_requests) == 1 + aws_managed_service._cleanup_aws_profile_workspace() + + +@pytest.mark.parametrize("reference", [" reference with spaces", "reference\nother"]) +def test_aws_secret_reference_is_passed_to_reader_unchanged( + aws_managed_service, mock_bloodhound_api, monkeypatch, reference +): + state = mock_bloodhound_api.app.state + state.collector_job_claim_response["secret_key_id"] = reference + seen = [] + + def retrieve(value): + seen.append(value) + return AWS_SECRET_KEY + + monkeypatch.setattr(aws_managed_service.secrets_manager, "get_secret", retrieve) + aws_managed_service._poll() + assert seen == [reference] + assert len(state.collector_job_start_requests) == 1 + aws_managed_service._cleanup_aws_profile_workspace() + + +@pytest.mark.parametrize( + "document", [{}, {"custom": {"nested": [None, {}, [], False, 0, ""]}}] +) +def test_nested_credential_values_reach_runner_unchanged(managed_service, document): + managed_service.secrets_manager = AWSSecretsManager( + SimpleNamespace( + get_secret_value=lambda **kwargs: {"SecretString": json.dumps(document)} + ) + ) + managed_service._poll() + assert managed_service.handoffs[0].credentials == document + + +@pytest.mark.parametrize( + "association_index", [0, 1], ids=["null-associations", "uuid-associations"] +) +def test_managed_poll_starts_claimed_job_and_tracks_worker( + managed_service, mock_bloodhound_api, monkeypatch, association_index +): + service = managed_service + state = mock_bloodhound_api.app.state + associated_job = load_json("collector_jobs_available_with_jobs.json")["data"][ + "jobs" + ][association_index] + for field in ("job_schedule_id", "job_profile_id"): + state.collector_job_queue_response["data"]["jobs"][0][field] = associated_job[ + field + ] + state.collector_job_claim_response[field] = associated_job[field] + service.managed_job_runner = service._run_prepared_managed_job + future = Future() + submitted = [] + + def submit(*args): + assert state.collector_job_start_requests == [ + (state.collector_job_claim_response["id"], b"") + ] + submitted.append(args) + return future + + monkeypatch.setattr(service.executor, "submit", submit) + + service._poll() + + job_id = state.collector_job_claim_response["id"] + assert state.managed_events == ["claim", "retrieve"] + assert submitted == [ + ( + scheduler_service._subprocess_collect, + "openhound-faker", + job_id, + CLAIMED_PARAMS, + COLLECTOR_SECRETS, + True, + ) + ] + assert service.job_running == service.managed_job_id == job_id + assert service.future is future + + service.next_heartbeat_at = 0 + service._poll() + assert len(state.collector_job_claim_requests) == 1 + assert len(state.collector_job_queue_requests) == 1 + assert state.collector_job_heartbeat_requests == [] + + state.collector_job_queue_response = load_json( + "collector_jobs_available_empty.json" + ) + future.set_result(Result(results={}, job_id=job_id)) + service._poll() + assert state.collector_job_end_requests == [ + {"job_id": job_id, "body": {"outcome": "succeeded"}} + ] + assert service.job_running is service.future is None + + +@pytest.mark.parametrize( + "committed", [False, True], ids=["before-start", "lost-response"] +) +@pytest.mark.parametrize( + "error", + [ + requests.Timeout("raw-provider-secret"), + requests.ConnectionError("raw-provider-secret"), + BloodHoundHTTPError("raw-provider-secret", 503), + ], + ids=["timeout", "connection-error", "http-503"], +) +def test_managed_start_failure_retries_prepared_job_before_discovery( + stateful_managed_service, mock_bloodhound_api, monkeypatch, caplog, committed, error +): + service = stateful_managed_service + state = mock_bloodhound_api.app.state + state.collector_job_lost_claim_responses = 0 + service.managed_job_runner = service._run_prepared_managed_job + job_id = state.collector_job_claim_response["id"] + original = service.client.request + starts = [] + + def flaky_request(**kwargs): + if kwargs["path"].endswith("/start"): + assert service._pending_managed_claim == job_id + assert service._pending_managed_start is not None + starts.append(kwargs["path"]) + if len(starts) <= 3: + if committed: + original(**kwargs) + raise error + return original(**kwargs) + + monkeypatch.setattr(service.client, "request", flaky_request) + + with caplog.at_level(logging.DEBUG): + for _ in range(3): + with pytest.raises(ManagedJobProcessingError) as raised: + service._poll() + assert "raw-provider-secret" not in "".join( + traceback.format_exception(raised.value) + ) + assert service._pending_managed_claim == job_id + assert service._pending_managed_start is not None + assert str(service._pending_managed_start.job.id) == job_id + assert "private-credential-value" not in repr( + service._pending_managed_start + ) + assert ( + service.future is service.job_running is service.managed_job_id is None + ) + assert service.executor.submitted == [] + assert state.collector_job_end_requests == [] + assert state.collector_job_claim_response["status"] == ( + "running" if committed else "claimed" + ) + + # Recovery must repeat start, not claim: BHE rejects re-claim of a running job. + service._poll() + + assert starts == [f"/api/v2/collie-job-queue/{job_id}/start"] * 4 + assert state.collector_job_claim_requests == [{"job_id": job_id, "body": b""}] + assert len(state.collector_job_queue_requests) == 1 + assert state.management_available_requests == 0 + assert service.secret_requests == ["private-secret-reference"] + assert service._pending_managed_claim is service._pending_managed_start is None + assert service.job_running == service.managed_job_id == job_id + assert state.collector_job_claim_response["status"] == "running" + (submission,) = service.executor.submitted + args, _, future = submission + assert args[1:] == ( + "openhound-faker", + job_id, + CLAIMED_PARAMS, + COLLECTOR_SECRETS, + True, + ) + assert service.future is future + for forbidden in ("raw-provider-secret", "private-credential-value"): + assert forbidden not in caplog.text + assert all(forbidden not in repr(record.__dict__) for record in caplog.records) + + service._poll() + assert len(starts) == 4 + assert len(service.executor.submitted) == 1 + + state.collector_job_queue_response = load_json( + "collector_jobs_available_empty.json" + ) + future.set_result(Result(results={}, job_id=job_id)) + service._poll() + assert state.collector_job_end_requests == [ + {"job_id": job_id, "body": {"outcome": "succeeded"}} + ] + assert service.job_running is service.future is None + + +@pytest.mark.parametrize("status", [404, 409]) +@pytest.mark.parametrize("lost_response_first", [False, True]) +def test_managed_start_rejection_clears_preparation_and_allows_next_job( + stateful_managed_service, + mock_bloodhound_api, + monkeypatch, + status, + lost_response_first, +): + service = stateful_managed_service + state = mock_bloodhound_api.app.state + state.collector_job_lost_claim_responses = 0 + service.managed_job_runner = service._run_prepared_managed_job + original = service.client.start_collector_job + + def lost_response(job_id): + original(job_id) + raise requests.Timeout("raw-provider-secret") + + if lost_response_first: + monkeypatch.setattr(service.client, "start_collector_job", lost_response) + with pytest.raises(ManagedJobProcessingError): + service._poll() + assert service._pending_managed_start is not None + + monkeypatch.setattr(service.client, "start_collector_job", original) + state.collector_job_start_statuses = [status] + with pytest.raises(ManagedJobProcessingError): + service._poll() + + assert service._pending_managed_claim is service._pending_managed_start is None + assert service.future is service.job_running is service.managed_job_id is None + assert service.executor.submitted == [] + assert len(state.collector_job_queue_requests) == 1 + assert len(state.collector_job_claim_requests) == 1 + assert state.collector_job_end_requests == [] + + # BHE no longer permits this start; discovery can now select another job. + job_b = state.collector_job_queue_response["data"]["jobs"][1] + state.collector_job_queue_response["data"]["jobs"] = [job_b] + state.collector_job_claim_response.update(id=job_b["id"], status="claimed") + state.collector_job_active_id = None + service._poll() + assert len(state.collector_job_queue_requests) == 2 + assert service.job_running == job_b["id"] + assert len(service.executor.submitted) == 1 + + +def test_managed_expired_claim_start_rejection_clears_preparation( + stateful_managed_service, mock_bloodhound_api, monkeypatch +): + service = stateful_managed_service + state = mock_bloodhound_api.app.state + state.collector_job_lost_claim_responses = 0 + service.managed_job_runner = service._run_prepared_managed_job + original = service.client.start_collector_job + + def unavailable(job_id): + raise requests.Timeout("raw-provider-secret") + + monkeypatch.setattr(service.client, "start_collector_job", unavailable) + with pytest.raises(ManagedJobProcessingError): + service._poll() + assert service._pending_managed_start is not None + + state.collector_job_claim_response["claim_expires_at"] = ( + datetime.now(UTC) - timedelta(seconds=1) + ).isoformat() + monkeypatch.setattr(service.client, "start_collector_job", original) + with pytest.raises(ManagedJobProcessingError): + service._poll() + assert service._pending_managed_claim is service._pending_managed_start is None + assert service.executor.submitted == [] + assert state.collector_job_end_requests == [] + + +@pytest.mark.parametrize( + "document", + [ + "plain-secret", + '{"malformed":', + '{"value":" "}', + '{"value":null}', + '{"value":[]}', + '{"value":NaN}', + '{" ":"private-credential-value"}', + ], +) +def test_invalid_managed_credentials_end_failed_without_start( + managed_service, mock_bloodhound_api, document, caplog +): + state = mock_bloodhound_api.app.state + + def retrieve(*, SecretId): + state.managed_events.append("retrieve") + return {"SecretString": document} + + managed_service.secrets_manager = AWSSecretsManager( + SimpleNamespace(get_secret_value=retrieve) + ) + with caplog.at_level(logging.DEBUG): + managed_service._poll() + + assert state.managed_events == ["claim", "retrieve", "end"] + assert state.collector_job_end_requests == [ + { + "job_id": state.collector_job_claim_response["id"], + "body": { + "outcome": "failed", + "failure_message": "Credential document is invalid", + }, + } + ] + assert managed_service.handoffs == [] + assert state.job_started is False + assert "Credential retrieval valid" not in caplog.text + assert "private-credential-value" not in caplog.text + assert all(r.exc_info is None for r in caplog.records) + + +@pytest.mark.parametrize("reference", [None, " "]) +def test_missing_managed_secret_reference_fails_without_local_fallback( + managed_service, mock_bloodhound_api, monkeypatch, reference +): + monkeypatch.setenv("SOURCES__SOURCE__FAKER__CREDENTIALS", "local-credentials") + state = mock_bloodhound_api.app.state + state.collector_job_claim_response["secret_key_id"] = reference + managed_service._poll() + + assert state.managed_events == ["claim", "end"] + assert managed_service.secret_requests == managed_service.handoffs == [] + assert ( + state.collector_job_end_requests[0]["body"]["failure_message"] + == "Credential document is invalid" + ) + + +@pytest.mark.parametrize( + "error,summary", + [ + ( + ClientError( + { + "Error": { + "Code": "AccessDeniedException", + "Message": "The workload role cannot read openhound/job", + } + }, + "GetSecretValue", + ), + "AWS secret access denied", + ), + ( + ClientError( + { + "Error": { + "Code": "ResourceNotFoundException", + "Message": "Secret does not exist in this region", + } + }, + "GetSecretValue", + ), + "Credential secret not found", + ), + ( + RuntimeError("private-credential-value"), + "AWS secret retrieval failed", + ), + ], + ids=["aws-permission-error", "aws-missing-secret", "unexpected-reader-error"], +) +def test_secret_retrieval_failure_keeps_local_cause_and_reports_safe_summary( + managed_service, mock_bloodhound_api, caplog, error, summary +): + state = mock_bloodhound_api.app.state + + def retrieve(*, SecretId): + state.managed_events.append("retrieve") + raise error + + managed_service.secrets_manager = AWSSecretsManager( + SimpleNamespace(get_secret_value=retrieve) + ) + with caplog.at_level(logging.DEBUG): + managed_service._poll() + + assert state.managed_events == ["claim", "retrieve", "end"] + assert state.collector_job_end_requests[0]["body"] == { + "outcome": "failed", + "failure_message": summary, + } + assert managed_service.handoffs == [] + assert "Reporting this job failed" in caplog.text + failure_log = next( + r for r in caplog.records if "Cannot prepare managed job" in r.getMessage() + ) + assert failure_log.secret_id == state.collector_job_claim_response["secret_key_id"] + if isinstance(error, ClientError): + assert error.response["Error"]["Message"] in caplog.text + assert "GetSecretValue" in caplog.text + assert failure_log.getMessage().count(failure_log.secret_id) == 1 + assert "private-credential-value" not in caplog.text + assert all( + "private-credential-value" not in repr(r.__dict__) for r in caplog.records + ) + + +@pytest.mark.parametrize("status", [409, 404, 400, 503]) +def test_failed_claim_does_not_retrieve_or_end( + managed_service, mock_bloodhound_api, monkeypatch, caplog, status +): + state = mock_bloodhound_api.app.state + attempts = ( + bloodhound_enterprise.MANAGED_JOB_CLAIM_MAX_ATTEMPTS if status == 503 else 1 + ) + state.collector_job_claim_statuses = [status] * attempts + delays = [] + monkeypatch.setattr(bloodhound_enterprise.time, "sleep", delays.append) + + with caplog.at_level(logging.DEBUG): + if status == 409: + managed_service._poll() + else: + with pytest.raises(ManagedJobProcessingError) as raised: + managed_service._poll() + assert "raw-provider-secret" not in "".join( + traceback.format_exception(raised.value) + ) + + assert state.managed_events == ["claim"] * attempts + assert delays == [bloodhound_enterprise.MANAGED_JOB_CLAIM_RETRY_DELAY_SECONDS] * ( + attempts - 1 + ) + assert managed_service.secret_requests == managed_service.handoffs == [] + assert state.collector_job_end_requests == [] + assert "Job successfully picked up" not in caplog.text + assert "Credential validation failed" not in caplog.text + assert "raw-provider-secret" not in caplog.text + + +@pytest.mark.parametrize( + "error", + [ + BloodHoundHTTPError("raw-provider-secret", 429), + BloodHoundHTTPError("raw-provider-secret", 502), + requests.ConnectionError("raw-provider-secret"), + requests.Timeout("raw-provider-secret"), + ], +) +def test_claim_retries_ambiguous_transient_failure_for_same_id( + managed_service, mock_bloodhound_api, monkeypatch, error +): + original = managed_service.client.request + claims = [] + delays = [] + + def flaky_request(**kwargs): + if kwargs["path"].endswith("/claim"): + claims.append(kwargs["path"]) + if len(claims) == 1: + raise error + return original(**kwargs) + + monkeypatch.setattr(managed_service.client, "request", flaky_request) + monkeypatch.setattr(bloodhound_enterprise.time, "sleep", delays.append) + managed_service._poll() + + assert len(claims) == 2 and claims[0] == claims[1] + assert delays == [bloodhound_enterprise.MANAGED_JOB_CLAIM_RETRY_DELAY_SECONDS] + assert mock_bloodhound_api.app.state.managed_events == [ + "claim", + "retrieve", + "start_handoff", + ] + + +def test_invalid_claim_response_redacts_parsing_error( + mock_service, mock_bloodhound_api, caplog +): + mock_bloodhound_api.app.state.collector_job_claim_response["id"] = ( + "private-credential-value" + ) + with caplog.at_level(logging.DEBUG), pytest.raises(RuntimeError) as raised: + mock_service.client.claim_managed_job("job-id") + assert "private-credential-value" not in "".join( + traceback.format_exception(raised.value) + ) + assert "private-credential-value" not in caplog.text + + +@pytest.mark.parametrize("resolved_status", [200, 404]) +def test_exhausted_credential_failure_end_is_retained_and_retried_before_discovery( + managed_service, mock_bloodhound_api, monkeypatch, caplog, resolved_status +): + state = mock_bloodhound_api.app.state + from botocore.exceptions import NoCredentialsError + + def unavailable(**kwargs): + raise NoCredentialsError() + + managed_service.secrets_manager = AWSSecretsManager( + SimpleNamespace(get_secret_value=unavailable) + ) + max_attempts = bloodhound_enterprise.MANAGED_JOB_END_MAX_ATTEMPTS + state.collector_job_end_statuses = [503] * max_attempts + delays = [] + monkeypatch.setattr(bloodhound_enterprise.time, "sleep", delays.append) + + with ( + caplog.at_level(logging.DEBUG), + pytest.raises(ManagedJobProcessingError) as raised, + ): + managed_service._poll() + assert ( + managed_service._pending_managed_failure + == state.collector_job_claim_response["id"] + ) + assert "raw-provider-secret" not in "".join( + traceback.format_exception(raised.value) + ) + assert "raw-provider-secret" not in caplog.text + assert all("raw-provider-secret" not in repr(r.__dict__) for r in caplog.records) + assert delays == [bloodhound_enterprise.MANAGED_JOB_END_RETRY_DELAY_SECONDS] * ( + max_attempts - 1 + ) + + state.collector_job_end_statuses = [resolved_status] + managed_service._poll() + assert managed_service._pending_managed_failure is None + assert all( + request["body"]["failure_message"] == "AWS credentials unavailable" + for request in state.collector_job_end_requests + ) + assert ( + len(state.collector_job_queue_requests) + == len(state.collector_job_claim_requests) + == 1 + ) + assert state.managed_events == ["claim"] + ["end"] * (max_attempts + 1) + assert managed_service.handoffs == [] + assert delays == [bloodhound_enterprise.MANAGED_JOB_END_RETRY_DELAY_SECONDS] * ( + max_attempts - 1 + ) + + +def test_managed_poll_skips_nonmatching_queue_entries( + managed_service, mock_bloodhound_api +): + state = mock_bloodhound_api.app.state + jobs = state.collector_job_queue_response["data"]["jobs"] + jobs.insert( + 0, + { + **jobs[0], + "job_key": "other-collector", + "id": "22222222-2222-2222-2222-222222222222", + }, + ) + managed_service._poll() + assert ( + state.collector_job_claim_requests[0]["job_id"] + == "11111111-1111-1111-1111-111111111111" + ) + + +def test_managed_runtime_dependency_fails_before_claim_or_legacy_work( + managed_service, mock_bloodhound_api +): + managed_service.managed_job_runner = None + with pytest.raises( + ManagedRuntimeUnavailableError, match="Managed collection runtime is required" + ): + managed_service._poll() + state = mock_bloodhound_api.app.state + assert ( + state.collector_job_claim_requests == state.collector_job_queue_requests == [] + ) + assert state.jobs_available_requests == state.jobs_current_requests == 0 + assert state.job_started is False + + +def test_mode_config_routes_scheduler_to_managed_dependency(mock_service, monkeypatch): + monkeypatch.setattr(scheduler_service, "is_managed", lambda: True) + service = Service( + bhe_uri="http://localhost:8000", + token_key="test-key", + token_id="test-id", + collector_name="openhound-faker", + ) + assert service.managed is True + assert service.managed_job_runner == service._run_prepared_managed_job + + +def test_runtime_failure_is_sanitized_and_not_reported_as_credential_failure( + managed_service, mock_bloodhound_api, caplog +): + def fail_start(prepared): + raise ValueError(str(prepared.credentials)) + + managed_service.managed_job_runner = fail_start + with ( + caplog.at_level(logging.DEBUG), + pytest.raises(ManagedJobProcessingError) as raised, + ): + managed_service._poll() + assert "private-credential-value" not in "".join( + traceback.format_exception(raised.value) + ) + assert "private-credential-value" not in caplog.text + assert "Credential validation failed" not in caplog.text + assert mock_bloodhound_api.app.state.collector_job_end_requests == [] + + +def test_valid_generic_nested_json_is_handed_off_without_field_schema(managed_service): + credentials = { + "sources": { + "source": {"aws": {"arbitrary_field": {"nested": ["value", False, 0, 1.5]}}} + } + } + managed_service.secrets_manager = AWSSecretsManager( + SimpleNamespace( + get_secret_value=lambda **kwargs: {"SecretString": json.dumps(credentials)} + ) + ) + managed_service._poll() + assert managed_service.handoffs[0].credentials == credentials + + +def test_managed_scheduler_logs_failure_once_and_retries_reporting_next_poll( + managed_service, mock_bloodhound_api, monkeypatch, caplog +): + state = mock_bloodhound_api.app.state + state.collector_job_claim_response["secret_key_id"] = None + original_end = managed_service.client.end_managed_job + reports = [] + pending_at_sleep = [] + + def end_job(*args, **kwargs): + reports.append((args, kwargs)) + if len(reports) == 1: + raise requests.Timeout("raw-provider-secret") + return original_end(*args, **kwargs) + + def next_poll(seconds): + assert seconds == managed_service.interval + pending_at_sleep.append(managed_service._pending_managed_failure) + if len(pending_at_sleep) == 2: + managed_service.exit = True + + monkeypatch.setattr(managed_service.client, "end_managed_job", end_job) + monkeypatch.setattr(scheduler_service.time, "sleep", next_poll) + monkeypatch.setattr(scheduler_service.signal, "signal", lambda *args: None) + + with caplog.at_level(logging.DEBUG): + managed_service.start() + + assert reports == [reports[0], reports[0]] + assert pending_at_sleep == [state.collector_job_claim_response["id"], None] + assert ( + len(state.collector_job_queue_requests) + == len(state.collector_job_claim_requests) + == 1 + ) + assert managed_service.handoffs == [] + assert [ + r.getMessage() for r in caplog.records if "poll failed" in r.getMessage() + ] == [ + f"Managed scheduler poll failed; retrying next poll in {managed_service.interval} seconds." + ] + assert not any( + r.getMessage() == "Managed collector job processing failed." + for r in caplog.records + ) + assert "raw-provider-secret" not in caplog.text + assert all("raw-provider-secret" not in repr(r.__dict__) for r in caplog.records) + + +@pytest.mark.parametrize( + "stage,failure,expected", + [ + ("poll", "region", "AWS_DEFAULT_REGION"), + ("poll", "provider", "custom-process"), + ("end", "document", "non-empty token_id"), + ], +) +def test_managed_401_refresh_failure_is_visible_without_debug_logging( + managed_service, mock_bloodhound_api, monkeypatch, caplog, stage, failure, expected +): + from botocore.exceptions import CredentialRetrievalError, NoRegionError + + secret_requests = [] + + def retrieve(*, SecretId): + secret_requests.append(SecretId) + if failure == "region": + raise NoRegionError() + if failure == "provider": + raise CredentialRetrievalError( + provider="custom-process", error_msg="private-refresh-token" + ) + return {"SecretString": '{"token_key":"private-refresh-token"}'} + + provider = AWSBHECredentials( + "openhound/bhe", AWSSecretsManager(SimpleNamespace(get_secret_value=retrieve)) + ) + managed_service.client._credential_refresh = provider.refresh + requests_made = [] + + def unauthorized(self, **kwargs): + requests_made.append(kwargs["path"]) + raise BloodHoundHTTPError("API token expired", 401) + + def next_poll(seconds): + assert seconds == managed_service.interval + managed_service.exit = True + + monkeypatch.setattr(bloodhound.BloodHound, "request", unauthorized) + monkeypatch.setattr(managed_service, "_poll_management", lambda: False) + monkeypatch.setattr(scheduler_service.signal, "signal", lambda *args: None) + monkeypatch.setattr(scheduler_service.time, "sleep", next_poll) + if stage == "end": + managed_service._pending_managed_failure = ( + "11111111-1111-1111-1111-111111111111" + ) + + with caplog.at_level(logging.INFO): + managed_service.start() + + assert len(requests_made) == 1 + assert secret_requests == ["openhound/bhe"] + failures = [ + r for r in caplog.records if "Cannot refresh BHE credentials" in r.getMessage() + ] + assert len(failures) == 1 + assert expected in failures[0].getMessage() + assert "openhound/bhe" in failures[0].getMessage() + assert all(r.exc_info is None for r in caplog.records) + assert "private-refresh-token" not in repr([r.__dict__ for r in caplog.records]) + assert managed_service.handoffs == [] + assert mock_bloodhound_api.app.state.collector_job_end_requests == [] + assert f"retrying next poll in {managed_service.interval} seconds" in caplog.text + if stage == "end": + assert managed_service._pending_managed_failure is not None + + +@pytest.fixture +def stateful_managed_service(managed_service, mock_bloodhound_api): + state = mock_bloodhound_api.app.state + state.collector_job_stateful_claims = True + job_a = state.collector_job_queue_response["data"]["jobs"][0] + job_b = {**job_a, "id": "22222222-2222-2222-2222-222222222222"} + state.collector_job_queue_response["data"]["jobs"].append(job_b) + state.collector_job_lost_claim_responses = ( + bloodhound_enterprise.MANAGED_JOB_CLAIM_MAX_ATTEMPTS + ) + return managed_service + + +def test_committed_claim_with_lost_responses_is_reconciled_before_discovery( + stateful_managed_service, mock_bloodhound_api, monkeypatch, caplog +): + service = stateful_managed_service + state = mock_bloodhound_api.app.state + job_id = state.collector_job_claim_response["id"] + delays = [] + monkeypatch.setattr(bloodhound_enterprise.time, "sleep", delays.append) + + with caplog.at_level(logging.DEBUG), pytest.raises(ManagedJobProcessingError): + service._poll() + + assert state.collector_job_active_id == job_id + assert service._pending_managed_claim == job_id + assert service.secret_requests == service.handoffs == [] + assert state.collector_job_end_requests == [] + assert delays == [bloodhound_enterprise.MANAGED_JOB_CLAIM_RETRY_DELAY_SECONDS] * 3 + + # A fresh discovery would offer B, which this client cannot claim while + # holding A. Recovery must use A's retained ID instead of this result. + available = mock_bloodhound_api.get("/api/v2/collie-job-queue/available").json() + assert [job["id"] for job in available["data"]["jobs"]] == [ + "22222222-2222-2222-2222-222222222222" + ] + discovery_count = len(state.collector_job_queue_requests) + with caplog.at_level(logging.DEBUG): + service._poll() + assert service._pending_managed_claim is None + assert service.secret_requests == ["private-secret-reference"] + (prepared,) = service.handoffs + assert str(prepared.job.id) == job_id + assert prepared.job.params == CLAIMED_PARAMS + + assert len(state.collector_job_queue_requests) == discovery_count + assert state.collector_job_claim_requests == [{"job_id": job_id, "body": b""}] * 5 + assert state.collector_job_end_requests == [] + assert state.job_started is False + assert "raw-provider-secret" not in caplog.text + assert all("raw-provider-secret" not in repr(r.__dict__) for r in caplog.records) + + +@pytest.mark.parametrize("status", [404, 409]) +def test_definitive_claim_rejection_clears_pending_id_without_credentials_or_end( + stateful_managed_service, mock_bloodhound_api, monkeypatch, status +): + service = stateful_managed_service + state = mock_bloodhound_api.app.state + monkeypatch.setattr(bloodhound_enterprise.time, "sleep", lambda seconds: None) + with pytest.raises(ManagedJobProcessingError): + service._poll() + + state.collector_job_claim_statuses = [status] + if status == 404: + with pytest.raises(ManagedJobProcessingError): + service._poll() + else: + service._poll() + + assert service._pending_managed_claim is None + assert service.secret_requests == service.handoffs == [] + assert state.collector_job_end_requests == [] + assert len(state.collector_job_queue_requests) == 1 + + # Once BHE confirms A is no longer claimable, the next poll can discover B. + job_b = state.collector_job_queue_response["data"]["jobs"][1] + state.collector_job_queue_response["data"]["jobs"] = [job_b] + state.collector_job_claim_response["id"] = job_b["id"] + state.collector_job_active_id = None + service._poll() + assert len(state.collector_job_queue_requests) == 2 + assert str(service.handoffs[0].job.id) == job_b["id"] + + +@pytest.mark.parametrize( + "field,value", + [ + ("id", "22222222-2222-2222-2222-222222222222"), + ("job_key", "other-collector"), + ("status", "running"), + ("claimed_by", None), + ("claimed_at", None), + ("claim_expires_at", None), + ("claim_expires_at", "2020-01-01T00:00:00Z"), + ("claim_expires_at", "2099-01-01T00:00:00"), + ], +) +def test_unusable_claim_response_blocks_credentials_and_retains_id( + managed_service, mock_bloodhound_api, field, value, caplog +): + state = mock_bloodhound_api.app.state + job_id = state.collector_job_claim_response["id"] + state.collector_job_claim_response[field] = value + with caplog.at_level(logging.DEBUG), pytest.raises(ManagedJobProcessingError): + managed_service._poll() + + assert managed_service._pending_managed_claim == job_id + assert managed_service.secret_requests == managed_service.handoffs == [] + assert state.collector_job_end_requests == [] + assert "Managed claim response does not confirm a usable claim." in caplog.text + assert "Credential validation failed" not in caplog.text + assert all( + "private-secret-reference" not in repr(r.__dict__) for r in caplog.records + ) + + +def test_lease_expiring_during_secret_retrieval_blocks_handoff( + managed_service, mock_bloodhound_api, monkeypatch +): + state = mock_bloodhound_api.app.state + expires_at = datetime.fromisoformat( + state.collector_job_claim_response["claim_expires_at"] + ) + clock = iter([expires_at - timedelta(seconds=1), expires_at + timedelta(seconds=1)]) + monkeypatch.setattr( + scheduler_service, "datetime", SimpleNamespace(now=lambda timezone: next(clock)) + ) + with pytest.raises(ManagedJobProcessingError): + managed_service._poll() + + assert ( + managed_service._pending_managed_claim + == state.collector_job_claim_response["id"] + ) + assert managed_service.secret_requests == ["private-secret-reference"] + assert managed_service.handoffs == [] + assert state.collector_job_end_requests == [] + + +def test_malformed_committed_claim_response_keeps_id_until_reconciled( + stateful_managed_service, mock_bloodhound_api, caplog +): + service = stateful_managed_service + state = mock_bloodhound_api.app.state + job_id = state.collector_job_claim_response["id"] + state.collector_job_lost_claim_responses = 0 + state.collector_job_claim_response["id"] = "raw-provider-secret" + with caplog.at_level(logging.DEBUG), pytest.raises(ManagedJobProcessingError): + service._poll() + + assert state.collector_job_active_id == job_id + assert service._pending_managed_claim == job_id + assert service.secret_requests == service.handoffs == [] + assert state.collector_job_end_requests == [] + assert "raw-provider-secret" not in caplog.text + + state.collector_job_claim_response["id"] = job_id + service._poll() + assert service._pending_managed_claim is None + assert len(state.collector_job_queue_requests) == 1 + assert len(service.handoffs) == 1 + assert str(service.handoffs[0].job.id) == job_id + + +def test_managed_scheduler_skips_legacy_calls_without_collection_runtime( + managed_service, mock_bloodhound_api, monkeypatch, caplog +): + service = managed_service + state = mock_bloodhound_api.app.state + service.managed_job_runner = None + state.management_operations = [_support_bundle_operation()] + polls = [] + + def next_poll(seconds): + assert seconds == service.interval + polls.append(seconds) + if len(polls) == 2: + service.exit = True + + monkeypatch.setattr(scheduler_service.time, "sleep", next_poll) + monkeypatch.setattr(scheduler_service.signal, "signal", lambda *args: None) + with caplog.at_level(logging.DEBUG): + service.start() + + assert len(polls) == 2 + assert state.management_available_requests == 0 + assert state.client_update_payload is None + assert state.operation_started is False + assert state.artifact_completed is False + assert state.operation_end_payload is None + assert ( + state.collector_job_queue_requests == state.collector_job_claim_requests == [] + ) + assert state.jobs_available_requests == state.jobs_current_requests == 0 + assert state.job_started is False + assert service.handoffs == [] + assert ( + "Managed collection runtime is unavailable; retrying next poll." in caplog.text + ) + + +@pytest.mark.parametrize("pending", ["claim", "failure"]) +def test_managed_recovery_continues_without_legacy_management( + managed_service, mock_bloodhound_api, monkeypatch, pending +): + service = managed_service + state = mock_bloodhound_api.app.state + state.management_operations = [_support_bundle_operation()] + job_id = state.collector_job_claim_response["id"] + if pending == "claim": + service._pending_managed_claim = job_id + else: + service._pending_managed_failure = job_id + sent = [] + monkeypatch.setattr(service, "_send_support_bundle", sent.append) + + service._poll() + + assert state.management_available_requests == 0 + assert sent == [] + assert state.collector_job_queue_requests == [] + assert service._pending_managed_claim is service._pending_managed_failure is None + assert state.managed_events == ( + ["claim", "retrieve", "start_handoff"] if pending == "claim" else ["end"] + ) + + service._poll() + assert state.management_available_requests == 0 + assert sent == [] + assert len(state.collector_job_queue_requests) == 1 + + +def test_support_bundle_tracebacks_preserve_aws_auth_context( + mock_service, monkeypatch, caplog +): + secret_reference = "private-bootstrap-secret-reference" + provider_message = "private-aws-provider-message" + + def get_secret_value(*, SecretId): + raise ClientError( + {"Error": {"Code": "AccessDeniedException", "Message": provider_message}}, + "GetSecretValue", + ) + + credentials = AWSBHECredentials( + secret_reference, + AWSSecretsManager(SimpleNamespace(get_secret_value=get_secret_value)), + ) + mock_service.client._credential_refresh = credentials.refresh + + def expired_token(self, **kwargs): + raise BloodHoundHTTPError("API token expired", 401) + + monkeypatch.setattr(bloodhound.BloodHound, "request", expired_token) + operation = ManagementOperation.model_validate(_support_bundle_operation()) + with caplog.at_level(logging.DEBUG), pytest.raises(SecretPermissionError) as raised: + mock_service._send_support_bundle(operation) + + assert "Secrets Manager GetSecretValue" in caplog.text + assert "AccessDeniedException" in caplog.text + assert any(record.exc_info is not None for record in caplog.records) + formatted_error = "".join(traceback.format_exception(raised.value)) + for context in (secret_reference, provider_message): + assert context in caplog.text + assert context in formatted_error + + +def test_managed_worker_resolves_nested_configuration_before_pipeline( + monkeypatch, caplog +): + collector = SimpleNamespace( + name="faker", + dlt_source=SimpleNamespace(section="source", name="random"), + package_version="1.0.0", + metadata=None, + ) + observed_configuration = {} + monkeypatch.setattr( + scheduler_service.CollectorManager, + "from_entrypoint", + lambda: SimpleNamespace(collectors=[collector]), + ) + + def capture_pipeline(extension): + observed_configuration.update( + node_count=dlt.config.get("sources.source.random.node_count", int), + filters=dlt.config.get("sources.source.random.filters", list), + token=dlt.secrets.get("sources.source.random.credentials.token", str), + ) + return {"collect": []} + + monkeypatch.setattr(scheduler_service.dataflow, "pipeline", capture_pipeline) + + with caplog.at_level(logging.INFO, logger="openhound.scheduler.service"): + result = _subprocess_collect( + "faker", + "collector-job-123", + { + "sources": { + "source": {"random": {"node_count": 4, "filters": ["one", "two"]}} + } + }, + { + "sources": { + "source": {"random": {"credentials": {"token": "secret-value"}}} + } + }, + ) + + assert result.job_id == "collector-job-123" + assert observed_configuration == { + "node_count": 4, + "filters": ["one", "two"], + "token": "secret-value", + } + assert "secret-value" not in caplog.text + + +@pytest.mark.parametrize( + "params,secrets,reason", + [ + ( + {"private-key": "private-credential-value"}, + {}, + "flat configuration is not supported", + ), + ( + CLAIMED_PARAMS, + {"private-key": "private-credential-value"}, + "flat configuration is not supported", + ), + ( + {}, + COLLECTOR_SECRETS, + "must be a non-empty nested JSON object", + ), + ( + { + "sources": { + "source": {"random": {"node_count": "private-credential-value"}} + } + }, + {}, + "could not be resolved; check setting types and required fields", + ), + ( + CLAIMED_PARAMS, + { + "sources": { + "source": {"random": {"node_count": "private-credential-value"}} + } + }, + "could not be resolved; check setting types and required fields", + ), + ], + ids=[ + "flat-params", + "flat-secrets", + "empty-params", + "params-resolution", + "secrets-resolution", + ], +) +def test_managed_configuration_failure_is_logged_and_reported_with_the_same_reason( + mock_service, + mock_bloodhound_api, + monkeypatch, + caplog, + params, + secrets, + reason, +): + monkeypatch.setattr( + scheduler_service.dataflow, + "pipeline", + lambda extension: { + "node_count": dlt.config.get("sources.source.random.node_count", int) + }, + ) + with ( + caplog.at_level(logging.DEBUG), + pytest.raises(ManagedConfigurationError) as error, + ): + _subprocess_collect( + "faker", + "collector-job-123", + params, + secrets, + ) + + future = Future() + future.set_exception(error.value) + mock_service.managed = True + mock_service.job_running = mock_service.managed_job_id = "collector-job-123" + mock_service.future = future + mock_service._handle_completed_job(future) + + (request,) = mock_bloodhound_api.app.state.collector_job_end_requests + assert request["body"]["outcome"] == "failed" + message = request["body"]["failure_message"] + assert reason in message + (record,) = [ + record + for record in caplog.records + if record.name == "openhound.scheduler.service" + and record.levelno == logging.ERROR + and record.getMessage().startswith( + "Managed job collector-job-123 configuration" + ) + ] + assert record.getMessage() == ( + f"Managed job collector-job-123 configuration is invalid: {message}" + ) + assert record.job_id == "collector-job-123" + assert record.collector_extension == mock_service.collector_name + assert record.exc_info is None + assert ( + mock_service.future + is mock_service.job_running + is mock_service.managed_job_id + is None + ) + assert "private-key" not in caplog.text + assert "private-credential-value" not in caplog.text + assert all( + "private-credential-value" not in repr(record.__dict__) + for record in caplog.records + ) + assert "private-credential-value" not in json.dumps(request) + + +def test_managed_worker_registers_real_faker_source_in_a_fresh_process(monkeypatch): + monkeypatch.setenv("SOURCES__SOURCE__RANDOM__NODE_COUNT", "999") + code = """ +from openhound.scheduler.service import _subprocess_collect, dataflow + +def inspect_collection(extension): + return {"node_count": len(list(extension.dlt_source()))} + +dataflow.pipeline = inspect_collection +result = _subprocess_collect( + "faker", "e0d43dbd-2cdf-4d68-9a81-2f95cd58a92f", + {"sources": {"source": {"random": {"node_count": 4}}}}, {} +) +assert result.results == {"node_count": 4}, result.results +result = _subprocess_collect( + "faker", "second-job", {"sources": {"source": {"random": {"node_count": 2}}}}, {} +) +assert result.results == {"node_count": 2}, result.results +import dlt +assert dlt.config.get("sources.source.random.node_count", int) == 999 +""" + subprocess.run( + [sys.executable, "-c", code], check=True, capture_output=True, timeout=30 + ) + + +class FakeClock: + def __init__(self, on_sleep=lambda now: None): + self.now = 0 + self.on_sleep = on_sleep + + def monotonic(self): + return self.now + + def sleep(self, seconds): + assert seconds > 0 + self.now += seconds + self.on_sleep(self.now) + + +@pytest.fixture +def enabled_managed_heartbeats(monkeypatch): + """Keep the deferred heartbeat implementation covered until BHE supports it.""" + monkeypatch.setattr(scheduler_service, "_MANAGED_HEARTBEATS_DISABLED", False) + + +@pytest.mark.usefixtures("enabled_managed_heartbeats") +def test_managed_job_uses_shared_poll_loop_without_polling_for_new_jobs( + mock_service, mock_bloodhound_api, monkeypatch, caplog +): + future = Future() + submitted = [] + + def submit(*args): + assert mock_bloodhound_api.app.state.collector_job_start_requests == [ + ("collector-job-123", b"") + ] + submitted.append(args) + return future + + clock = FakeClock() + monkeypatch.setattr(mock_service.executor, "submit", submit) + monkeypatch.setattr(scheduler_service, "time", clock) + monkeypatch.setattr( + mock_service, + "check_jobs", + lambda: pytest.fail("job polling must stop during managed collection"), + ) + + def fail_unmanaged_checkin(self): + pytest.fail("managed jobs must use the heartbeat endpoint") + + monkeypatch.setattr( + mock_service.client.__class__, + "jobs_current", + property(fail_unmanaged_checkin), + ) + + with caplog.at_level(logging.INFO, logger="openhound.scheduler.service"): + mock_service.run_claimed_job( + "collector-job-123", {"node_count": 4}, {"token": "secret-value"} + ) + assert mock_service.job_running == "collector-job-123" + for now, expected_heartbeats in ( + (0, 0), + (299, 0), + (300, 1), + (599, 1), + (600, 2), + ): + clock.now = now + mock_service._poll() + assert ( + len(mock_bloodhound_api.app.state.collector_job_heartbeat_requests) + == expected_heartbeats + ) + future.set_result(Result(results={"collect": []}, job_id="collector-job-123")) + clock.now = 630 + monkeypatch.setattr(mock_service, "check_jobs", lambda: None) + mock_service._poll() + + assert submitted[0][1:] == ( + "openhound-faker", + "collector-job-123", + {"node_count": 4}, + {"token": "secret-value"}, + True, + ) + assert mock_bloodhound_api.app.state.collector_job_start_requests == [ + ("collector-job-123", b"") + ] + assert mock_bloodhound_api.app.state.collector_job_heartbeat_requests == [ + ("collector-job-123", b""), + ("collector-job-123", b""), + ] + assert mock_bloodhound_api.app.state.collector_job_end_requests == [ + {"job_id": "collector-job-123", "body": {"outcome": "succeeded"}} + ] + assert mock_service.future is None + assert mock_service.job_running is None + assert "claimed successfully" not in caplog.text + assert "started successfully" in caplog.text + assert "heartbeat succeeded" in caplog.text + assert "secret-value" not in caplog.text + + +def test_managed_start_failure_does_not_launch_worker(mock_service, monkeypatch): + monkeypatch.setattr( + mock_service.client, + "start_collector_job", + lambda job_id: (_ for _ in ()).throw(BloodHoundHTTPError("conflict", 409)), + ) + monkeypatch.setattr( + mock_service.executor, + "submit", + lambda *args: pytest.fail("worker must not launch after start failure"), + ) + + with pytest.raises(BloodHoundHTTPError) as error: + mock_service.run_claimed_job("collector-job-123", {}, {}) + + assert error.value.code == 409 + assert mock_service.future is None + assert mock_service.job_running is None + assert mock_service.managed_job_id is None + assert mock_service.next_heartbeat_at is None + + +@pytest.mark.parametrize( + ("job_id", "managed", "message"), + [ + (123, True, "Managed job IDs must be strings"), + ("collector-job-123", False, "Unmanaged job IDs must be integers"), + ], +) +def test_launch_job_rejects_invalid_id_type_before_start( + mock_service, mock_bloodhound_api, job_id, managed, message +): + with pytest.raises(TypeError, match=message): + mock_service._launch_job(job_id, managed=managed) + + assert mock_bloodhound_api.app.state.collector_job_start_requests == [] + assert mock_bloodhound_api.app.state.job_started is False + assert mock_service.executor.submitted == [] + assert mock_service.future is mock_service.job_running is None + + +def test_managed_runtime_normalizes_uuid_at_entry(mock_service, mock_bloodhound_api): + job_id = UUID("e0d43dbd-2cdf-4d68-9a81-2f95cd58a92f") + + mock_service.run_claimed_job(job_id, {}, {}) + + assert mock_bloodhound_api.app.state.collector_job_start_requests == [ + (str(job_id), b"") + ] + assert mock_service.job_running == mock_service.managed_job_id == str(job_id) + args, _, _ = mock_service.executor.submitted[0] + assert args[2] == str(job_id) + + +def test_claimed_job_cannot_start_while_another_job_runs( + mock_service, mock_bloodhound_api +): + mock_service.job_running = 123 + mock_service.future = Future() + + with pytest.raises(RuntimeError, match="Job 123 is already running"): + mock_service.run_claimed_job("collector-job-123", {}, {}) + + assert mock_bloodhound_api.app.state.collector_job_start_requests == [] + assert mock_service.job_running == 123 + + +@pytest.mark.usefixtures("enabled_managed_heartbeats") +def test_transient_managed_heartbeat_failures_keep_collection_running( + mock_service, mock_bloodhound_api, monkeypatch, caplog +): + future = Future() + attempts = [] + + clock = FakeClock() + monkeypatch.setattr(mock_service.executor, "submit", lambda *args: future) + monkeypatch.setattr(scheduler_service, "time", clock) + monkeypatch.setattr(mock_service, "check_jobs", lambda: None) + + def heartbeat(job_id): + attempts.append(clock.now) + if len(attempts) == 1: + raise requests.Timeout("no response") + if len(attempts) == 2: + raise BloodHoundHTTPError("temporary failure", 503) + + monkeypatch.setattr(mock_service.client, "heartbeat_collector_job", heartbeat) + + with caplog.at_level(logging.INFO, logger="openhound.scheduler.service"): + mock_service.run_claimed_job("collector-job-123", {}, {}) + for now in (300, 600, 900): + clock.now = now + mock_service._poll() + future.set_result(Result(results={}, job_id="collector-job-123")) + clock.now = 930 + mock_service._poll() + + assert attempts == [300, 600, 900] + assert mock_bloodhound_api.app.state.collector_job_end_requests == [ + {"job_id": "collector-job-123", "body": {"outcome": "succeeded"}} + ] + assert caplog.text.count("heartbeat request failed") == 2 + assert "heartbeat succeeded" in caplog.text + + +@pytest.mark.parametrize("during_shutdown", [False, True], ids=["poll", "shutdown"]) +@pytest.mark.usefixtures("enabled_managed_heartbeats") +def test_managed_heartbeat_claim_loss_stops_worker_without_reporting_outcome( + mock_service, mock_bloodhound_api, monkeypatch, during_shutdown +): + future = Future() + future.set_running_or_notify_cancel() + events = [] + executor = mock_service.executor + + class Worker: + def kill(self): + events.append("kill") + future.set_exception(BrokenProcessPool("worker killed")) + + executor._processes = {123: Worker()} + monkeypatch.setattr(executor, "submit", lambda *args: future) + + def shutdown(*, wait, cancel_futures): + assert wait is True + assert cancel_futures is True + events.append("join") + + monkeypatch.setattr(executor, "shutdown", shutdown) + clock = FakeClock() + monkeypatch.setattr(scheduler_service, "time", clock) + + def lose_claim(job_id): + raise BloodHoundHTTPError("claim lost", 404) + + monkeypatch.setattr(mock_service.client, "heartbeat_collector_job", lose_claim) + mock_service.run_claimed_job("collector-job-123", {}, {}) + clock.now = 300 + if during_shutdown: + mock_service._shutdown() + else: + mock_service._poll() + monkeypatch.setattr(mock_service, "check_jobs", lambda: None) + mock_service._poll() + + assert events == ["kill", "join"] + assert future.cancelled() is False # A running future cannot simply be cancelled. + assert mock_service.executor is not executor + assert mock_service.future is None + assert mock_service.job_running is None + assert mock_service.managed_job_id is None + assert mock_service.next_heartbeat_at is None + assert mock_bloodhound_api.app.state.collector_job_end_requests == [] + assert mock_bloodhound_api.app.state.job_ended is False + + +@pytest.mark.usefixtures("enabled_managed_heartbeats") +def test_managed_heartbeat_claim_loss_kills_a_real_running_process( + mock_service, mock_bloodhound_api, monkeypatch +): + executor = ProcessPoolExecutor(max_workers=1, max_tasks_per_child=1) + original_submit = executor.submit + mock_service.executor = executor + monkeypatch.setattr( + executor, "submit", lambda *args: original_submit(time.sleep, 60) + ) + + def lose_claim(job_id): + raise BloodHoundHTTPError("claim lost", 404) + + monkeypatch.setattr(mock_service.client, "heartbeat_collector_job", lose_claim) + try: + mock_service.run_claimed_job("collector-job-123", {}, {}) + future = mock_service.future + deadline = time.monotonic() + 10 + while not future.running() and time.monotonic() < deadline: + time.sleep(0.01) + assert future.running() + workers = list(executor._processes.values()) + assert workers and all(worker.is_alive() for worker in workers) + + mock_service.next_heartbeat_at = 0 + mock_service._poll() + + assert all(not worker.is_alive() for worker in workers) + assert future.done() + assert mock_service.future is None + assert mock_service.executor is not executor + assert mock_bloodhound_api.app.state.collector_job_end_requests == [] + finally: + for worker in list((executor._processes or {}).values()): + worker.kill() + executor.shutdown(wait=True, cancel_futures=True) + mock_service.executor.shutdown(wait=True, cancel_futures=True) + + +def test_managed_worker_failure_ends_job_failed( + mock_service, mock_bloodhound_api, monkeypatch +): + future = Future() + future.set_exception(RuntimeError("worker error")) + monkeypatch.setattr(mock_service.executor, "submit", lambda *args: future) + monkeypatch.setattr(mock_service, "check_jobs", lambda: None) + + mock_service.run_claimed_job("collector-job-123", {}, {}) + mock_service._poll() + + assert mock_bloodhound_api.app.state.collector_job_heartbeat_requests == [] + assert mock_bloodhound_api.app.state.collector_job_end_requests == [ + { + "job_id": "collector-job-123", + "body": { + "outcome": "failed", + "failure_message": "Unexpected error while running 'openhound-faker' collector", + }, + } + ] + + +@pytest.mark.usefixtures("enabled_managed_heartbeats") +def test_shutdown_keeps_heartbeating_until_managed_worker_finishes( + mock_service, mock_bloodhound_api, monkeypatch +): + future = Future() + + def finish_worker(now): + if now >= 630 and not future.done(): + future.set_result(Result(results={}, job_id="collector-job-123")) + + clock = FakeClock(finish_worker) + monkeypatch.setattr(mock_service.executor, "submit", lambda *args: future) + monkeypatch.setattr(scheduler_service, "time", clock) + + mock_service.run_claimed_job("collector-job-123", {}, {}) + mock_service._shutdown() + + assert mock_bloodhound_api.app.state.collector_job_heartbeat_requests == [ + ("collector-job-123", b""), + ("collector-job-123", b""), + ] + assert mock_bloodhound_api.app.state.collector_job_end_requests == [ + {"job_id": "collector-job-123", "body": {"outcome": "succeeded"}} + ] + assert mock_service.future is None + + +@pytest.mark.parametrize( + "failure_phase", ["upload", "success_report", "failure_report"] +) +@pytest.mark.usefixtures("enabled_managed_heartbeats") +def test_managed_shutdown_retries_completion_and_renews_claim( + mock_service, mock_bloodhound_api, monkeypatch, tmp_path, caplog, failure_phase +): + job_id = "6bc0857e-9d9b-48c6-bbd8-77a31c5d351d" + output = tmp_path / "output" + output.mkdir() + files = [output / f"part-{index}.json" for index in range(2)] + for index, path in enumerate(files): + path.write_bytes(f"file-{index}".encode()) + result = Result(results={}, job_id=job_id, output_path=output, data_files=files) + future = Future() + future.set_result(result) + mock_service.managed = True + mock_service.job_running = mock_service.managed_job_id = job_id + mock_service.future = future + mock_service.interval = 300 + mock_service.next_heartbeat_at = 300 + clock = FakeClock() + monkeypatch.setattr(scheduler_service, "time", clock) + uploads = [] + ends = [] + original_end = mock_service.client.end_managed_job + + def upload(job, data): + assert job == job_id + uploads.append(data.read()) + if len(uploads) == 2: + if failure_phase == "upload": + raise requests.Timeout("sensitive-response-body") + if failure_phase == "failure_report": + raise BloodHoundHTTPError("sensitive-response-body", 415) + + def end(job, outcome, **kwargs): + ends.append(outcome) + if failure_phase != "upload" and len(ends) == 1: + raise BloodHoundHTTPError("sensitive-response-body", 500) + return original_end(job, outcome, **kwargs) + + def check_retry_state(now): + assert now == 300 # Only one retry delay should be needed. + assert mock_service.future is future + assert mock_service.job_running == job_id + assert result.uploaded_files == (2 if failure_phase == "success_report" else 1) + assert output.exists() + assert mock_bloodhound_api.app.state.collector_job_end_requests == [] + + clock.on_sleep = check_retry_state + monkeypatch.setattr(mock_service.client, "upload_collector_job_data", upload) + monkeypatch.setattr(mock_service.client, "end_managed_job", end) + + mock_service._shutdown() + + expected_outcome = ( + ManagedJobOutcome.FAILED + if failure_phase == "failure_report" + else ManagedJobOutcome.SUCCEEDED + ) + assert uploads == [b"file-0", b"file-1"] + ( + [b"file-1"] if failure_phase == "upload" else [] + ) + assert ends == [expected_outcome] * (1 if failure_phase == "upload" else 2) + body = {"outcome": expected_outcome.value} + if failure_phase == "failure_report": + body["failure_message"] = "Collector data upload rejected (HTTP 415)" + assert mock_bloodhound_api.app.state.collector_job_end_requests == [ + {"job_id": job_id, "body": body} + ] + assert mock_bloodhound_api.app.state.collector_job_heartbeat_requests == [ + (job_id, b"") + ] + assert mock_bloodhound_api.app.state.job_ended is False + assert mock_service.future is None + assert mock_service.job_running is None + assert mock_service.managed_job_id is None + assert not output.exists() + assert "sensitive-response-body" not in caplog.text + + +def test_managed_completion_fails_job_when_remaining_staged_file_disappears( + mock_service, monkeypatch, tmp_path +): + job_id = "6bc0857e-9d9b-48c6-bbd8-77a31c5d351d" + output = tmp_path / "output" + output.mkdir() + files = [output / f"part-{index}.json" for index in range(2)] + for index, path in enumerate(files): + path.write_bytes(f"file-{index}".encode()) + result = Result(results={}, job_id=job_id, output_path=output, data_files=files) + future = Future() + future.set_result(result) + mock_service.managed = True + mock_service.job_running = mock_service.managed_job_id = job_id + mock_service.future = future + uploads = [] + ends = [] + monkeypatch.setattr(scheduler_service, "time", FakeClock()) + monkeypatch.setattr(mock_service, "_poll_management", lambda: False) + monkeypatch.setattr(mock_service, "_poll_managed_jobs", lambda: None) + + def upload(job, data): + assert job == job_id + uploads.append(data.read()) + # Lose the remaining file after the first upload has been accepted. + files[1].unlink() + + def end(job, outcome, **kwargs): + ends.append((job, outcome, kwargs)) + + monkeypatch.setattr(mock_service.client, "upload_collector_job_data", upload) + monkeypatch.setattr(mock_service.client, "end_managed_job", end) + + mock_service._poll() + mock_service._poll() + + assert uploads == [b"file-0"] + assert len(ends) == 1, ( + "Missing staged data must fail the job instead of retrying forever" + ) + assert ends[0][:2] == (job_id, ManagedJobOutcome.FAILED) + assert ends[0][2]["failure_message"] + assert mock_service.future is None + assert mock_service.job_running is None + assert mock_service.managed_job_id is None + assert not output.exists() + + +@pytest.mark.usefixtures("enabled_managed_heartbeats") +def test_managed_poll_cleans_partial_upload_output_after_heartbeat_claim_loss( + mock_service, monkeypatch, tmp_path +): + job_id = "6bc0857e-9d9b-48c6-bbd8-77a31c5d351d" + output = tmp_path / "output" + output.mkdir() + files = [output / f"part-{index}.json" for index in range(2)] + for index, path in enumerate(files): + path.write_bytes(f"file-{index}".encode()) + result = Result(results={}, job_id=job_id, output_path=output, data_files=files) + future = Future() + future.set_result(result) + mock_service.managed = True + mock_service.job_running = mock_service.managed_job_id = job_id + mock_service.future = future + mock_service.next_heartbeat_at = 300 + clock = FakeClock() + clock.now = 300 + monkeypatch.setattr(scheduler_service, "time", clock) + monkeypatch.setattr(mock_service, "_poll_management", lambda: False) + monkeypatch.setattr(mock_service, "_poll_managed_jobs", lambda: None) + events = [] + ends = [] + + def upload(job, data): + assert job == job_id + content = data.read() + events.append(content) + if content == b"file-1": + raise requests.Timeout("upload interrupted") + + def lose_claim(job): + assert job == job_id + events.append("heartbeat") + raise BloodHoundHTTPError("claim lost", 404) + + monkeypatch.setattr(mock_service.client, "upload_collector_job_data", upload) + monkeypatch.setattr(mock_service.client, "heartbeat_collector_job", lose_claim) + monkeypatch.setattr( + mock_service.client, + "end_managed_job", + lambda *args, **kwargs: ends.append((args, kwargs)), + ) + + mock_service._poll() + mock_service._poll() + + assert events == [b"file-0", b"file-1", "heartbeat"] + assert result.uploaded_files == 1 + assert ends == [] + assert mock_service.future is None + assert mock_service.job_running is None + assert mock_service.managed_job_id is None + assert mock_service.next_heartbeat_at is None + assert not output.exists(), ( + "Claim loss must also remove partially uploaded staged data" + ) + + +@pytest.mark.usefixtures("enabled_managed_heartbeats") +def test_managed_shutdown_retries_upload_across_multiple_heartbeat_intervals( + mock_service, mock_bloodhound_api, monkeypatch, tmp_path +): + job_id = "6bc0857e-9d9b-48c6-bbd8-77a31c5d351d" + output = tmp_path / "output" + output.mkdir() + files = [output / f"part-{index}.json" for index in range(2)] + for index, path in enumerate(files): + path.write_bytes(f"file-{index}".encode()) + result = Result(results={}, job_id=job_id, output_path=output, data_files=files) + future = Future() + future.set_result(result) + mock_service.managed = True + mock_service.job_running = mock_service.managed_job_id = job_id + mock_service.future = future + mock_service.interval = 30 + mock_service.next_heartbeat_at = 300 + clock = FakeClock() + monkeypatch.setattr(scheduler_service, "time", clock) + uploads = [] + heartbeats = [] + retry_times = [] + original_heartbeat = mock_service.client.heartbeat_collector_job + + def upload(job, data): + assert job == job_id + content = data.read() + uploads.append((clock.now, content)) + if content == b"file-1" and clock.now < 930: + raise BloodHoundHTTPError("temporary outage", 503) + + def heartbeat(job): + heartbeats.append(clock.now) + return original_heartbeat(job) + + def check_retry_state(now): + assert now <= 930, "Shutdown should finish once the outage ends" + retry_times.append(now) + assert mock_service.future is future + assert mock_service.job_running == job_id + assert result.uploaded_files == 1 + assert output.exists() + assert mock_bloodhound_api.app.state.collector_job_end_requests == [] + + clock.on_sleep = check_retry_state + monkeypatch.setattr(mock_service.client, "upload_collector_job_data", upload) + monkeypatch.setattr(mock_service.client, "heartbeat_collector_job", heartbeat) + + mock_service._shutdown() + + assert retry_times == list(range(30, 931, 30)) + assert uploads == [(0, b"file-0")] + [(now, b"file-1") for now in range(0, 931, 30)] + assert heartbeats == [300, 600, 900] + assert mock_bloodhound_api.app.state.collector_job_end_requests == [ + {"job_id": job_id, "body": {"outcome": "succeeded"}} + ] + assert result.uploaded_files == 2 + assert mock_service.future is None + assert mock_service.job_running is None + assert mock_service.managed_job_id is None + assert not output.exists() + + +@pytest.mark.usefixtures("enabled_managed_heartbeats") +def test_managed_shutdown_cleans_staged_output_when_heartbeat_loses_claim( + mock_service, mock_bloodhound_api, monkeypatch, tmp_path +): + job_id = "6bc0857e-9d9b-48c6-bbd8-77a31c5d351d" + data_file = tmp_path / "output" / "graph.json" + data_file.parent.mkdir() + data_file.write_bytes(b"graph") + future = Future() + future.set_result( + Result( + results={}, + job_id=job_id, + output_path=data_file.parent, + data_files=[data_file], + ) + ) + mock_service.managed = True + mock_service.job_running = mock_service.managed_job_id = job_id + mock_service.future = future + mock_service.next_heartbeat_at = 0 + monkeypatch.setattr(scheduler_service, "time", FakeClock()) + + def lose_claim(job): + assert job == job_id + raise BloodHoundHTTPError("claim lost", 404) + + def unexpected_upload(*args, **kwargs): + raise AssertionError("A lost claim must not upload data") + + monkeypatch.setattr(mock_service.client, "heartbeat_collector_job", lose_claim) + monkeypatch.setattr( + mock_service.client, "upload_collector_job_data", unexpected_upload + ) + mock_service._shutdown() + + assert mock_service.future is None + assert mock_service.job_running is None + assert mock_service.managed_job_id is None + assert not data_file.parent.exists() + assert mock_bloodhound_api.app.state.collector_job_end_requests == [] + + +@pytest.mark.parametrize("end_fails", [False, True]) +def test_unmanaged_shutdown_waits_for_worker_and_attempts_completion_once( + mock_service, mock_bloodhound_api, monkeypatch, end_fails +): + future = Future() + mock_service.job_running = 123 + mock_service.future = future + ends = [] + checkins = [] + original_end = mock_service.client.end_job + original_checkin = mock_service._poll_running_job + + def finish_worker(now): + assert now == mock_service.interval + future.set_result(Result(results={}, job_id=123)) + + clock = FakeClock(finish_worker) + monkeypatch.setattr(scheduler_service, "time", clock) + + def end(status, message): + assert future.done() + ends.append(status) + if end_fails: + raise BloodHoundHTTPError("temporary failure", 500) + return original_end(status, message) + + def checkin(): + checkins.append(clock.now) + original_checkin() + + def unexpected_managed_completion(): + raise AssertionError("Unmanaged jobs must retain their existing shutdown path") + + monkeypatch.setattr(mock_service.client, "end_job", end) + monkeypatch.setattr(mock_service, "_poll_running_job", checkin) + monkeypatch.setattr( + mock_service, "_finish_managed_job_on_shutdown", unexpected_managed_completion + ) + mock_service._shutdown() + + assert ends == [JobStatus.COMPLETE] + assert checkins == [0] + assert mock_bloodhound_api.app.state.job_ended is (not end_fails) + assert mock_bloodhound_api.app.state.collector_job_heartbeat_requests == [] + assert mock_bloodhound_api.app.state.collector_job_end_requests == [] + assert mock_service.future is None + assert mock_service.job_running is None + + +@pytest.mark.parametrize( + "upload_error", + [ + BloodHoundHTTPError("temporary error", 500), + BloodHoundHTTPError("rate limited", 429), + requests.ConnectionError("connection interrupted"), + requests.Timeout("upload timed out"), + ], +) +def test_managed_completion_retries_remaining_uploads_before_ending( + mock_service, monkeypatch, tmp_path, upload_error +): + job_id = "6bc0857e-9d9b-48c6-bbd8-77a31c5d351d" + output = tmp_path / "output" + output.mkdir() + files = [output / "nodes.json", output / "edges.json"] + for index, path in enumerate(files): + path.write_bytes(f"file-{index}".encode()) + result = Result(results={}, job_id=job_id, output_path=output, data_files=files) + future = Future() + future.set_result(result) + mock_service.managed = True + mock_service.job_running = job_id + mock_service.managed_job_id = job_id + mock_service.future = future + calls = [] + + def upload(job, data): + assert job == job_id + calls.append(data.read()) + if len(calls) == 2: + raise upload_error + + def end(job, outcome, **kwargs): + calls.append((job, outcome)) + + monkeypatch.setattr(mock_service.client, "upload_collector_job_data", upload) + monkeypatch.setattr(mock_service.client, "end_managed_job", end) + + with pytest.raises(type(upload_error)): + mock_service._handle_completed_job(future) + assert result.uploaded_files == 1 + assert mock_service.future is future + assert mock_service.job_running == job_id + assert output.exists() + + mock_service._handle_completed_job(future) + assert calls == [ + b"file-0", + b"file-1", + b"file-1", + (str(job_id), ManagedJobOutcome.SUCCEEDED), + ] + assert mock_service.future is None + assert mock_service.job_running is None + assert not output.exists() + + +@pytest.mark.parametrize("status", [400, 401, 403, 415]) +@pytest.mark.parametrize("retry_end", [False, True]) +def test_permanent_upload_rejection_ends_failed_without_retrying_upload( + mock_service, monkeypatch, tmp_path, caplog, status, retry_end +): + job_id = "6bc0857e-9d9b-48c6-bbd8-77a31c5d351d" + output = tmp_path / "output" + output.mkdir() + files = [output / f"part-{index}.json" for index in range(3)] + for index, path in enumerate(files): + path.write_bytes(f"file-{index}".encode()) + result = Result(results={}, job_id=job_id, output_path=output, data_files=files) + future = Future() + future.set_result(result) + mock_service.managed = True + mock_service.job_running = mock_service.managed_job_id = job_id + mock_service.future = future + uploads = [] + ends = [] + + def upload(job, data): + uploads.append((job, data.read())) + if len(uploads) == 2: + raise BloodHoundHTTPError("sensitive-response-body", status) + + def end(job, outcome, **kwargs): + ends.append((job, outcome, kwargs)) + if retry_end and len(ends) == 1: + raise BloodHoundHTTPError("temporary error", 500) + + monkeypatch.setattr(mock_service.client, "upload_collector_job_data", upload) + monkeypatch.setattr(mock_service.client, "end_managed_job", end) + + if retry_end: + with pytest.raises(BloodHoundHTTPError): + mock_service._handle_completed_job(future) + assert mock_service.future is future + assert mock_service.job_running == job_id + assert output.exists() + + mock_service._handle_completed_job(future) + + assert uploads == [(job_id, b"file-0"), (job_id, b"file-1")] + assert ends == [ + ( + job_id, + ManagedJobOutcome.FAILED, + {"failure_message": f"Collector data upload rejected (HTTP {status})"}, + ) + ] * (2 if retry_end else 1) + assert "sensitive-response-body" not in caplog.text + assert mock_service.future is None + assert mock_service.job_running is None + assert mock_service.managed_job_id is None + assert not output.exists() + + +@pytest.mark.parametrize("status", [404, 409]) +def test_upload_for_missing_or_stopped_job_aborts_without_reporting_outcome( + mock_service, monkeypatch, tmp_path, caplog, status +): + job_id = "6bc0857e-9d9b-48c6-bbd8-77a31c5d351d" + output = tmp_path / "output" + output.mkdir() + data_file = output / "graph.json" + data_file.write_bytes(b"graph") + future = Future() + future.set_result( + Result(results={}, job_id=job_id, output_path=output, data_files=[data_file]) + ) + mock_service.managed = True + mock_service.job_running = mock_service.managed_job_id = job_id + mock_service.future = future + + def upload(job, data): + raise BloodHoundHTTPError("sensitive-response-body", status) + + def unexpected_end(*args, **kwargs): + raise AssertionError("A missing or stopped job must not receive an outcome") + + monkeypatch.setattr(mock_service.client, "upload_collector_job_data", upload) + monkeypatch.setattr(mock_service.client, "end_managed_job", unexpected_end) + mock_service._handle_completed_job(future) + + assert mock_service.future is None + assert mock_service.job_running is None + assert mock_service.managed_job_id is None + assert mock_service.next_heartbeat_at is None + assert not output.exists() + assert "sensitive-response-body" not in caplog.text + + +def test_managed_end_retry_does_not_upload_accepted_files_again( + mock_service, monkeypatch, tmp_path +): + job_id = "6bc0857e-9d9b-48c6-bbd8-77a31c5d351d" + output = tmp_path / "output" + output.mkdir() + data_file = output / "nodes.json" + data_file.write_bytes(b"graph") + result = Result( + results={}, job_id=job_id, output_path=output, data_files=[data_file] + ) + future = Future() + future.set_result(result) + mock_service.managed = True + mock_service.job_running = job_id + mock_service.managed_job_id = job_id + mock_service.future = future + uploads = [] + ends = [] + monkeypatch.setattr( + mock_service.client, + "upload_collector_job_data", + lambda job, data: uploads.append(data.read()), + ) + + def end(job, outcome, **kwargs): + ends.append((job, outcome)) + if len(ends) == 1: + raise BloodHoundHTTPError("temporary error", 500) + + monkeypatch.setattr(mock_service.client, "end_managed_job", end) + with pytest.raises(BloodHoundHTTPError): + mock_service._handle_completed_job(future) + assert output.exists() + assert mock_service.future is future + mock_service._handle_completed_job(future) + assert uploads == [b"graph"] + assert ends == [(str(job_id), ManagedJobOutcome.SUCCEEDED)] * 2 + assert not output.exists() + + +def test_managed_collection_failure_ends_failed_without_upload( + mock_service, monkeypatch +): + job_id = "6bc0857e-9d9b-48c6-bbd8-77a31c5d351d" + future = Future() + future.set_exception(ExtensionNotFoundError("missing collector")) + mock_service.managed = True + mock_service.job_running = job_id + mock_service.managed_job_id = job_id + mock_service.future = future + ends = [] + monkeypatch.setattr( + mock_service.client, + "end_managed_job", + lambda *args, **kwargs: ends.append((args, kwargs)), + ) + + def unexpected_upload(*args, **kwargs): + raise AssertionError("Failed collections must not upload") + + monkeypatch.setattr( + mock_service.client, "upload_collector_job_data", unexpected_upload + ) + mock_service._handle_completed_job(future) + assert ends[0][0] == (str(job_id), ManagedJobOutcome.FAILED) + assert ends[0][1]["failure_message"] == "Collector 'openhound-faker' not found" + assert mock_service.future is None + assert mock_service.job_running is None + + +def test_managed_collection_stages_then_uploads_before_job_end( + mock_service, mock_bloodhound_api, monkeypatch, tmp_path +): + monkeypatch.chdir(tmp_path) + monkeypatch.setenv("DLT_DATA_DIR", str(tmp_path / ".dlt")) + monkeypatch.setattr( + scheduler_service.dataflow, + "Path", + lambda value: tmp_path if value == "/tmp/openhound" else Path(value), + ) + job_id = "6bc0857e-9d9b-48c6-bbd8-77a31c5d351d" + result = _subprocess_collect("faker", job_id, managed=True) + assert result.data_files + assert mock_bloodhound_api.app.state.managed_ingest_requests == [] + assert mock_bloodhound_api.app.state.ingested_nodes == 0 + assert mock_bloodhound_api.app.state.ingested_edges == 0 + + future = Future() + future.set_result(result) + mock_service.managed = True + mock_service.job_running = job_id + mock_service.managed_job_id = job_id + mock_service.future = future + mock_service._handle_completed_job(future) + + state = mock_bloodhound_api.app.state + assert state.managed_ingest_requests == [str(job_id)] * len(result.data_files) + assert state.managed_ingested_nodes == 1000 + assert state.managed_ingested_edges == 10000 + assert state.collector_job_end_requests == [ + {"job_id": str(job_id), "body": {"outcome": "succeeded"}} + ] + assert state.job_ended is False + assert not result.output_path.exists() + + +def test_failed_managed_conversion_cleans_up_staged_files(monkeypatch, tmp_path): + output = tmp_path / "output" + output.mkdir() + monkeypatch.setattr( + scheduler_service.tempfile, "mkdtemp", lambda **kwargs: str(output) + ) + collector = SimpleNamespace(name="example", package_version="1", metadata=None) + monkeypatch.setattr( + scheduler_service.CollectorManager, + "from_entrypoint", + lambda: SimpleNamespace(collectors=[collector]), + ) + + def fail(extension, output_path): + (output_path / "partial.json").write_text("partial") + raise RuntimeError("conversion failed") + + monkeypatch.setattr(scheduler_service.dataflow, "pipeline", fail) + with pytest.raises(ManagedConfigurationError, match="check worker logs"): + _subprocess_collect( + "example", "6bc0857e-9d9b-48c6-bbd8-77a31c5d351d", managed=True + ) + assert not output.exists() + + +def test_managed_completion_streams_graph_to_api_before_ending( + mock_service, mock_bloodhound_api, tmp_path +): + job_id = "6bc0857e-9d9b-48c6-bbd8-77a31c5d351d" + output = tmp_path / "output" + output.mkdir() + data_file = output / "graph.json" + payload = json.loads( + (Path(__file__).parent / "test_data" / "sample_graph.json").read_text() + ) + for node in payload["graph"]["nodes"]: + node["properties"]["environmentid"] = "test-environment" + content = json.dumps(payload).encode() + data_file.write_bytes(content) + graph = Graph.model_validate_json(content) + future = Future() + future.set_result( + Result(results={}, job_id=job_id, output_path=output, data_files=[data_file]) + ) + mock_service.managed = True + mock_service.job_running = job_id + mock_service.managed_job_id = job_id + mock_service.future = future + + mock_service._handle_completed_job(future) + + state = mock_bloodhound_api.app.state + assert state.managed_ingest_requests == [str(job_id)] + assert state.managed_ingested_nodes == len(graph.graph.nodes) + assert state.managed_ingested_edges == len(graph.graph.edges) + assert state.collector_job_end_requests == [ + {"job_id": str(job_id), "body": {"outcome": "succeeded"}} + ] + assert state.ingested_nodes == 0 + assert state.ingested_edges == 0 + assert state.job_ended is False + assert not output.exists() diff --git a/tests/test_collector_job_ingest.py b/tests/test_collector_job_ingest.py new file mode 100644 index 00000000..527e1777 --- /dev/null +++ b/tests/test_collector_job_ingest.py @@ -0,0 +1,114 @@ +from uuid import UUID + +import pytest +import requests + +import openhound +from openhound.core.clients.bloodhound import BloodHoundHTTPError +from openhound.core.clients.bloodhound_enterprise import BloodHoundEnterprise + +JOB_ID = UUID("6bc0857e-9d9b-48c6-bbd8-77a31c5d351d") + + +@pytest.fixture(autouse=True) +def supported_package_version(monkeypatch): + monkeypatch.setattr(openhound, "__version__", "0.3.0rc1") + + +@pytest.mark.parametrize("encoding", [None, "gzip"]) +def test_upload_uses_job_endpoint_and_accepts_empty_202(monkeypatch, encoding): + calls = [] + + def request(**kwargs): + calls.append(kwargs) + response = requests.Response() + response.status_code = 202 + response._content = b"" + return response + + monkeypatch.setattr(requests, "request", request) + client = BloodHoundEnterprise("key", "token-id", "https://bhe.example") + data = b"collected file" + + client.upload_collector_job_data(JOB_ID, data, "application/json", encoding) + + assert len(calls) == 1 + call = calls[0] + assert call["method"] == "POST" + assert call["url"] == ( + f"https://bhe.example/api/v2/collie-job-queue/{JOB_ID}/ingest" + ) + assert call["data"] == data + assert call["headers"]["Content-Type"] == "application/json" + assert call["headers"].get("Content-Encoding") == encoding + assert call["headers"]["Authorization"] == "bhesignature token-id" + assert call["timeout"] == (10, 120) + + +@pytest.mark.parametrize("status", [400, 401, 403, 404, 409, 415, 500]) +def test_upload_propagates_endpoint_errors(monkeypatch, status): + def request(**kwargs): + response = requests.Response() + response.status_code = status + response._content = b"upload rejected" + return response + + monkeypatch.setattr(requests, "request", request) + client = BloodHoundEnterprise("key", "token-id") + + with pytest.raises(BloodHoundHTTPError) as error: + client.upload_collector_job_data(JOB_ID, b"{}", "application/json") + + assert error.value.code == status + + +def test_upload_rejects_legacy_numeric_job_id_before_request(monkeypatch): + def unexpected_request(**kwargs): + raise AssertionError("Legacy jobs must not use the managed ingest endpoint") + + monkeypatch.setattr(requests, "request", unexpected_request) + client = BloodHoundEnterprise("key", "token-id") + + with pytest.raises(ValueError): + client.upload_collector_job_data("123", b"{}", "application/json") + + +def test_stream_upload_signs_entire_body_and_rewinds_after_401(monkeypatch): + import base64 + import hashlib + import hmac + from io import BytesIO + + from openhound.core.clients.bhe_credentials import BHECredentials + + content = b'{"graph":{"nodes":[],"edges":[]}}' * 50000 + stream = BytesIO(content) + calls = [] + + def request(**kwargs): + # The HMAC pass must leave the file ready for the HTTP transport. + assert kwargs["data"].tell() == 0 + uploaded = kwargs["data"].read() + assert uploaded == content + headers = kwargs["headers"] + token_key = "old-key" if not calls else "new-key" + digest = hmac.new(token_key.encode(), None, hashlib.sha256) + digest.update(f"POST/api/v2/collie-job-queue/{JOB_ID}/ingest".encode()) + digest = hmac.new(digest.digest(), None, hashlib.sha256) + digest.update(headers["RequestDate"][:13].encode()) + digest = hmac.new(digest.digest(), uploaded, hashlib.sha256) + assert headers["Signature"] == base64.b64encode(digest.digest()).decode() + calls.append(headers["Authorization"]) + response = requests.Response() + response.status_code = 401 if len(calls) == 1 else 202 + response._content = b"" + return response + + monkeypatch.setattr(requests, "request", request) + client = BloodHoundEnterprise( + "old-key", + "old-id", + credential_refresh=lambda: BHECredentials("new-key", "new-id"), + ) + client.upload_collector_job_data(JOB_ID, stream) + assert calls == ["bhesignature old-id", "bhesignature new-id"] diff --git a/tests/test_config.py b/tests/test_config.py new file mode 100644 index 00000000..ac133896 --- /dev/null +++ b/tests/test_config.py @@ -0,0 +1,97 @@ +import dlt +import pytest +from dlt.common.configuration.exceptions import ConfigValueCannotBeCoercedException + +from openhound.config import ( + DEFAULT_MANAGED_BHE_POLL_INTERVAL, + ManagedBHEConfig, + ModeDefaults, + get_managed_bhe_config, + get_mode_defaults, + is_managed, +) + + +def test_is_managed_defaults_to_unmanaged(monkeypatch): + monkeypatch.setattr(dlt.config, "get", lambda field, expected_type: None) + + assert is_managed() is False + + +@pytest.mark.parametrize( + ("value", "expected"), + [ + ("true", True), + ("false", False), + ], +) +def test_is_managed_resolves_environment_configuration(monkeypatch, value, expected): + monkeypatch.setenv("OPENHOUND__MANAGED", value) + + assert is_managed() is expected + + +def test_is_managed_rejects_invalid_boolean(monkeypatch): + monkeypatch.setenv("OPENHOUND__MANAGED", "maybe") + + with pytest.raises(ConfigValueCannotBeCoercedException): + is_managed() + + +def test_get_mode_defaults_selects_mode_specific_log_formats(monkeypatch): + monkeypatch.setattr("openhound.config.is_managed", lambda: False) + unmanaged = get_mode_defaults() + + monkeypatch.setattr("openhound.config.is_managed", lambda: True) + managed = get_mode_defaults() + + assert isinstance(unmanaged, ModeDefaults) + assert isinstance(managed, ModeDefaults) + assert unmanaged is not managed + assert unmanaged.log_format == "text" + assert managed.log_format == "json" + + +def test_get_managed_bhe_config_uses_configured_secret_and_default_interval( + monkeypatch, +): + def get(field, expected_type): + values = { + "openhound.aws_secrets_manager.secret_name": "production/bhe", + "openhound.aws_secrets_manager.poll_interval": None, + } + return values[field] + + monkeypatch.setattr(dlt.config, "get", get) + + assert get_managed_bhe_config() == ManagedBHEConfig( + secret_name="production/bhe", + poll_interval=DEFAULT_MANAGED_BHE_POLL_INTERVAL, + ) + + +def test_get_managed_bhe_config_rejects_missing_secret_name(monkeypatch): + monkeypatch.setattr(dlt.config, "get", lambda field, expected_type: None) + + with pytest.raises(ValueError, match="secret_name"): + get_managed_bhe_config() + + +def test_get_managed_bhe_config_rejects_non_positive_interval(monkeypatch): + def get(field, expected_type): + return "production/bhe" if field.endswith("secret_name") else 0 + + monkeypatch.setattr(dlt.config, "get", get) + + with pytest.raises(ValueError, match="between 1 and 599 seconds"): + get_managed_bhe_config() + + +def test_get_managed_bhe_config_rejects_interval_at_bhe_timeout(monkeypatch): + def get(field, expected_type): + return "production/bhe" if field.endswith("secret_name") else 600 + + monkeypatch.setattr(dlt.config, "get", get) + + with pytest.raises(ValueError, match="between 1 and 599 seconds"): + get_managed_bhe_config() diff --git a/tests/test_data/api/jobs/collector_jobs_available_empty.json b/tests/test_data/api/jobs/collector_jobs_available_empty.json new file mode 100644 index 00000000..4b71256f --- /dev/null +++ b/tests/test_data/api/jobs/collector_jobs_available_empty.json @@ -0,0 +1,8 @@ +{ + "count": 0, + "skip": 0, + "limit": 1, + "data": { + "jobs": [] + } +} diff --git a/tests/test_data/api/jobs/collector_jobs_available_with_job.json b/tests/test_data/api/jobs/collector_jobs_available_with_job.json new file mode 100644 index 00000000..0c3610c1 --- /dev/null +++ b/tests/test_data/api/jobs/collector_jobs_available_with_job.json @@ -0,0 +1,43 @@ +{ + "count": 1, + "skip": 0, + "limit": 1, + "data": { + "jobs": [ + { + "id": "11111111-1111-1111-1111-111111111111", + "job_schedule_id": null, + "job_profile_id": null, + "job_type_id": 7, + "job_key": "openhound-faker", + "version": "v2", + "params": { + "sources": { + "source": { + "random": { + "domains": ["example.com", "example.org"], + "include_deleted": false, + "nested": {"depth": 2} + } + } + } + }, + "scope_client_id": null, + "external_secret_key_id": null, + "secret_key_id": null, + "priority": 10, + "status": "ready", + "run_at": "2026-02-20T10:00:00Z", + "unclaimed_deadline_at": "2026-02-20T10:05:00Z", + "attempts": 0, + "max_attempts": 3, + "last_failure": null, + "claimed_by": null, + "claimed_at": null, + "claim_expires_at": null, + "created_at": "2026-02-20T09:00:00Z", + "updated_at": "2026-02-20T09:30:00Z" + } + ] + } +} diff --git a/tests/test_data/api/jobs/collector_jobs_available_with_jobs.json b/tests/test_data/api/jobs/collector_jobs_available_with_jobs.json new file mode 100644 index 00000000..e1ad94f0 --- /dev/null +++ b/tests/test_data/api/jobs/collector_jobs_available_with_jobs.json @@ -0,0 +1,67 @@ +{ + "count": 2, + "skip": 0, + "limit": 2, + "data": { + "jobs": [ + { + "id": "11111111-1111-1111-1111-111111111111", + "job_schedule_id": null, + "job_profile_id": null, + "job_type_id": 7, + "job_key": "openhound-faker", + "version": "v2", + "params": { + "sources": { + "source": { + "random": { + "domains": ["example.com", "example.org"], + "include_deleted": false, + "nested": {"depth": 2} + } + } + } + }, + "scope_client_id": null, + "external_secret_key_id": null, + "secret_key_id": null, + "priority": 10, + "status": "ready", + "run_at": "2026-02-20T10:00:00Z", + "unclaimed_deadline_at": "2026-02-20T10:05:00Z", + "attempts": 0, + "max_attempts": 3, + "last_failure": null, + "claimed_by": null, + "claimed_at": null, + "claim_expires_at": null, + "created_at": "2026-02-20T09:00:00Z", + "updated_at": "2026-02-20T09:30:00Z" + }, + { + "id": "22222222-2222-2222-2222-222222222222", + "job_schedule_id": "44444444-4444-4444-4444-444444444444", + "job_profile_id": "55555555-5555-5555-5555-555555555555", + "job_type_id": 8, + "job_key": "openhound-faker", + "version": "v3", + "params": {"sources": {"source": {"random": {"mode": "full"}}}}, + "scope_client_id": "33333333-3333-3333-3333-333333333333", + "external_secret_key_id": "external-key", + "secret_key_id": "secret-1", + "priority": 5, + "status": "ready", + "run_at": "2026-02-20T11:00:00Z", + "unclaimed_deadline_at": "2026-02-20T11:05:00Z", + "attempts": 1, + "max_attempts": 4, + "last_failure": "previous attempt failed", + "claimed_by": null, + "claimed_at": null, + "claim_expires_at": null, + "created_at": "2026-02-20T09:45:00Z", + "updated_at": "2026-02-20T10:50:00Z" + } + ] + } +} diff --git a/tests/test_data/managed_aws_configuration.json b/tests/test_data/managed_aws_configuration.json new file mode 100644 index 00000000..411b72be --- /dev/null +++ b/tests/test_data/managed_aws_configuration.json @@ -0,0 +1,21 @@ +{ + "sources": { + "source": { + "aws": { + "credentials": { + "profile_names": ["bhe-profile"], + "region_name": "us-west-2", + "endpoint_url": "https://aws.example.test" + }, + "bucket_policy": { + "include": true, + "max_objects": 50, + "rule": [ + {"include": true, "match_bucket": ".*"}, + {"include": false, "match_object": "private/.*"} + ] + } + } + } + } +} diff --git a/tests/test_data/managed_aws_secrets.json b/tests/test_data/managed_aws_secrets.json new file mode 100644 index 00000000..20a8ec05 --- /dev/null +++ b/tests/test_data/managed_aws_secrets.json @@ -0,0 +1,11 @@ +{ + "sources": { + "source": { + "aws": { + "credentials": { + "profile_names": ["retrieved-profile-a", "retrieved-profile-b"] + } + } + } + } +} diff --git a/tests/test_log_handlers.py b/tests/test_log_handlers.py index d83373e6..a6792ad3 100644 --- a/tests/test_log_handlers.py +++ b/tests/test_log_handlers.py @@ -1,19 +1,79 @@ import json import logging +import os +import threading +from concurrent.futures.process import BrokenProcessPool +from contextlib import nullcontext +from decimal import Decimal +from io import StringIO from pathlib import Path +import dlt import pytest +from rich.console import Console +from rich.logging import RichHandler from openhound.core.logging import ( CustomLogger, + LogMode, OpenHoundJSONFormatter, + OpenHoundRichFormatter, OpenHoundTextFormatter, RotatingFileHandler, - logger_override, + _ParentWorkerHandler, + _ValueRedactor, + redact_log_values, ) -def test_root_handler_setup(): +@pytest.fixture +def configured_logger(tmp_path, monkeypatch): + """Set up real root/DLT loggers without retaining process-wide changes.""" + original_get = dlt.config.get + + def get(field, expected_type): + if field.startswith("runtime."): + return { + "runtime.log_path": str(tmp_path), + "runtime.log_format": "JSON", + }.get(field) + return original_get(field, expected_type) + + monkeypatch.setattr(dlt.config, "get", get) + monkeypatch.setattr( + CustomLogger, "runtime_mode", property(lambda self: LogMode.SERVICE) + ) + custom_logger = CustomLogger("openhound.log", base_path=str(tmp_path)) + original = [ + (logger, logger.handlers[:], logger.level, logger.propagate) + for logger in (custom_logger.root_logger, custom_logger.dlt_logger) + ] + try: + custom_logger.setup() + yield custom_logger + finally: + for logger, handlers, level, propagate in original: + logger.handlers[:] = handlers + logger.setLevel(level) + logger.propagate = propagate + for handler in custom_logger._file_handlers.values(): + handler.close() + + +@pytest.fixture +def no_log_format_override(monkeypatch): + """Remove the resolved override, including values from local TOML files.""" + original_get = dlt.config.get + + def get(field, expected_type): + if field == "runtime.log_format": + return None + return original_get(field, expected_type) + + monkeypatch.setattr(dlt.config, "get", get) + + +def test_root_handler_setup(configured_logger): """Test that the root logger has a handler configured and that it is a RotatingFileHandler""" root_logger = logging.getLogger() assert len(root_logger.handlers) > 0, ( @@ -29,7 +89,7 @@ def test_root_handler_setup(): ) -def test_dlt_handler_setup(): +def test_dlt_handler_setup(configured_logger): """Test that the default DLT handler is overwritten by our RotatingFileHandler""" dlt_logger = logging.getLogger("dlt") assert len(dlt_logger.handlers) > 0, ( @@ -48,9 +108,9 @@ def test_dlt_handler_setup(): ) -def test_dlt_extension_handlers(): +def test_dlt_extension_handlers(configured_logger): """Test that the DLT handler is correctly set up for an extension and that the destination filename has changed""" - logger_override.set_handler("test_extension") + configured_logger.set_handler("test_extension") dlt_logger = logging.getLogger("dlt") assert len(dlt_logger.handlers) > 0, ( "The DLT logger should have handlers configured after setting an extension handler" @@ -68,21 +128,16 @@ def test_dlt_extension_handlers(): ) -def test_log_routing_content(tmp_path, caplog, monkeypatch): +def test_log_routing_content(configured_logger, tmp_path, caplog): """Test that logs are correctly routed and that the files are created for the expected paths""" - # Pin JSON (dlt's only crash-safe override value) so the JSON-parsing assertions are self-contained. - monkeypatch.setenv("RUNTIME__LOG_FORMAT", "JSON") - logger_override.base_path = tmp_path - logger_override.setup() - logger_override.set_handler("test_extension") + configured_logger.set_handler("test_extension") root_logger = logging.getLogger() dlt_logger = logging.getLogger("dlt") - with caplog.at_level(logging.INFO): - with caplog.at_level(logging.INFO, logger="dlt"): - root_logger.info("Core openhound log") - dlt_logger.info("Extension DLT log") + with caplog.at_level(logging.INFO), caplog.at_level(logging.INFO, logger="dlt"): + root_logger.info("Core openhound log") + dlt_logger.info("Extension DLT log") assert (tmp_path / "openhound.log").exists(), ( "The 'openhound.log' file should exist for the core logs" @@ -127,6 +182,97 @@ def test_file_formatter_selection(): ) +def test_managed_mode_writes_structured_file_logs( + tmp_path, monkeypatch, no_log_format_override +): + """Managed mode should default rotating file logs to newline-delimited JSON.""" + monkeypatch.setenv("OPENHOUND__MANAGED", "true") + monkeypatch.setenv("RUNTIME__LOG_PATH", str(tmp_path)) + custom_logger = CustomLogger("managed.log", base_path=str(tmp_path)) + # Keep this setup isolated from the process-wide loggers used by other tests. + custom_logger.root_logger = logging.getLogger("managed-test") + custom_logger.dlt_logger = logging.getLogger("managed-test-dlt") + custom_logger.root_logger.propagate = False + custom_logger.dlt_logger.propagate = False + + try: + custom_logger.setup() + custom_logger.root_logger.info( + "Collection started", extra={"job_id": 42, "collector": "aws"} + ) + + log_record = json.loads((tmp_path / "managed.log").read_text()) + assert log_record["message"] == "Collection started" + assert log_record["job_id"] == 42 + assert log_record["collector"] == "aws" + assert isinstance( + custom_logger.root_logger.handlers[0].formatter, + OpenHoundJSONFormatter, + ) + finally: + for handler in custom_logger._file_handlers.values(): + handler.close() + custom_logger.root_logger.handlers.clear() + custom_logger.dlt_logger.handlers.clear() + + +def test_explicit_log_format_overrides_managed_default(tmp_path, monkeypatch): + """Operators can retain the existing runtime override in managed mode.""" + monkeypatch.setenv("OPENHOUND__MANAGED", "true") + monkeypatch.setenv("RUNTIME__LOG_FORMAT", "text") + monkeypatch.setenv("RUNTIME__LOG_PATH", str(tmp_path)) + custom_logger = CustomLogger("managed.log", base_path=str(tmp_path)) + custom_logger.root_logger = logging.getLogger("managed-override-test") + custom_logger.dlt_logger = logging.getLogger("managed-override-test-dlt") + custom_logger.root_logger.propagate = False + custom_logger.dlt_logger.propagate = False + + try: + custom_logger.setup() + + assert isinstance( + custom_logger.root_logger.handlers[0].formatter, + OpenHoundTextFormatter, + ) + finally: + for handler in custom_logger._file_handlers.values(): + handler.close() + custom_logger.root_logger.handlers.clear() + custom_logger.dlt_logger.handlers.clear() + + +def test_mode_change_refreshes_cached_file_handler( + tmp_path, monkeypatch, no_log_format_override +): + """Re-running setup should apply a changed mode to an existing log file.""" + monkeypatch.setenv("RUNTIME__LOG_PATH", str(tmp_path)) + custom_logger = CustomLogger("mode-change.log", base_path=str(tmp_path)) + custom_logger.root_logger = logging.getLogger("mode-change-test") + custom_logger.dlt_logger = logging.getLogger("mode-change-test-dlt") + custom_logger.root_logger.propagate = False + custom_logger.dlt_logger.propagate = False + + try: + monkeypatch.setenv("OPENHOUND__MANAGED", "false") + custom_logger.setup() + file_handler = custom_logger._get_file_handler(tmp_path / "mode-change.log") + assert isinstance(file_handler.formatter, OpenHoundTextFormatter) + + monkeypatch.setenv("OPENHOUND__MANAGED", "true") + custom_logger.setup() + + assert ( + custom_logger._get_file_handler(tmp_path / "mode-change.log") + is file_handler + ) + assert isinstance(file_handler.formatter, OpenHoundJSONFormatter) + finally: + for handler in custom_logger._file_handlers.values(): + handler.close() + custom_logger.root_logger.handlers.clear() + custom_logger.dlt_logger.handlers.clear() + + def test_text_formatter_produces_plain_text(): """Test that the text formatter produces a readable, non-JSON line with extras""" formatter = OpenHoundTextFormatter() @@ -159,6 +305,111 @@ def test_text_formatter_produces_plain_text(): json.loads(output) +@pytest.mark.parametrize( + "message,args,secret,expected", + [ + ("Key %d", (12345,), "12345", "Key [REDACTED]"), + ( + "Key %s costs %.2f", + ("example-secret-key", Decimal("12.34")), + "example-secret-key", + "Key [REDACTED] costs 12.34", + ), + ( + "Key %(key)s costs %(cost).2f", + ({"key": "example-secret-key", "cost": Decimal("12.34")},), + "example-secret-key", + "Key [REDACTED] costs 12.34", + ), + ( + "Key %r", + (Path("example-secret-key"),), + "example-secret-key", + f"Key {Path('[REDACTED]')!r}", + ), + ], + ids=["numeric-secret", "typed-positional", "typed-mapping", "object-repr"], +) +def test_redaction_formats_original_arguments(message, args, secret, expected): + record = logging.LogRecord("test", logging.INFO, __file__, 1, message, args, None) + record.response = {"credential": secret, "count": 7} + redactor = _ValueRedactor((secret,)) + + for formatter in (OpenHoundTextFormatter(), OpenHoundJSONFormatter()): + # Each handler filters the same record; subsequent passes must be safe. + assert redactor.filter(record) + assert record.msg == record.getMessage() == expected + assert record.args == () + assert record.response == {"credential": "[REDACTED]", "count": 7} + # Numeric metadata (timestamps/PIDs) is preserved and may coincidentally + # contain a short numeric test secret. Check rendered/string fields. + assert all( + secret not in value + for value in record.__dict__.values() + if isinstance(value, str) + ) + assert secret not in formatter.format(record) + + +@pytest.mark.parametrize("redacted", [False, True], ids=["original", "redacted"]) +@pytest.mark.parametrize("rich_first", [False, True], ids=["files-first", "rich-first"]) +def test_exception_diagnostics_survive_all_handlers(redacted, rich_first): + """One shared record keeps its traceback exactly once in every output.""" + access_key, secret_key = "example-access-key", "example-secret-key" + outputs = {name: StringIO() for name in ("text", "json", "rich")} + handlers = [] + for name, formatter in ( + ("text", OpenHoundTextFormatter()), + ("json", OpenHoundJSONFormatter()), + ("rich", OpenHoundRichFormatter()), + ): + handler = ( + RichHandler( + console=Console(file=outputs[name], width=240, color_system=None), + show_time=False, + show_path=False, + markup=True, + rich_tracebacks=True, + ) + if name == "rich" + else logging.StreamHandler(outputs[name]) + ) + handler.setFormatter(formatter) + handlers.append(handler) + if rich_first: + handlers.reverse() + logger = logging.getLogger("exception-diagnostics-test") + previous = logger.handlers[:], logger.level, logger.propagate + logger.handlers = handlers + logger.setLevel(logging.ERROR) + logger.propagate = False + try: + with redact_log_values((access_key, secret_key)) if redacted else nullcontext(): + try: + raise RuntimeError(f"AccessDenied for {access_key}: {secret_key}") + except RuntimeError: + logger.exception("Collection failed") + for name, output in outputs.items(): + rendered = output.getvalue() + if name == "json": + document = json.loads(rendered) + assert document["message"] == "Collection failed" + rendered = document["exception"] + assert rendered.count("RuntimeError: AccessDenied") == 1 + assert rendered.count("Traceback (most recent call last)") == 1 + if redacted: + assert access_key not in rendered + assert secret_key not in rendered + assert "[REDACTED]" in rendered + else: + assert access_key in rendered + assert secret_key in rendered + finally: + logger.handlers, logger.level, logger.propagate = previous + for handler in handlers: + handler.close() + + def test_get_file_handler_caches_handler_per_path(tmp_path): """The shared handler cache should return the same instance for a given path and distinct instances for different paths, so each file is only opened once.""" @@ -190,51 +441,59 @@ def test_get_file_handler_caches_handler_per_path(tmp_path): handler.close() -def test_root_and_dlt_loggers_share_single_file_handler(tmp_path, monkeypatch): - """The root and dlt loggers writing to openhound.log must share one handler - instance so the file is only opened once, which is required for rotation on - Windows where an open handle blocks renaming the file.""" - # A sibling module may set RUNTIME__LOG_PATH on import; keep our explicit base_path. - monkeypatch.delenv("RUNTIME__LOG_PATH", raising=False) - custom_logger = CustomLogger("openhound.log", base_path=str(tmp_path)) - +@pytest.mark.parametrize("rich", [False, True], ids=["stdout", "rich-console"]) +def test_parent_worker_routing_deduplicates_shared_console_destinations( + configured_logger, rich +): + output = StringIO() + handlers = [ + RichHandler(console=Console(file=output), show_time=False, show_path=False) + if rich + else logging.StreamHandler(output) + for _ in range(2) + ] + configured_logger.root_logger.addHandler(handlers[0]) + configured_logger.dlt_logger.addHandler(handlers[1]) try: - custom_logger.setup() - - root_logger = logging.getLogger() - dlt_logger = logging.getLogger("dlt") - - root_file_handlers = [ - handler - for handler in root_logger.handlers - if isinstance(handler, RotatingFileHandler) - ] - dlt_file_handlers = [ - handler - for handler in dlt_logger.handlers - if isinstance(handler, RotatingFileHandler) - ] - - assert len(root_file_handlers) == 1, ( - "The root logger should have exactly one rotating file handler" - ) - assert len(dlt_file_handlers) == 1, ( - "The dlt logger should have exactly one rotating file handler" - ) - assert root_file_handlers[0] is dlt_file_handlers[0], ( - "Root and dlt loggers must share the same handler instance for openhound.log" - ) - assert root_file_handlers[0].baseFilename.endswith("openhound.log"), ( - "The shared handler should write to 'openhound.log'" + record = logging.LogRecord( + "dlt", logging.ERROR, __file__, 1, "one-delivery", (), None ) + _ParentWorkerHandler(configured_logger).handle(record) + assert output.getvalue().count("one-delivery") == 1 finally: - for handler in custom_logger._file_handlers.values(): + configured_logger.root_logger.removeHandler(handlers[0]) + configured_logger.dlt_logger.removeHandler(handlers[1]) + for handler in handlers: handler.close() - # Undo the monkeypatched env before restoring global state so setup() - # runs against the real environment, not the modified one. - monkeypatch.undo() - # Restore the shared global logging state for subsequent tests. - logger_override.setup() + + +def test_root_and_dlt_loggers_share_single_file_handler(configured_logger): + """The root and dlt loggers writing to openhound.log must share one handler + instance so the file is only opened once, which is required for rotation on + Windows where an open handle blocks renaming the file.""" + root_file_handlers = [ + handler + for handler in configured_logger.root_logger.handlers + if isinstance(handler, RotatingFileHandler) + ] + dlt_file_handlers = [ + handler + for handler in configured_logger.dlt_logger.handlers + if isinstance(handler, RotatingFileHandler) + ] + + assert len(root_file_handlers) == 1, ( + "The root logger should have exactly one rotating file handler" + ) + assert len(dlt_file_handlers) == 1, ( + "The dlt logger should have exactly one rotating file handler" + ) + assert root_file_handlers[0] is dlt_file_handlers[0], ( + "Root and dlt loggers must share the same handler instance for openhound.log" + ) + assert root_file_handlers[0].baseFilename.endswith("openhound.log"), ( + "The shared handler should write to 'openhound.log'" + ) def test_build_file_handler_applies_rotation_settings(tmp_path): @@ -267,3 +526,276 @@ def test_build_file_handler_applies_rotation_settings(tmp_path): ) finally: handler.close() + + +WORKER_ACCESS = "worker-example-access-key" +WORKER_SECRET = "worker-example-secret-key" + + +def _spawned_log_collection(collector, job_id, params, credentials, managed, directory): + """Run the real job/configuration bridge; replace only collection/network IO.""" + from dlt.common import logger as dlt_logger + + from openhound.core.logging import logger_override + from openhound.scheduler import service + + def collect(*args): + # DLT runtime initialization can configure its logger, as can each of + # OpenHound's collect/preprocess/convert stages. All must keep forwarding. + dlt.pipeline( + pipeline_name="worker_logs", pipelines_dir=str(directory / "pipelines") + ) + logger_override.setup() + CustomLogger("unused.log", base_path="/must/not/be/created").setup() + logger_override.set_handler("aws") + logger_override.set_handler("aws") + loggers = [logging.getLogger(), *logging.Logger.manager.loggerDict.values()] + assert not any( + isinstance(handler, logging.FileHandler) + for logger in loggers + if isinstance(logger, logging.Logger) + for handler in logger.handlers + ), "scheduler workers must not own file handlers" + logging.getLogger("openhound.core.collect").info( + "forwarded-core %s", + WORKER_ACCESS, + extra={ + "details": {"secret": WORKER_SECRET, "values": [False, 0, None]}, + "path": Path("structured-path"), + }, + ) + logging.getLogger("openhound_aws.auth").warning( + "forwarded-extension %s", + WORKER_SECRET, + extra={ + "job_id": "explicit-job", + "collector_extension": "explicit-collector", + }, + ) + dlt_logger.info( + "forwarded-dlt %s", WORKER_SECRET, extra={"metrics": {"rows": 7}} + ) + try: + raise RuntimeError(f"diagnostic {WORKER_SECRET}") + except RuntimeError: + logging.getLogger("openhound.core.collect").exception("forwarded-exception") + logging.getLogger("botocore.worker_test").debug("forbidden-sdk-response") + if params.get("test", {}).get("fail"): + raise RuntimeError(f"worker failure {WORKER_SECRET}") + return service.Result({}, job_id) + + service._collect_in_worker = collect + # Network IO is out of scope for this logging-only worker. + service.validate_aws_profile = lambda profile: None + return service._subprocess_collect( + collector, job_id, params, credentials, managed, directory + ) + + +def _spawned_log_then_die(): + logging.getLogger("openhound.core.collect").warning("before-abrupt-death") + os._exit(17) + + +def _spawned_busy_logger(): + logger = logging.getLogger("openhound.core.collect") + logger.warning("before-abrupt-death") + while True: + logger.info("worker-writing", extra={"payload": "x" * 65536}) + + +@pytest.fixture +def logging_scheduler(configured_logger, monkeypatch): + from openhound.core import logging as openhound_logging + from openhound.scheduler.service import Service + + monkeypatch.setattr(openhound_logging, "logger_override", configured_logger) + scheduler = Service("https://bhe.example.test", "token", "id", "aws", managed=False) + monkeypatch.setattr(scheduler.client, "start_collector_job", lambda job_id: None) + monkeypatch.setattr( + scheduler.client, "end_managed_job", lambda *args, **kwargs: None + ) + try: + yield scheduler + finally: + scheduler.executor.shutdown(wait=True, cancel_futures=True) + scheduler._cleanup_aws_profile_workspace() + scheduler._worker_logging.close() + + +def test_real_spawned_worker_routes_redacted_records_and_survives_replacement( + logging_scheduler, configured_logger, tmp_path, monkeypatch +): + from openhound.scheduler.aws_profile import AWSAccessKeyCredentials + from openhound.scheduler.configuration import ManagedConfigurationError + + scheduler = logging_scheduler + records = [] + + class ParentOnlyHandler(logging.Handler): + def emit(self, record): + records.append(record) + + handler = ParentOnlyHandler() + logging.getLogger().addHandler(handler) + listener_thread = scheduler._worker_logging.listener._thread + manager_process = scheduler._worker_logging.manager._process + try: + for index in range(3): + if index == 2: + scheduler._reset_executor() + executor = scheduler.executor + submit = executor.submit + with monkeypatch.context() as patch: + patch.setattr( + executor, + "submit", + lambda function, *args, submit=submit: submit( + _spawned_log_collection, *args + ), + ) + scheduler.run_claimed_job( + f"job-{index}", + {"test": {"fail": index == 1}}, + AWSAccessKeyCredentials(WORKER_ACCESS, WORKER_SECRET), + ) + if index == 1: + with pytest.raises(ManagedConfigurationError) as error: + scheduler.future.result(timeout=30) + assert WORKER_SECRET not in str(error.value) + else: + assert scheduler.future.result(timeout=30).job_id == f"job-{index}" + scheduler._handle_completed_job(scheduler.future) + scheduler._shutdown() + assert not listener_thread.is_alive() + assert not manager_process.is_alive() + assert scheduler._worker_logging.listener is None + assert scheduler._worker_logging.manager is None + # Closing twice is safe, including after Service shutdown. + scheduler._worker_logging.close() + + core = [ + record + for record in records + if record.getMessage().startswith("forwarded-core") + ] + extension = [ + record + for record in records + if record.getMessage().startswith("forwarded-extension") + ] + dlt_records = [ + record + for record in records + if record.getMessage().startswith("forwarded-dlt") + ] + exceptions = [ + record for record in records if record.getMessage() == "forwarded-exception" + ] + assert len(core) == len(extension) == len(dlt_records) == len(exceptions) == 3 + assert {record.job_id for record in core} == {"job-0", "job-1", "job-2"} + assert all( + record.collector_extension == "aws" + for record in core + dlt_records + exceptions + ) + assert all( + record.job_id == "explicit-job" + and record.collector_extension == "explicit-collector" + for record in extension + ) + assert all( + record.name == "openhound_aws.auth" and record.levelno == logging.WARNING + for record in extension + ) + assert all( + record.details == {"secret": "[REDACTED]", "values": [False, 0, None]} + and record.path == "structured-path" + for record in core + ) + assert all(record.metrics == {"rows": 7} for record in dlt_records) + assert all( + "RuntimeError: diagnostic [REDACTED]" in record.exc_text + and "collect" in record.exc_text + for record in exceptions + ) + assert any( + "worker failure [REDACTED]" in (record.exc_text or "") for record in records + ) + assert "forbidden-sdk-response" not in repr( + [record.__dict__ for record in records] + ) + for secret in (WORKER_ACCESS, WORKER_SECRET): + assert secret not in repr([record.__dict__ for record in records]) + assert secret not in (tmp_path / "openhound.log").read_text() + assert secret not in (tmp_path / "ext_aws.log").read_text() + main = [ + json.loads(line) + for line in (tmp_path / "openhound.log").read_text().splitlines() + ] + assert sum(row["message"].startswith("forwarded-dlt") for row in main) == 3 + extension_file = [ + json.loads(line) + for line in (tmp_path / "ext_aws.log").read_text().splitlines() + ] + assert ( + sum(row["message"].startswith("forwarded-dlt") for row in extension_file) + == 3 + ) + finally: + logging.getLogger().removeHandler(handler) + + +@pytest.mark.parametrize("claim_lost", [False, True], ids=["broken-pool", "claim-loss"]) +def test_worker_logging_abrupt_termination_and_shutdown_do_not_hang( + logging_scheduler, claim_lost, tmp_path +): + scheduler = logging_scheduler + listener_thread = scheduler._worker_logging.listener._thread + if claim_lost: + ready = threading.Event() + + class ReadyHandler(logging.Handler): + def emit(self, record): + if record.getMessage() == "worker-writing": + ready.set() + + handler = ReadyHandler() + logging.getLogger().addHandler(handler) + try: + future = scheduler.executor.submit(_spawned_busy_logger) + scheduler.future = future + assert ready.wait(30), "worker did not begin forwarding records" + # Kill a live producer while it is continuously forwarding records, + # rather than only exercising orderly queue submissions before exit. + scheduler._abort_managed_job("job") + with pytest.raises(BrokenProcessPool): + future.result(timeout=30) + finally: + logging.getLogger().removeHandler(handler) + else: + scheduler.future = scheduler.executor.submit(_spawned_log_then_die) + with pytest.raises(BrokenProcessPool): + scheduler.future.result(timeout=30) + scheduler._reset_executor() + scheduler.future = None + # A replacement initialized through either reset path still forwards. + future = scheduler.executor.submit(_spawned_log_then_die) + with pytest.raises(BrokenProcessPool): + future.result(timeout=30) + scheduler.future = None + finished = threading.Event() + + def shutdown(): + scheduler._shutdown() + finished.set() + + thread = threading.Thread(target=shutdown, daemon=True) + thread.start() + assert finished.wait(10), "shutdown hung after abrupt worker termination" + thread.join(timeout=1) + assert not listener_thread.is_alive() + main = [ + json.loads(line) + for line in (tmp_path / "openhound.log").read_text().splitlines() + ] + assert sum(row["message"] == "before-abrupt-death" for row in main) == 2 diff --git a/tests/test_managed_configuration.py b/tests/test_managed_configuration.py new file mode 100644 index 00000000..3ec429f0 --- /dev/null +++ b/tests/test_managed_configuration.py @@ -0,0 +1,331 @@ +import json +import logging +import os +import traceback +from copy import deepcopy +from pathlib import Path +from types import SimpleNamespace + +import dlt +import pytest +from dlt.common.configuration import configspec +from dlt.common.configuration.container import Container +from dlt.common.configuration.specs import BaseConfiguration, CredentialsConfiguration +from dlt.common.configuration.specs.pluggable_run_context import PluggableRunContext +from dlt.common.configuration.utils import TraceLogContext + +from openhound.scheduler import service +from openhound.scheduler.configuration import ( + ManagedConfigurationError, + managed_job_configuration, +) + +DATA = Path(__file__).parent / "test_data" + + +# The relevant openhound-aws model shapes. Exercise the actual DLT source +# resolver without requiring the AWS extension's native build or AWS requests. +@configspec +class BucketPolicyRule(BaseConfiguration): + include: bool = True + match_bucket: str | None = None + match_object: str | None = None + + +@configspec +class BucketPolicy(BaseConfiguration): + include: bool = False + max_objects: int = 10 + rule: list[BucketPolicyRule] | None = None + + +@configspec +class AWSProfiles(CredentialsConfiguration): + profile_names: list[str] = None + region_name: str | None = "us-east-1" + endpoint_url: str | None = None + + +@dlt.source(name="aws", section="source") +def aws_configuration_source( + credentials: AWSProfiles = dlt.secrets.value, + bucket_policy: BucketPolicy = dlt.config.value, +): + rules = [ + BucketPolicyRule.from_init_value(rule) for rule in bucket_policy.rule or [] + ] + return dlt.resource( + [ + { + "profiles": credentials.profile_names, + "region": credentials.region_name, + "endpoint": credentials.endpoint_url, + "include": bucket_policy.include, + "max_objects": bucket_policy.max_objects, + "bucket_rules": [ + (rule.match_bucket, rule.include) + for rule in rules + if rule.match_bucket + ], + "object_rules": [ + (rule.match_object, rule.include) + for rule in rules + if rule.match_object + ], + } + ], + name="resolved_settings", + ) + + +@pytest.fixture +def params(): + return json.loads((DATA / "managed_aws_configuration.json").read_text()) + + +@pytest.fixture +def secrets(): + return json.loads((DATA / "managed_aws_secrets.json").read_text()) + + +@pytest.mark.parametrize("include,max_objects", [(True, 50), (False, 0)]) +def test_dlt_source_resolves_nested_aws_settings_and_rule_arrays( + params, include, max_objects +): + bucket = params["sources"]["source"]["aws"]["bucket_policy"] + bucket.update(include=include, max_objects=max_objects) + with managed_job_configuration(params, {}): + (row,) = list(aws_configuration_source()) + + assert row == { + "profiles": ["bhe-profile"], + "region": "us-west-2", + "endpoint": "https://aws.example.test", + "include": include, + "max_objects": max_objects, + "bucket_rules": [(".*", True)], + "object_rules": [("private/.*", False)], + } + + +def test_dlt_resolves_retrieved_secrets_per_setting_and_preserves_siblings( + params, secrets, monkeypatch, caplog +): + monkeypatch.setenv("SOURCES__SOURCE__AWS__CREDENTIALS__PROFILE_NAMES", '["local"]') + # The BHE URL and bootstrap secret reference belong to deployment config; + # the job documents supply only collector settings and collector secrets. + monkeypatch.setenv( + "DESTINATION__BLOODHOUNDENTERPRISE__URL", "https://bhe.example.test" + ) + monkeypatch.setenv("OPENHOUND__AWS_SECRETS_MANAGER__SECRET_NAME", "production/bhe") + original_params, original_secrets = deepcopy(params), deepcopy(secrets) + providers = Container()[PluggableRunContext].providers + previous = providers.providers + previous_traces = Container()[TraceLogContext] + with caplog.at_level(logging.DEBUG), managed_job_configuration(params, secrets): + (row,) = list(aws_configuration_source()) + assert row["profiles"] == ["retrieved-profile-a", "retrieved-profile-b"] + assert row["region"] == "us-west-2" + assert row["endpoint"] == "https://aws.example.test" + assert row["include"] is True + assert row["max_objects"] == 50 + assert ( + dlt.config.get("destination.bloodhoundenterprise.url", str) + == "https://bhe.example.test" + ) + assert ( + dlt.config.get("openhound.aws_secrets_manager.secret_name", str) + == "production/bhe" + ) + assert Container()[TraceLogContext].all_traces == [] + + assert providers.providers is previous + assert Container()[TraceLogContext] is previous_traces + assert params == original_params + assert secrets == original_secrets + for credential in ("retrieved-profile-a", "retrieved-profile-b"): + assert credential not in caplog.text + + +def test_dynamic_nested_mappings_lists_and_falsy_values_resolve(params): + options = { + "new_nested_section": {"enabled": False, "limit": 0}, + "labels": ["one", "two"], + "empty_list": [], + "empty_mapping": {}, + } + params["sources"]["source"]["aws"]["future_options"] = options + overrides = { + "sources": { + "source": { + "aws": {"future_options": {"new_nested_section": {"enabled": True}}} + } + } + } + with managed_job_configuration(params, overrides): + resolved = dlt.config.get("sources.source.aws.future_options", dict) + assert resolved == { + **options, + "new_nested_section": {"enabled": True, "limit": 0}, + } + assert dlt.config.get("sources.source.aws.future_options.labels", list) == [ + "one", + "two", + ] + assert ( + dlt.config.get("sources.source.aws.future_options.empty_list", list) == [] + ) + + +def test_retrieved_false_zero_and_rule_list_replace_matching_settings(params, secrets): + secrets["sources"]["source"]["aws"]["bucket_policy"] = { + "include": False, + "max_objects": 0, + "rule": [{"include": True, "match_object": "public/.*"}], + } + with managed_job_configuration(params, secrets): + (row,) = list(aws_configuration_source()) + assert row["include"] is False + assert row["max_objects"] == 0 + assert row["bucket_rules"] == [] + assert row["object_rules"] == [("public/.*", True)] + assert row["endpoint"] == "https://aws.example.test" + + +@pytest.mark.parametrize("failed", [False, True], ids=["successful-job", "failed-job"]) +def test_successive_jobs_do_not_inherit_configuration_or_secrets( + params, secrets, failed +): + original_environment = dict(os.environ) + try: + with managed_job_configuration(params, secrets): + assert next(iter(aws_configuration_source()))["profiles"] == [ + "retrieved-profile-a", + "retrieved-profile-b", + ] + if failed: + raise RuntimeError("collection failed") + except RuntimeError: + assert failed + + second_job = { + "sources": { + "source": {"aws": {"credentials": {"profile_names": ["second-job"]}}} + } + } + with managed_job_configuration(second_job, {}): + (row,) = list(aws_configuration_source()) + assert row["profiles"] == ["second-job"] + assert row["region"] == "us-east-1" + assert row["endpoint"] is None + assert row["include"] is False + assert row["max_objects"] == 10 + assert row["bucket_rules"] == row["object_rules"] == [] + assert dict(os.environ) == original_environment + + +@pytest.mark.parametrize( + "document", + [ + {"max_objects": 50}, + {"sources.source.aws.bucket_policy.include": True}, + {"SOURCES__SOURCE__AWS__BUCKET_POLICY__INCLUDE": True}, + {"profile_names": ["private-credential"]}, + {"sources": {}, "private-key": "private-credential"}, + ], +) +@pytest.mark.parametrize("secret_document", [False, True]) +def test_flat_configuration_is_rejected_without_exposing_keys_or_values( + params, document, secret_document +): + with ( + pytest.raises( + ManagedConfigurationError, match="flat configuration is not supported" + ) as error, + managed_job_configuration( + params if secret_document else document, document if secret_document else {} + ), + ): + pytest.fail("flat configuration must not reach the pipeline") + assert "private-credential" not in str(error.value) + assert "private-key" not in str(error.value) + + +def test_empty_params_are_rejected_even_with_retrieved_secrets(secrets): + with ( + pytest.raises(ManagedConfigurationError, match="non-empty nested JSON object"), + managed_job_configuration({}, secrets), + ): + pytest.fail("empty params must not reach the pipeline") + + +def test_custom_sections_and_values_are_passed_to_dlt_without_schema_validation(): + settings = { + "future_options": {"optional_value": None, "labels": [], "limit": 0}, + "enabled": False, + } + with managed_job_configuration({"custom_pipeline": settings}, {}): + assert dlt.config.get("custom_pipeline", dict) == settings + + +def test_unused_nested_settings_are_ignored_during_dlt_source_resolution(params): + with managed_job_configuration(params, {}): + expected = list(aws_configuration_source()) + + params["unrelated_pipeline"] = {"extra_metadata": ["ignored"]} + aws = params["sources"]["source"]["aws"] + aws["unused_source_setting"] = {"arbitrary": None} + aws["credentials"]["unused_credential_setting"] = [False, 0] + aws["bucket_policy"]["unused_bucket_setting"] = "not a collector option" + for rule in aws["bucket_policy"]["rule"]: + rule["unused_rule_setting"] = {"arbitrary": None} + + with managed_job_configuration(params, {}): + assert list(aws_configuration_source()) == expected + + +def test_dlt_type_errors_are_sanitized_and_provider_is_removed(params, caplog): + providers = Container()[PluggableRunContext].providers + previous = providers.providers + params["sources"]["source"]["aws"]["bucket_policy"][ + "max_objects" + ] = "private-credential" + with ( + caplog.at_level(logging.DEBUG), + pytest.raises(ManagedConfigurationError) as error, + managed_job_configuration(params, {}), + ): + list(aws_configuration_source()) + assert "could not be resolved" in str(error.value) + assert "private-credential" not in "".join(traceback.format_exception(error.value)) + assert "private-credential" not in caplog.text + assert providers.providers is previous + + +def test_worker_applies_configuration_before_collector_loading( + params, secrets, monkeypatch +): + observed = [] + + def load_collectors(): + observed.append(next(iter(aws_configuration_source()))) + return SimpleNamespace( + collectors=[ + SimpleNamespace( + name="aws", + dlt_source=aws_configuration_source, + package_version="test", + metadata=None, + ) + ] + ) + + monkeypatch.setattr(service.CollectorManager, "from_entrypoint", load_collectors) + monkeypatch.setattr( + service.dataflow, + "pipeline", + lambda extension: {"settings": list(extension.dlt_source())}, + ) + result = service._subprocess_collect("aws", "job-1", params, secrets) + assert result.results["settings"] == observed + assert observed[0]["profiles"] == ["retrieved-profile-a", "retrieved-profile-b"] diff --git a/tests/test_scheduler_dataflow.py b/tests/test_scheduler_dataflow.py new file mode 100644 index 00000000..3df45af9 --- /dev/null +++ b/tests/test_scheduler_dataflow.py @@ -0,0 +1,59 @@ +from pathlib import Path +from types import SimpleNamespace + +import pytest + +from openhound.core.convert import Method +from openhound.core.progress import Progress +from openhound.scheduler import dataflow + + +@pytest.mark.parametrize("output_path", [None, Path("job-output")]) +def test_conversion_stages_files_only_when_output_path_is_provided( + monkeypatch, output_path +): + calls = [] + extension = SimpleNamespace( + name="example", converter=lambda **kwargs: calls.append(kwargs) + ) + monkeypatch.setattr(dataflow, "extracted_resources", lambda info: ["nodes"]) + + assert dataflow.convert(extension, output_path) == ["nodes"] + + expected = { + "input_path": Path("/tmp/openhound/example"), + "method": Method.ingest if output_path is None else Method.write, + "progress": Progress.log, + } + if output_path is not None: + expected["output_path"] = output_path + assert calls == [expected] + + +def test_pipeline_finishes_collection_and_preprocessing_before_conversion(monkeypatch): + calls = [] + extension = SimpleNamespace(preprocessor=object()) + output_path = Path("job-output") + + def collect(extension): + calls.append("collect") + return ["source"] + + def preprocess(extension): + calls.append("preprocess") + return ["lookup"] + + def convert(extension, destination): + calls.append(("convert", destination)) + return ["nodes"] + + monkeypatch.setattr(dataflow, "collect", collect) + monkeypatch.setattr(dataflow, "preprocess", preprocess) + monkeypatch.setattr(dataflow, "convert", convert) + + assert dataflow.pipeline(extension, output_path) == { + "collect": ["source"], + "preprocess": ["lookup"], + "convert": ["nodes"], + } + assert calls == ["collect", "preprocess", ("convert", output_path)] diff --git a/tests/test_scheduler_start.py b/tests/test_scheduler_start.py new file mode 100644 index 00000000..0c0938a6 --- /dev/null +++ b/tests/test_scheduler_start.py @@ -0,0 +1,104 @@ +import scheduler +from openhound.config import ManagedBHEConfig +from openhound.core.clients.bhe_credentials import BHECredentials + + +def test_managed_scheduler_loads_aws_secret_before_starting_service(monkeypatch): + constructed = [] + started = [] + + class FakeCredentials: + def __init__(self, secret_name): + constructed.append(secret_name) + + def poll_until_available(self, interval): + assert interval == 45 + return BHECredentials(token_id="aws-id", token_key="aws-key") + + def refresh(self): + return BHECredentials(token_id="refreshed-id", token_key="refreshed-key") + + class FakeService: + def __init__(self, **kwargs): + started.append(kwargs) + + def start(self): + started.append("started") + + monkeypatch.setattr(scheduler.config, "is_managed", lambda: True) + monkeypatch.setattr( + scheduler.config, + "get_managed_bhe_config", + lambda: ManagedBHEConfig(secret_name="production/bhe", poll_interval=45), + ) + monkeypatch.setattr( + scheduler.dlt, + "config", + { + "destination.bloodhoundenterprise.url": "https://bhe.example", + "destination.bloodhoundenterprise.collector_name": "collector", + }, + ) + monkeypatch.setattr(scheduler, "AWSBHECredentials", FakeCredentials) + monkeypatch.setattr(scheduler, "Service", FakeService) + + scheduler.start() + + assert constructed == ["production/bhe"] + assert started[0]["bhe_uri"] == "https://bhe.example" + assert started[0]["token_key"] == "aws-key" + assert started[0]["token_id"] == "aws-id" + assert started[0]["collector_name"] == "collector" + assert started[0]["interval"] == 45 + assert started[0]["credential_refresh"]() == BHECredentials( + token_id="refreshed-id", token_key="refreshed-key" + ) + assert started[1] == "started" + + +def test_unmanaged_scheduler_uses_dlt_secrets_without_credential_refresh(monkeypatch): + started = [] + + class FakeService: + def __init__(self, **kwargs): + started.append(kwargs) + + def start(self): + started.append("started") + + def unexpected_aws_credentials(*args, **kwargs): + raise AssertionError("unmanaged scheduler must not load AWS credentials") + + monkeypatch.setattr(scheduler.config, "is_managed", lambda: False) + monkeypatch.setattr( + scheduler.dlt, + "config", + { + "destination.bloodhoundenterprise.url": "https://bhe.example", + "destination.bloodhoundenterprise.collector_name": "collector", + }, + ) + monkeypatch.setattr( + scheduler.dlt, + "secrets", + { + "destination.bloodhoundenterprise.token_key": "dlt-key", + "destination.bloodhoundenterprise.token_id": "dlt-id", + }, + ) + monkeypatch.setattr(scheduler, "AWSBHECredentials", unexpected_aws_credentials) + monkeypatch.setattr(scheduler, "Service", FakeService) + + scheduler.start() + + assert started == [ + { + "bhe_uri": "https://bhe.example", + "token_key": "dlt-key", + "token_id": "dlt-id", + "collector_name": "collector", + "interval": 30, + "credential_refresh": None, + }, + "started", + ] diff --git a/uv.lock b/uv.lock index a4e444e2..9aac18a0 100644 --- a/uv.lock +++ b/uv.lock @@ -159,6 +159,34 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/cd/0c/31cfaa6b56fe23488ecb993bc9fc526c0d84d89607decdf2a10776426c2e/binaryornot-0.6.0-py3-none-any.whl", hash = "sha256:900adfd5e1b821255ba7e63139b0396b14c88b9286e74e03b6f51e0200331337", size = 14185, upload-time = "2026-03-08T16:26:27.466Z" }, ] +[[package]] +name = "boto3" +version = "1.43.90" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "botocore" }, + { name = "jmespath" }, + { name = "s3transfer" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/67/0b/53f833efea7331aab106a65f1156e0035beed9bacbfcb309b128204468d4/boto3-1.43.90.tar.gz", hash = "sha256:4b669742d5b45b8fd20ca50ac414a4e4cf995ebb8f280d21be28676e71c97594", size = 112689, upload-time = "2026-09-08T19:22:30.875Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/4b/c3/1d5c5ea7f599e2b80d122beeb2a9c1a8daff78325ea2391521e5fd5af1fd/boto3-1.43.90-py3-none-any.whl", hash = "sha256:aaaa1216d65ddb3dcf86bf9d93cfa436af05cbcd1e1f8b2847f2b83116012186", size = 140026, upload-time = "2026-09-08T19:22:28.555Z" }, +] + +[[package]] +name = "botocore" +version = "1.43.90" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "jmespath" }, + { name = "python-dateutil" }, + { name = "urllib3" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/6b/3e/29872261075d878581a31e9e98c512c0b85322a36e6d6784b9f0694a918d/botocore-1.43.90.tar.gz", hash = "sha256:a139ed601e8b8fb1d730022355fe2b284b8c15cfe9ac0f100254a35d2273e1d3", size = 16081908, upload-time = "2026-09-08T19:22:25.645Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/f6/b8/65cbd0d65553c3e51d27159de4c74308ea008c652b05ed3f830f32667efb/botocore-1.43.90-py3-none-any.whl", hash = "sha256:65f3394ef07314e45c92a90120988531d651136390d464a94938a92f665893f7", size = 15775313, upload-time = "2026-09-08T19:22:22.713Z" }, +] + [[package]] name = "certifi" version = "2026.10.23" @@ -800,6 +828,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/62/a1/3d680cbfd5f4b8f15abc1d571870c5fc3e594bb582bc3b64ea099db13e56/jinja2-3.1.6-py3-none-any.whl", hash = "sha256:85ece4451f492d0c13c5dd7c13a64681a86afae63a5f347908daf103ce6d2f67", size = 134899, upload-time = "2025-03-05T20:05:00.369Z" }, ] +[[package]] +name = "jmespath" +version = "1.1.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/d3/59/322338183ecda247fb5d1763a6cbe46eff7222eaeebafd9fa65d4bf5cb11/jmespath-1.1.0.tar.gz", hash = "sha256:472c87d80f36026ae83c6ddd0f1d05d4e510134ed462851fd5f754c8c3cbb88d", size = 27377, upload-time = "2026-01-22T16:35:26.279Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/14/2f/967ba146e6d58cf6a652da73885f52fc68001525b4197effc174321d70b4/jmespath-1.1.0-py3-none-any.whl", hash = "sha256:a5663118de4908c91729bea0acadca56526eb2698e83de10cd116ae0f4e97c64", size = 20419, upload-time = "2026-01-22T16:35:24.919Z" }, +] + [[package]] name = "joserfc" version = "1.7.4" @@ -1269,9 +1306,11 @@ name = "openhound" source = { editable = "." } dependencies = [ { name = "alive-progress" }, + { name = "boto3" }, { name = "cookiecutter" }, { name = "dlt" }, { name = "duckdb" }, + { name = "filelock" }, { name = "griffe" }, { name = "griffe-fieldz" }, { name = "jinja2" }, @@ -1320,9 +1359,11 @@ dev = [ [package.metadata] requires-dist = [ { name = "alive-progress", specifier = ">=3.3.0" }, + { name = "boto3", specifier = ">=1.43.90" }, { name = "cookiecutter", specifier = ">=2.6.0" }, { name = "dlt", specifier = "==1.26.0" }, { name = "duckdb", specifier = "==1.5.2" }, + { name = "filelock", specifier = ">=3.32.0" }, { name = "griffe", specifier = ">=1.15.0" }, { name = "griffe-fieldz", specifier = ">=0.5.0" }, { name = "jinja2", specifier = ">=3.1.6" }, @@ -2091,6 +2132,18 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/cb/9a/8415f2657cbe200f41a4531ccededf135505a92d4a012229121f885b26f9/ruff-0.16.0-py3-none-win_arm64.whl", hash = "sha256:14296fedcd2705c77ab8235439278bbb38f285cf7da5528b00b3e330c3d4872d", size = 11273407, upload-time = "2026-07-23T19:11:28.705Z" }, ] +[[package]] +name = "s3transfer" +version = "0.19.2" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "botocore" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/76/43/35e4d8aa320bffe8287fe8f65f578fa2d2db0a64212f0e710dce58267854/s3transfer-0.19.2.tar.gz", hash = "sha256:ba0309fd86be3c27dbf78cdd813c13c5e1df16e5874b99d2535ebbdfb9892993", size = 165592, upload-time = "2026-07-22T19:30:44.432Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/bc/e7/5c595c75e9f41a44f30e526eda465ea0b4eec93470e074e4a111b253f13a/s3transfer-0.19.2-py3-none-any.whl", hash = "sha256:d8168eccca828cbb2cd573675333f3bddd254313a9c42494b84c76b539e8ba25", size = 90216, upload-time = "2026-07-22T19:30:43.251Z" }, +] + [[package]] name = "semver" version = "3.0.4"