From 6e6064d291381090f19be62ba26eb923b2ec3ccd Mon Sep 17 00:00:00 2001 From: Adrian Hannah Date: Tue, 8 Sep 2026 09:45:29 -0400 Subject: [PATCH 01/34] feat(config): add managed mode foundation Add a shared managed-mode check and empty mode profiles for future managed collector behavior. ref: BED-9511 --- src/openhound/config.py | 32 +++++++++++++++++++++++++++++ tests/test_config.py | 45 +++++++++++++++++++++++++++++++++++++++++ 2 files changed, 77 insertions(+) create mode 100644 src/openhound/config.py create mode 100644 tests/test_config.py diff --git a/src/openhound/config.py b/src/openhound/config.py new file mode 100644 index 0000000..b826c89 --- /dev/null +++ b/src/openhound/config.py @@ -0,0 +1,32 @@ +"""OpenHound managed-mode configuration.""" + +from dataclasses import dataclass + +import dlt + + +@dataclass(frozen=True) +class ModeDefaults: + """Placeholder for defaults that differ between OpenHound modes.""" + + +_UNMANAGED_DEFAULTS = ModeDefaults() +_MANAGED_DEFAULTS = ModeDefaults() + + +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 diff --git a/tests/test_config.py b/tests/test_config.py new file mode 100644 index 0000000..3b632e8 --- /dev/null +++ b/tests/test_config.py @@ -0,0 +1,45 @@ +import dlt +import pytest +from dlt.common.configuration.exceptions import ConfigValueCannotBeCoercedException + +from openhound.config import ModeDefaults, 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_distinct_empty_profiles(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 vars(unmanaged) == {} + assert vars(managed) == {} From 6caeb2cb4ed43b58da3e8d46b224fc62a0197560 Mon Sep 17 00:00:00 2001 From: Adrian Hannah Date: Tue, 8 Sep 2026 09:46:36 -0400 Subject: [PATCH 02/34] docs(config): document managed mode default Show the unmanaged default explicitly in the BloodHound Enterprise example configuration. ref: BED-9511 --- .../bloodhound-enterprise/.dlt-example/config.toml | 3 +++ 1 file changed, 3 insertions(+) diff --git a/example-configurations/bloodhound-enterprise/.dlt-example/config.toml b/example-configurations/bloodhound-enterprise/.dlt-example/config.toml index 8da76f8..0cdc72d 100644 --- a/example-configurations/bloodhound-enterprise/.dlt-example/config.toml +++ b/example-configurations/bloodhound-enterprise/.dlt-example/config.toml @@ -1,4 +1,7 @@ # Example configuration: https://bloodhound.specterops.io/openhound/enterprise#full-configuration-example +[openhound] +managed = false + [runtime] http_show_error_body = true log_cli_level = "WARNING" From 4e430663cd858b95e014ffbe1e5060877e9eca67 Mon Sep 17 00:00:00 2001 From: Adrian Hannah Date: Tue, 15 Sep 2026 10:19:02 -0400 Subject: [PATCH 03/34] feat(logging): default managed mode to structured logs Use the managed-mode profile to select newline-delimited JSON file logging while preserving text output for self-managed collectors and explicit runtime overrides. Refresh cached file handlers when the selected mode changes and cover the behavior with focused tests. ref: BED-9511 --- .../.dlt-example/config.toml | 5 +- src/openhound/config.py | 8 +- src/openhound/core/logging.py | 19 +++- tests/test_config.py | 6 +- tests/test_log_handlers.py | 96 ++++++++++++++++++- 5 files changed, 117 insertions(+), 17 deletions(-) diff --git a/example-configurations/bloodhound-enterprise/.dlt-example/config.toml b/example-configurations/bloodhound-enterprise/.dlt-example/config.toml index 0cdc72d..c486a50 100644 --- a/example-configurations/bloodhound-enterprise/.dlt-example/config.toml +++ b/example-configurations/bloodhound-enterprise/.dlt-example/config.toml @@ -13,8 +13,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/src/openhound/config.py b/src/openhound/config.py index b826c89..f40d1e1 100644 --- a/src/openhound/config.py +++ b/src/openhound/config.py @@ -7,11 +7,13 @@ @dataclass(frozen=True) class ModeDefaults: - """Placeholder for defaults that differ between OpenHound modes.""" + """Configuration defaults that differ between OpenHound modes.""" + log_format: str -_UNMANAGED_DEFAULTS = ModeDefaults() -_MANAGED_DEFAULTS = ModeDefaults() + +_UNMANAGED_DEFAULTS = ModeDefaults(log_format="text") +_MANAGED_DEFAULTS = ModeDefaults(log_format="json") def is_managed() -> bool: diff --git a/src/openhound/core/logging.py b/src/openhound/core/logging.py index 55ca595..7e932db 100644 --- a/src/openhound/core/logging.py +++ b/src/openhound/core/logging.py @@ -13,10 +13,12 @@ from pathlib import Path import dlt -import openhound from rich.console import Console from rich.logging import RichHandler +import openhound +from openhound.config import get_mode_defaults + VALID_LEVELS = ["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"] VALID_FORMATS = ["json", "text"] @@ -226,7 +228,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 +239,15 @@ 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.handlers = { LogMode.CLI: self.cli_handlers, @@ -342,8 +345,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 @@ -402,6 +407,10 @@ def _get_file_handler(self, file_path: Path) -> "RotatingFileHandler": 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 handler def container_handlers(self, logger: logging.Logger, file_path: Path) -> None: diff --git a/tests/test_config.py b/tests/test_config.py index 3b632e8..87a673f 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -31,7 +31,7 @@ def test_is_managed_rejects_invalid_boolean(monkeypatch): is_managed() -def test_get_mode_defaults_selects_distinct_empty_profiles(monkeypatch): +def test_get_mode_defaults_selects_mode_specific_log_formats(monkeypatch): monkeypatch.setattr("openhound.config.is_managed", lambda: False) unmanaged = get_mode_defaults() @@ -41,5 +41,5 @@ def test_get_mode_defaults_selects_distinct_empty_profiles(monkeypatch): assert isinstance(unmanaged, ModeDefaults) assert isinstance(managed, ModeDefaults) assert unmanaged is not managed - assert vars(unmanaged) == {} - assert vars(managed) == {} + assert unmanaged.log_format == "text" + assert managed.log_format == "json" diff --git a/tests/test_log_handlers.py b/tests/test_log_handlers.py index d83373e..cdb4c30 100644 --- a/tests/test_log_handlers.py +++ b/tests/test_log_handlers.py @@ -79,10 +79,9 @@ def test_log_routing_content(tmp_path, caplog, monkeypatch): 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 +126,95 @@ def test_file_formatter_selection(): ) +def test_managed_mode_writes_structured_file_logs(tmp_path, monkeypatch): + """Managed mode should default rotating file logs to newline-delimited JSON.""" + monkeypatch.setenv("OPENHOUND__MANAGED", "true") + monkeypatch.delenv("RUNTIME__LOG_FORMAT", raising=False) + 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): + """Re-running setup should apply a changed mode to an existing log file.""" + monkeypatch.setenv("RUNTIME__LOG_PATH", str(tmp_path)) + monkeypatch.delenv("RUNTIME__LOG_FORMAT", raising=False) + 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() From 605bcf4f6cbc0552f71dac5172393b45947075ea Mon Sep 17 00:00:00 2001 From: Stran Dutton <65475509+StranDutton@users.noreply.github.com> Date: Tue, 15 Sep 2026 14:00:01 -0500 Subject: [PATCH 04/34] BED-9509: add aws secrets manager service (#83) * BED-9509: add aws secrets manager service --- .github/workflows/test.yml | 4 + pyproject.toml | 1 + .../core/clients/aws_secrets_manager.py | 295 ++++++++++++++++ tests/test_aws_secrets_manager.py | 333 ++++++++++++++++++ uv.lock | 51 +++ 5 files changed, 684 insertions(+) create mode 100644 src/openhound/core/clients/aws_secrets_manager.py create mode 100644 tests/test_aws_secrets_manager.py diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 43575ac..bdc53f0 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -42,6 +42,10 @@ 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 log handling tests run: | .venv/bin/pytest tests/test_log_handlers.py -v diff --git a/pyproject.toml b/pyproject.toml index 064a710..79ee26d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -19,6 +19,7 @@ dependencies = [ "jinja2>=3.1.6", "types-requests==2.33.0.20260503", "pydantic-extra-types>=2.11.1", + "boto3>=1.43.90", ] [project.optional-dependencies] 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 0000000..567369c --- /dev/null +++ b/src/openhound/core/clients/aws_secrets_manager.py @@ -0,0 +1,295 @@ +import json +import logging +from collections.abc import Iterable, Mapping +from typing import Any, Protocol + +import boto3 +from botocore.exceptions import ( + BotoCoreError, + ClientError, + ConfigNotFound, + CredentialRetrievalError, + InvalidConfigError, + NoCredentialsError, + NoRegionError, + PartialCredentialsError, + ProfileNotFound, +) + +logger = logging.getLogger(__name__) + +MAX_BATCH_SIZE = 20 +_PROVIDER = "aws.secretsmanager" + + +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.""" + + +class SecretNotFoundError(SecretRetrievalError): + """The requested secret does not exist.""" + + +class SecretPermissionError(SecretRetrievalError): + """The configured AWS identity cannot read the requested secret.""" + + +class AWSConfigurationError(SecretRetrievalError): + """The AWS SDK or Secrets Manager client is not configured correctly.""" + + +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__("One or more AWS secret retrievals failed") + + +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) as error: + classified = _classify_aws_error(error) + 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) + + 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: + if self._client is not None: + return self._client + try: + return boto3.client("secretsmanager") + except BotoCoreError 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) as error: + classified = _classify_aws_error(error) + _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 _parse_secret_value(response: Mapping[str, Any]) -> SecretValue: + secret_string = response.get("SecretString") + if not isinstance(secret_string, str): + raise SecretRequestError( + "AWS Secrets Manager response did not contain a SecretString payload" + ) + + 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], +) -> 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(entry.get("ErrorCode")) + + 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 | str | None, +) -> SecretRetrievalError: + """Map a boto error or AWS error code to a public retrieval exception.""" + if isinstance(error, ClientError): + error = error.response.get("Error", {}).get("Code") + elif isinstance( + error, + ( + ConfigNotFound, + CredentialRetrievalError, + InvalidConfigError, + NoCredentialsError, + NoRegionError, + PartialCredentialsError, + ProfileNotFound, + ), + ): + return AWSConfigurationError( + "AWS Secrets Manager client configuration is invalid" + ) + if error == "ResourceNotFoundException": + return SecretNotFoundError("AWS secret was not found") + if error in { + "AccessDenied", + "AccessDeniedException", + "AuthorizationError", + "UnauthorizedOperation", + "UnrecognizedClientException", + }: + return SecretPermissionError("AWS Secrets Manager denied access to the secret") + return SecretRequestError("AWS Secrets Manager request could not be completed") + + +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 + logger.error("Secret retrieval failed", extra=extra) diff --git a/tests/test_aws_secrets_manager.py b/tests/test_aws_secrets_manager.py new file mode 100644 index 0000000..98b01e1 --- /dev/null +++ b/tests/test_aws_secrets_manager.py @@ -0,0 +1,333 @@ +import logging +from collections.abc import Mapping, Sequence +from typing import Any, NoReturn, cast + +import pytest +from botocore.exceptions import ( + BotoCoreError, + ClientError, + ConfigNotFound, + CredentialRetrievalError, + EndpointConnectionError, + InvalidConfigError, + NoCredentialsError, + NoRegionError, + 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 = "secret-id-that-must-not-be-logged" +SECRET_VALUE = "secret-value-that-must-not-be-logged" +PROVIDER_MESSAGE = "provider message that must not be exposed" + + +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 = set(secret_ids) | {SECRET_VALUE, PROVIDER_MESSAGE} + + 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 record.levelno == logging.ERROR + ] + assert len(error_records) == 2 + 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_safe_failure(raised.value, {SECRET_ID, PROVIDER_MESSAGE}) + + +@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): + AWSSecretsManager().get_secret(SECRET_ID) + + +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) diff --git a/uv.lock b/uv.lock index a4e444e..80fc76b 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,6 +1306,7 @@ name = "openhound" source = { editable = "." } dependencies = [ { name = "alive-progress" }, + { name = "boto3" }, { name = "cookiecutter" }, { name = "dlt" }, { name = "duckdb" }, @@ -1320,6 +1358,7 @@ 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" }, @@ -2091,6 +2130,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" From f514f16aa134288aeaea95af81dccc7c6ea12556 Mon Sep 17 00:00:00 2001 From: Adrian Hannah Date: Mon, 21 Sep 2026 13:53:07 -0400 Subject: [PATCH 05/34] refactor(managed): generalize transient request detection Generalize the transient request error classifier for reuse across BHE client workflows. ref: BED-9513 --- src/openhound/core/clients/bloodhound_enterprise.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/openhound/core/clients/bloodhound_enterprise.py b/src/openhound/core/clients/bloodhound_enterprise.py index 78de5e7..c4338c1 100644 --- a/src/openhound/core/clients/bloodhound_enterprise.py +++ b/src/openhound/core/clients/bloodhound_enterprise.py @@ -218,7 +218,7 @@ def _file_checksum(path: Path) -> str: return digest.hexdigest() @staticmethod - def _is_transient_support_bundle_error(error: Exception) -> bool: + def _is_transient_request_error(error: Exception) -> bool: if isinstance(error, requests.RequestException): return True return isinstance(error, BloodHoundHTTPError) and error.code in { @@ -237,7 +237,7 @@ 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): raise if retry == SUPPORT_BUNDLE_MAX_RETRIES: raise From a61364b33f61e195f88b0fd83d788ae71cfab193 Mon Sep 17 00:00:00 2001 From: Adrian Hannah Date: Mon, 21 Sep 2026 14:22:54 -0400 Subject: [PATCH 06/34] feat(managed): add managed job completion Add managed job completion with outcome payloads, bounded timeouts, transient retries, and structured logging. ref: BED-9513 --- .../core/clients/bloodhound_enterprise.py | 100 +++++++++++++++++- 1 file changed, 97 insertions(+), 3 deletions(-) diff --git a/src/openhound/core/clients/bloodhound_enterprise.py b/src/openhound/core/clients/bloodhound_enterprise.py index c4338c1..9cc2b57 100644 --- a/src/openhound/core/clients/bloodhound_enterprise.py +++ b/src/openhound/core/clients/bloodhound_enterprise.py @@ -6,19 +6,20 @@ 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, TypeVar -import openhound import requests + from openhound.core.clients.bloodhound import BloodHound, BloodHoundHTTPError from openhound.core.clients.models.jobs import ( + ArtifactUploadSession, JobsAvailable, JobsCurrent, JobsEnd, JobStart, - ArtifactUploadSession, ManagementAvailable, ManagementOperationResult, ManagementOperationStatus, @@ -32,6 +33,15 @@ class JobStatus(str, Enum): FAILED = "failed" +class ManagedJobOutcome(str, Enum): + SUCCEEDED = "succeeded" + FAILED = "failed" + + +MANAGED_JOB_END_MAX_RETRIES = 3 +MANAGED_JOB_END_RETRY_DELAY_SECONDS = 2 +MANAGED_JOB_END_CONNECT_TIMEOUT_SECONDS = 10 +MANAGED_JOB_END_READ_TIMEOUT_SECONDS = 120 SUPPORT_BUNDLE_PART_SIZE = 8 * 1024 * 1024 # 8 MiB SUPPORT_BUNDLE_MAX_RETRIES = 3 SUPPORT_BUNDLE_RETRY_DELAY_SECONDS = 2 @@ -67,6 +77,90 @@ 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 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/collector-jobs/{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() + max_attempts = MANAGED_JOB_END_MAX_RETRIES + 1 + + for attempt in range(1, max_attempts + 1): + log_context = { + "job_id": job_id, + "outcome": outcome.value, + "attempt": attempt, + "max_attempts": max_attempts, + } + logger.info( + "Attempting to end managed collector job %s (%s/%s).", + job_id, + attempt, + 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: + retryable = self._is_transient_request_error(error) + 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}, + exc_info=True, + ) + + if not retryable or attempt == max_attempts: + raise + + 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 = { From d7882cf259b6e832012f2031a6f1f18951aeaa0f Mon Sep 17 00:00:00 2001 From: Adrian Hannah Date: Mon, 21 Sep 2026 14:23:16 -0400 Subject: [PATCH 07/34] test(managed): cover managed job completion Cover managed job payloads, bounded timeouts, retry behavior, structured logging, and terminal errors. ref: BED-9513 --- tests/test_bhe_job_scheduling.py | 175 ++++++++++++++++++++++++++++++- 1 file changed, 174 insertions(+), 1 deletion(-) diff --git a/tests/test_bhe_job_scheduling.py b/tests/test_bhe_job_scheduling.py index 7e3746f..1cc5f1b 100644 --- a/tests/test_bhe_job_scheduling.py +++ b/tests/test_bhe_job_scheduling.py @@ -10,12 +10,16 @@ from urllib.parse import urlsplit import pytest +import requests from fastapi import FastAPI, Request, Response from fastapi.testclient import TestClient from openhound.core.clients import bloodhound, bloodhound_enterprise 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 ( ManagementOperation, ManagementOperationStatus, @@ -51,6 +55,7 @@ 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.start_payload = None app.state.client_update_payload = None app.state.ingested_edges = 0 @@ -88,6 +93,13 @@ async def end_job(body: dict): app.state.end_payload = body return load_json("job_end.json") + @app.post("/api/v2/collector-jobs/{job_id}/end") + async def end_managed_job(job_id: str, body: dict): + app.state.collector_job_end_requests.append( + {"job_id": job_id, "body": body} + ) + return Response(status_code=200) + @app.post("/api/v2/ingest") async def ingest(request: Request): body = await request.body() @@ -324,6 +336,167 @@ 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_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"), + ], + ids=["http-500", "connection-error"], +) +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/collector-jobs/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 == bloodhound_enterprise.MANAGED_JOB_END_MAX_RETRIES + 1 + assert delays == [ + bloodhound_enterprise.MANAGED_JOB_END_RETRY_DELAY_SECONDS + ] * bloodhound_enterprise.MANAGED_JOB_END_MAX_RETRIES + + +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 == [] + + def test_jobs_starts_new_job(mock_service, mock_bloodhound_api): """Runs _check_jobs and checks if the new job is started when available.""" From f7185d36bbda9bcbf89ccaae25830facbea2494a Mon Sep 17 00:00:00 2001 From: Adrian Hannah Date: Tue, 22 Sep 2026 14:22:25 -0400 Subject: [PATCH 08/34] fix(client): narrow managed job retry exceptions Allow each caller to select retryable Requests exceptions while preserving the support-bundle retry policy. ref: BED-9513 --- .../core/clients/bloodhound_enterprise.py | 17 +++++++++++++---- 1 file changed, 13 insertions(+), 4 deletions(-) diff --git a/src/openhound/core/clients/bloodhound_enterprise.py b/src/openhound/core/clients/bloodhound_enterprise.py index 9cc2b57..fb61c55 100644 --- a/src/openhound/core/clients/bloodhound_enterprise.py +++ b/src/openhound/core/clients/bloodhound_enterprise.py @@ -129,7 +129,10 @@ def end_managed_job( ), ) except Exception as error: - retryable = self._is_transient_request_error(error) + retryable = self._is_transient_request_error( + error, + (requests.ConnectionError, requests.Timeout), + ) logger.info( "Managed collector job %s end attempt failed.", job_id, @@ -312,8 +315,11 @@ def _file_checksum(path: Path) -> str: return digest.hexdigest() @staticmethod - def _is_transient_request_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, @@ -331,7 +337,10 @@ def _retry_support_bundle_request( try: return request() except Exception as error: - if not self._is_transient_request_error(error): + if not self._is_transient_request_error( + error, + (requests.RequestException,), + ): raise if retry == SUPPORT_BUNDLE_MAX_RETRIES: raise From 08e66060da168287662e0659c11cf9a2005c655e Mon Sep 17 00:00:00 2001 From: Adrian Hannah Date: Tue, 22 Sep 2026 14:22:56 -0400 Subject: [PATCH 09/34] test(client): cover request exception retry policies Verify managed jobs retry connection and timeout failures but reject invalid URLs and redirects while support bundles retain broad retries. ref: BED-9513 --- tests/test_bhe_job_scheduling.py | 58 +++++++++++++++++++++++++++++++- 1 file changed, 57 insertions(+), 1 deletion(-) diff --git a/tests/test_bhe_job_scheduling.py b/tests/test_bhe_job_scheduling.py index 1cc5f1b..b917794 100644 --- a/tests/test_bhe_job_scheduling.py +++ b/tests/test_bhe_job_scheduling.py @@ -400,8 +400,9 @@ def capture_request(*args, **kwargs): [ BloodHoundHTTPError("server error", 500), requests.ConnectionError("connection failed"), + requests.Timeout("request timed out"), ], - ids=["http-500", "connection-error"], + ids=["http-500", "connection-error", "timeout"], ) def test_end_managed_job_retries_transient_errors( mock_service, monkeypatch, caplog, transient_error @@ -497,6 +498,37 @@ def bad_request(*args, **kwargs): 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.""" @@ -895,6 +927,30 @@ def flaky_request(method, path, **kwargs): assert delays == [2, 2] +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 ): From 6555c64583e5d17de8c342f45c35228ab2664f8c Mon Sep 17 00:00:00 2001 From: Stran Dutton <65475509+StranDutton@users.noreply.github.com> Date: Fri, 25 Sep 2026 11:42:41 -0500 Subject: [PATCH 10/34] Check available jobs (managed OH)- BED-9512 (#86) * Check for available managed collector jobs - BED-9512 * Add timeout, cleanup, and tests * separating into independent func, removing calls * align test data and models with contract --- .../core/clients/bloodhound_enterprise.py | 36 +++++ src/openhound/core/clients/models/jobs.py | 40 ++++- src/openhound/scheduler/service.py | 12 ++ tests/test_bhe_job_scheduling.py | 143 +++++++++++++++++- .../jobs/collector_jobs_available_empty.json | 8 + .../collector_jobs_available_with_job.json | 36 +++++ .../collector_jobs_available_with_jobs.json | 59 ++++++++ 7 files changed, 330 insertions(+), 4 deletions(-) create mode 100644 tests/test_data/api/jobs/collector_jobs_available_empty.json create mode 100644 tests/test_data/api/jobs/collector_jobs_available_with_job.json create mode 100644 tests/test_data/api/jobs/collector_jobs_available_with_jobs.json diff --git a/src/openhound/core/clients/bloodhound_enterprise.py b/src/openhound/core/clients/bloodhound_enterprise.py index fb61c55..8d86eda 100644 --- a/src/openhound/core/clients/bloodhound_enterprise.py +++ b/src/openhound/core/clients/bloodhound_enterprise.py @@ -10,12 +10,14 @@ from enum import Enum from pathlib import Path from typing import Any, TypeVar +from urllib.parse import quote import requests from openhound.core.clients.bloodhound import BloodHound, BloodHoundHTTPError from openhound.core.clients.models.jobs import ( ArtifactUploadSession, + CollectorJobsAvailable, JobsAvailable, JobsCurrent, JobsEnd, @@ -47,6 +49,8 @@ class ManagedJobOutcome(str, Enum): 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 T = TypeVar("T") @@ -64,6 +68,38 @@ 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/collector-job-queue/available?limit=1&job_key=" + f"{encoded_job_key}" + ) + 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): + logger.exception( + "Managed collector job queue request failed.", + extra={"endpoint": path, "job_key": job_key}, + ) + raise + return CollectorJobsAvailable.model_validate(response.json()) + def start_job(self, job_id: int) -> JobStart: path = "/api/v2/jobs/start" body = json.dumps({"id": job_id}) diff --git a/src/openhound/core/clients/models/jobs.py b/src/openhound/core/clients/models/jobs.py index cc218f7..8ecd5a5 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,41 @@ class JobsEnd(BaseModel): data: Job +class CollectorJob(BaseModel): + id: UUID + job_schedule_id: int | None + job_profile_id: int | None + job_type_id: int + job_key: str + params_version: str + params: dict[str, Any] + scope_client_id: UUID | 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/scheduler/service.py b/src/openhound/scheduler/service.py index e64a4b6..a8831f0 100644 --- a/src/openhound/scheduler/service.py +++ b/src/openhound/scheduler/service.py @@ -10,6 +10,7 @@ import openhound.core.logging as openhound_logging from openhound.core.clients.bloodhound_enterprise import BloodHoundEnterprise, JobStatus from openhound.core.clients.models.jobs import ( + CollectorJob, Job, ManagementOperation, ManagementOperationStatus, @@ -170,6 +171,17 @@ 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) + if available_jobs.data.jobs: + logger.info( + "New managed queue job available: %s", available_jobs.data.jobs[0].id + ) + return available_jobs.data.jobs[0] + return 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.") diff --git a/tests/test_bhe_job_scheduling.py b/tests/test_bhe_job_scheduling.py index b917794..ba0cd5c 100644 --- a/tests/test_bhe_job_scheduling.py +++ b/tests/test_bhe_job_scheduling.py @@ -8,6 +8,7 @@ from pathlib import Path from types import SimpleNamespace from urllib.parse import urlsplit +from uuid import UUID import pytest import requests @@ -21,6 +22,8 @@ ManagedJobOutcome, ) from openhound.core.clients.models.jobs import ( + CollectorJob, + CollectorJobsAvailable, ManagementOperation, ManagementOperationStatus, ManagementOperationType, @@ -57,6 +60,13 @@ def mock_bloodhound_api(): app.state.end_payload = None app.state.collector_job_end_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 = [] @@ -73,14 +83,23 @@ def mock_bloodhound_api(): @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/collector-job-queue/available") + async def collector_job_queue(request: Request): + app.state.collector_job_queue_requests.append(dict(request.query_params)) + if app.state.collector_job_queue_error_status is not None: + return Response(status_code=app.state.collector_job_queue_error_status) + return app.state.collector_job_queue_response + @app.post("/api/v2/jobs/start") async def start_job(body: dict): app.state.job_started = True @@ -223,7 +242,10 @@ def shutdown(self, *args, **kwargs): return None def mock_request(method, url, **kwargs): - path = urlsplit(url).path + 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": @@ -242,8 +264,6 @@ def mock_request(method, url, **kwargs): token_id="test-id", collector_name="openhound-faker", ) - - def test_client_update_sends_metadata(mock_service, mock_bloodhound_api, monkeypatch): monkeypatch.setattr( bloodhound_enterprise.socket, "gethostname", lambda: "test-host" @@ -298,6 +318,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 ): @@ -545,6 +584,104 @@ 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"} + ] + assert mock_bloodhound_api.app.state.jobs_available_requests == 0 + + +def test_collector_job_queue_response_parses_contract_fields(): + payload = load_json("collector_jobs_available_with_jobs.json") + response = CollectorJobsAvailable.model_validate(payload) + 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 job.job_key == "openhound-faker" + assert job.params_version == "v2" + assert job.params == { + "domains": ["example.com", "example.org"], + "include_deleted": False, + "nested": {"depth": 2}, + } + assert job.scope_client_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" + + +@pytest.mark.parametrize("field", ["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", ["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) + + +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() 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 0000000..4b71256 --- /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 0000000..7e59a5c --- /dev/null +++ b/tests/test_data/api/jobs/collector_jobs_available_with_job.json @@ -0,0 +1,36 @@ +{ + "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", + "params_version": "v2", + "params": { + "domains": ["example.com", "example.org"], + "include_deleted": false, + "nested": {"depth": 2} + }, + "scope_client_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 0000000..7b68a33 --- /dev/null +++ b/tests/test_data/api/jobs/collector_jobs_available_with_jobs.json @@ -0,0 +1,59 @@ +{ + "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", + "params_version": "v2", + "params": { + "domains": ["example.com", "example.org"], + "include_deleted": false, + "nested": {"depth": 2} + }, + "scope_client_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": 42, + "job_profile_id": 84, + "job_type_id": 8, + "job_key": "openhound-faker", + "params_version": "v3", + "params": {"mode": "full"}, + "scope_client_id": "33333333-3333-3333-3333-333333333333", + "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" + } + ] + } +} From 214a5d430bd5f81a23d63d086b9503a5d75bc06f Mon Sep 17 00:00:00 2001 From: Adrian Hannah Date: Tue, 6 Oct 2026 14:09:47 -0400 Subject: [PATCH 11/34] fix: corrected 'end job' endpoint --- src/openhound/core/clients/bloodhound_enterprise.py | 2 +- tests/test_bhe_job_scheduling.py | 6 ++++-- 2 files changed, 5 insertions(+), 3 deletions(-) diff --git a/src/openhound/core/clients/bloodhound_enterprise.py b/src/openhound/core/clients/bloodhound_enterprise.py index 8d86eda..274bffe 100644 --- a/src/openhound/core/clients/bloodhound_enterprise.py +++ b/src/openhound/core/clients/bloodhound_enterprise.py @@ -121,7 +121,7 @@ def end_managed_job( metadata: dict[str, Any] | None = None, ) -> None: """End a claimed managed collector job.""" - path = f"/api/v2/collector-jobs/{job_id}/end" + path = f"/api/v2/collector-job-queue/{job_id}/end" payload: dict[str, Any] = {"outcome": outcome.value} if failure_message is not None: payload["failure_message"] = failure_message diff --git a/tests/test_bhe_job_scheduling.py b/tests/test_bhe_job_scheduling.py index ba0cd5c..cdab470 100644 --- a/tests/test_bhe_job_scheduling.py +++ b/tests/test_bhe_job_scheduling.py @@ -112,7 +112,7 @@ async def end_job(body: dict): app.state.end_payload = body return load_json("job_end.json") - @app.post("/api/v2/collector-jobs/{job_id}/end") + @app.post("/api/v2/collector-job-queue/{job_id}/end") async def end_managed_job(job_id: str, body: dict): app.state.collector_job_end_requests.append( {"job_id": job_id, "body": body} @@ -264,6 +264,8 @@ def mock_request(method, url, **kwargs): token_id="test-id", collector_name="openhound-faker", ) + + def test_client_update_sends_metadata(mock_service, mock_bloodhound_api, monkeypatch): monkeypatch.setattr( bloodhound_enterprise.socket, "gethostname", lambda: "test-host" @@ -482,7 +484,7 @@ def flaky_request(*args, **kwargs): assert any( record.levelno == logging.DEBUG and getattr(record, "endpoint", None) - == "/api/v2/collector-jobs/collector-job-123/end" + == "/api/v2/collector-job-queue/collector-job-123/end" for record in caplog.records ) From a801f169a29314d56d83e51a16b6ec1f79b85793 Mon Sep 17 00:00:00 2001 From: Justin LaRose Date: Wed, 7 Oct 2026 13:25:58 -0400 Subject: [PATCH 12/34] BED-9510 - Adding BHE Auth path (#88) * feat: Adding authentication path for initial BHE calls and subsequent managed BHE calls. BED-9510 * Added a todo comment while we iron out the exact shape of this during dev work. BED-9510 * BED-9510 - CodeRabbit suggestion fixes after review. * BED-9510 code rabbit and PR feedback. --- .../.dlt-example/config.toml | 7 + src/openhound/config.py | 42 ++++++ src/openhound/core/clients/bhe_credentials.py | 89 ++++++++++++ .../core/clients/bloodhound_enterprise.py | 44 ++++++ src/openhound/scheduler/service.py | 32 +++-- src/scheduler.py | 24 +++- tests/test_bhe_credentials.py | 44 ++++++ tests/test_bhe_job_scheduling.py | 136 ++++++++++++++++++ tests/test_config.py | 54 ++++++- tests/test_scheduler_start.py | 104 ++++++++++++++ 10 files changed, 558 insertions(+), 18 deletions(-) create mode 100644 src/openhound/core/clients/bhe_credentials.py create mode 100644 tests/test_bhe_credentials.py create mode 100644 tests/test_scheduler_start.py diff --git a/example-configurations/bloodhound-enterprise/.dlt-example/config.toml b/example-configurations/bloodhound-enterprise/.dlt-example/config.toml index c486a50..370cf51 100644 --- a/example-configurations/bloodhound-enterprise/.dlt-example/config.toml +++ b/example-configurations/bloodhound-enterprise/.dlt-example/config.toml @@ -2,6 +2,13 @@ [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" diff --git a/src/openhound/config.py b/src/openhound/config.py index f40d1e1..598c310 100644 --- a/src/openhound/config.py +++ b/src/openhound/config.py @@ -12,8 +12,18 @@ class ModeDefaults: 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: @@ -32,3 +42,35 @@ 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/bhe_credentials.py b/src/openhound/core/clients/bhe_credentials.py new file mode 100644 index 0000000..37b4d3b --- /dev/null +++ b/src/openhound/core/clients/bhe_credentials.py @@ -0,0 +1,89 @@ +"""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, 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. + """ + + 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.refresh() + except Exception: + # AWS errors can be transient, and an absent or not-yet-populated + # secret is expected while a managed deployment is provisioning. + logger.warning( + "BHE credentials are unavailable; retrying in %s seconds.", + interval, + ) + time.sleep(interval) diff --git a/src/openhound/core/clients/bloodhound_enterprise.py b/src/openhound/core/clients/bloodhound_enterprise.py index 274bffe..48d0874 100644 --- a/src/openhound/core/clients/bloodhound_enterprise.py +++ b/src/openhound/core/clients/bloodhound_enterprise.py @@ -14,6 +14,7 @@ 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, @@ -56,6 +57,49 @@ class ManagedJobOutcome(str, Enum): 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 | 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.""" + 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 + + 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" diff --git a/src/openhound/scheduler/service.py b/src/openhound/scheduler/service.py index a8831f0..99b0120 100644 --- a/src/openhound/scheduler/service.py +++ b/src/openhound/scheduler/service.py @@ -1,6 +1,7 @@ import logging import signal import time +from collections.abc import Callable from concurrent.futures import Future, ProcessPoolExecutor from concurrent.futures.process import BrokenProcessPool from dataclasses import dataclass @@ -8,6 +9,7 @@ import openhound import openhound.core.logging as openhound_logging +from openhound.core.clients.bhe_credentials import BHECredentials from openhound.core.clients.bloodhound_enterprise import BloodHoundEnterprise, JobStatus from openhound.core.clients.models.jobs import ( CollectorJob, @@ -102,15 +104,20 @@ def __init__( token_id: str, collector_name: str, log_base_path: Path | None = None, + interval: int = POLL_INTERVAL, + credential_refresh: Callable[[], BHECredentials] | 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 + self.interval = interval self.log_base_path = ( log_base_path or openhound_logging.logger_override.base_path ) @@ -258,11 +265,6 @@ def _handle_completed_job(self, future: Future[Result]) -> None: """ 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", - ) except ExtensionNotFoundError: logger.error( @@ -290,9 +292,15 @@ def _handle_completed_job(self, future: Future[Result]) -> None: f"Unexpected error while running '{self.collector_name}' collector", ) - finally: - self.future = None - self.job_running = None + else: + 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.future = None + self.job_running = None def _poll(self) -> None: """Checks if jobs are completed and if a job should be run.""" @@ -301,9 +309,7 @@ def _poll(self) -> None: 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 + 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: diff --git a/src/scheduler.py b/src/scheduler.py index d18fa16..5fc388c 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_bhe_credentials.py b/tests/test_bhe_credentials.py new file mode 100644 index 0000000..c3d27c3 --- /dev/null +++ b/tests/test_bhe_credentials.py @@ -0,0 +1,44 @@ +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() diff --git a/tests/test_bhe_job_scheduling.py b/tests/test_bhe_job_scheduling.py index cdab470..8f8a0c4 100644 --- a/tests/test_bhe_job_scheduling.py +++ b/tests/test_bhe_job_scheduling.py @@ -16,6 +16,7 @@ from fastapi.testclient import TestClient from openhound.core.clients import bloodhound, bloodhound_enterprise +from openhound.core.clients.bhe_credentials import BHECredentials from openhound.core.clients.bloodhound import BloodHoundHTTPError from openhound.core.clients.bloodhound_enterprise import ( JobStatus, @@ -304,6 +305,66 @@ 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") + + assert refreshes == [True] + assert len(requests) == 2 + assert requests[0]["headers"]["Authorization"] == "bhesignature old-id" + assert requests[1]["headers"]["Authorization"] == "bhesignature new-id" + 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 = {} @@ -1152,3 +1213,78 @@ 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")] diff --git a/tests/test_config.py b/tests/test_config.py index 87a673f..ac13389 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -2,7 +2,14 @@ import pytest from dlt.common.configuration.exceptions import ConfigValueCannotBeCoercedException -from openhound.config import ModeDefaults, get_mode_defaults, is_managed +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): @@ -43,3 +50,48 @@ def test_get_mode_defaults_selects_mode_specific_log_formats(monkeypatch): 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_scheduler_start.py b/tests/test_scheduler_start.py new file mode 100644 index 0000000..0c0938a --- /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", + ] From 6f59b04bc732d0674aefd2067b046eac97924cac Mon Sep 17 00:00:00 2001 From: Stran Dutton <65475509+StranDutton@users.noreply.github.com> Date: Thu, 8 Oct 2026 16:18:43 -0500 Subject: [PATCH 13/34] Claim a Job (managed OH) - BED-9515 (#94) * Claim a job (managed collector) = BED-9515 * refactor: clarify managed job handling and retry/attempt limits * update description for clarity * add claim-recovery logic, clean up scheduler constructor settings * ensure support bundles still work in both modes, tests --- .../core/clients/aws_secrets_manager.py | 63 +- .../core/clients/bloodhound_enterprise.py | 101 +- src/openhound/scheduler/service.py | 279 +++++- tests/test_aws_secrets_manager.py | 257 ++++- tests/test_bhe_job_scheduling.py | 880 +++++++++++++++++- tests/test_log_handlers.py | 142 +-- 6 files changed, 1575 insertions(+), 147 deletions(-) diff --git a/src/openhound/core/clients/aws_secrets_manager.py b/src/openhound/core/clients/aws_secrets_manager.py index 567369c..f600098 100644 --- a/src/openhound/core/clients/aws_secrets_manager.py +++ b/src/openhound/core/clients/aws_secrets_manager.py @@ -4,6 +4,7 @@ from typing import Any, Protocol import boto3 +from botocore.config import Config from botocore.exceptions import ( BotoCoreError, ClientError, @@ -21,6 +22,38 @@ MAX_BATCH_SIZE = 20 _PROVIDER = "aws.secretsmanager" +# 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]: ... @@ -107,9 +140,7 @@ def _get_secrets_in_batches( except (ClientError, BotoCoreError) as error: classified = _classify_aws_error(error) chunk_values: dict[str, SecretValue] = {} - chunk_failures = { - secret_id: classified for secret_id in chunk - } + chunk_failures = {secret_id: classified for secret_id in chunk} else: chunk_values, chunk_failures = _parse_batch_response(response, chunk) @@ -138,10 +169,14 @@ def _get_secrets_in_batches( } def _get_client(self) -> SecretsManagerClient: + _configure_safe_sdk_logging() if self._client is not None: return self._client try: - return boto3.client("secretsmanager") + return boto3.client( + "secretsmanager", + config=DEFAULT_AWS_CLIENT_CONFIG, + ) except BotoCoreError as error: raise _classify_aws_error(error) from None @@ -172,9 +207,17 @@ def _chunks(values: list[str], size: int) -> Iterable[list[str]]: def _parse_secret_value(response: Mapping[str, Any]) -> SecretValue: secret_string = response.get("SecretString") if not isinstance(secret_string, str): - raise SecretRequestError( - "AWS Secrets Manager response did not contain a SecretString payload" - ) + # 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: @@ -271,7 +314,11 @@ def _log_debug( } if batch_number is not None: extra["batch_number"] = batch_number - message = "Secret retrieval attempt" if outcome == "attempt" else "Secret retrieval result" + message = ( + "Secret retrieval attempt" + if outcome == "attempt" + else "Secret retrieval result" + ) logger.debug(message, extra=extra) diff --git a/src/openhound/core/clients/bloodhound_enterprise.py b/src/openhound/core/clients/bloodhound_enterprise.py index 48d0874..a9c2c90 100644 --- a/src/openhound/core/clients/bloodhound_enterprise.py +++ b/src/openhound/core/clients/bloodhound_enterprise.py @@ -18,6 +18,7 @@ from openhound.core.clients.bloodhound import BloodHound, BloodHoundHTTPError from openhound.core.clients.models.jobs import ( ArtifactUploadSession, + CollectorJob, CollectorJobsAvailable, JobsAvailable, JobsCurrent, @@ -41,7 +42,7 @@ class ManagedJobOutcome(str, Enum): FAILED = "failed" -MANAGED_JOB_END_MAX_RETRIES = 3 +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 @@ -52,10 +53,23 @@ class ManagedJobOutcome(str, Enum): 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, @@ -136,13 +150,63 @@ def available_collector_jobs(self, job_key: str) -> CollectorJobsAvailable: MANAGED_COLLECTOR_JOB_AVAILABLE_READ_TIMEOUT_SECONDS, ), ) - except (BloodHoundHTTPError, requests.RequestException): - logger.exception( + except (BloodHoundHTTPError, requests.RequestException) as error: + logger.error( "Managed collector job queue request failed.", extra={"endpoint": path, "job_key": job_key}, ) - raise - return CollectorJobsAvailable.model_validate(response.json()) + 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/collector-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" @@ -172,20 +236,18 @@ def end_managed_job( if metadata is not None: payload["metadata"] = metadata body = json.dumps(payload).encode() - max_attempts = MANAGED_JOB_END_MAX_RETRIES + 1 - - for attempt in range(1, max_attempts + 1): + for attempt in range(1, MANAGED_JOB_END_MAX_ATTEMPTS + 1): log_context = { "job_id": job_id, "outcome": outcome.value, "attempt": attempt, - "max_attempts": max_attempts, + "max_attempts": MANAGED_JOB_END_MAX_ATTEMPTS, } logger.info( "Attempting to end managed collector job %s (%s/%s).", job_id, attempt, - max_attempts, + MANAGED_JOB_END_MAX_ATTEMPTS, extra=log_context, ) logger.debug( @@ -208,7 +270,13 @@ def end_managed_job( MANAGED_JOB_END_READ_TIMEOUT_SECONDS, ), ) - except Exception as error: + 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), @@ -220,12 +288,15 @@ def end_managed_job( ) logger.debug( "Managed collector job end request failed.", - extra={**log_context, "retryable": retryable}, - exc_info=True, + extra={ + **log_context, + "retryable": retryable, + "status_code": getattr(error, "code", None), + }, ) - if not retryable or attempt == max_attempts: - raise + 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.", diff --git a/src/openhound/scheduler/service.py b/src/openhound/scheduler/service.py index 99b0120..c3303c7 100644 --- a/src/openhound/scheduler/service.py +++ b/src/openhound/scheduler/service.py @@ -1,16 +1,26 @@ import logging +import math import signal import time from collections.abc import Callable from concurrent.futures import Future, ProcessPoolExecutor from concurrent.futures.process import BrokenProcessPool -from dataclasses import dataclass +from dataclasses import dataclass, field +from datetime import UTC, datetime from pathlib import Path +from typing import Any import openhound import openhound.core.logging as openhound_logging +from openhound.config import is_managed +from openhound.core.clients.aws_secrets_manager import AWSSecretsManager from openhound.core.clients.bhe_credentials import BHECredentials -from openhound.core.clients.bloodhound_enterprise import BloodHoundEnterprise, JobStatus +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, @@ -30,7 +40,48 @@ 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.""" + + +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 memory; never serialize or log them. + """ + + job: CollectorJob = field(repr=False) + credentials: dict[str, Any] = field(repr=False) + + +def _usable_credential_value(value: Any) -> bool: + """Validate generic JSON values without inventing a collector field schema.""" + if isinstance(value, str): + return bool(value.strip()) + if isinstance(value, dict): + return bool(value) and all( + isinstance(key, str) + and bool(key.strip()) + and _usable_credential_value(item) + for key, item in value.items() + ) + if isinstance(value, list): + return bool(value) and all(_usable_credential_value(item) for item in value) + if isinstance(value, (bool, int)): + return True + if isinstance(value, float): + return math.isfinite(value) + return False @dataclass @@ -58,7 +109,9 @@ def _subprocess_collect(collector_name: str, job_id: int) -> Result: available_collectors = CollectorManager.from_entrypoint() for collector in available_collectors.collectors: - if collector.name == collector_name: # pyright: ignore[reportAttributeAccessIssue] + if ( + collector.name == collector_name + ): # pyright: ignore[reportAttributeAccessIssue] log_fields = { "collector_extension": collector.name, "collector_extension_version": ( @@ -106,22 +159,34 @@ def __init__( 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, credential_refresh=credential_refresh, ) - # Interval how often to check for a job + + # 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 ) + # Managed collection settings and recovery state + self.managed = is_managed() if managed is None else managed + self.secrets_manager = secrets_manager or AWSSecretsManager() + self.managed_job_runner = managed_job_runner + self._pending_managed_claim: str | None = None + self._pending_managed_failure: str | None = None + # Stores the ID of currently running BHE job self.job_running: int | None = None @@ -182,13 +247,128 @@ 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) - if available_jobs.data.jobs: - logger.info( - "New managed queue job available: %s", available_jobs.data.jobs[0].id - ) - return available_jobs.data.jobs[0] + 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="Credential validation failed", + ) + self._pending_managed_failure = None + + def _load_managed_credentials(self, job: CollectorJob) -> dict[str, Any]: + """Retrieve and validate credentials, exposing only a safe failure message.""" + if job.secret_key_id is None or not job.secret_key_id.strip(): + raise CredentialValidationError("Credential validation failed") + + logger.debug( + "Retrieving managed collector credentials.", + extra={ + "job_id": str(job.id), + "provider": "aws.secretsmanager", + "secret_count": 1, + }, + ) + try: + credentials = self.secrets_manager.get_secret(job.secret_key_id) + if isinstance(credentials, dict) and _usable_credential_value(credentials): + return credentials + except Exception: # noqa: BLE001 - redact secret reader and validation errors + raise CredentialValidationError("Credential validation failed") from None + raise CredentialValidationError("Credential validation failed") + + 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: + logger.error("Credential validation failed", extra=context) + self._pending_managed_claim = None + 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) + self._pending_managed_claim = None + logger.info("Credential retrieval valid", extra=context) + return PreparedManagedJob(job=job, credentials=credentials) + + 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 _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 + 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) + if prepared is not None: + runner(prepared) + 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.") @@ -200,6 +380,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 @@ -304,28 +500,30 @@ def _handle_completed_job(self, future: Future[Result]) -> None: 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("Error reporting completed job; retaining result for retry.") - - # Management operations have priority over new collections while idle. - if self.job_running is None: + if not self.managed: try: - operation = self.check_management() + if self.future is not None and self.future.done(): + self._handle_completed_job(self.future) except Exception: - logger.exception("Error checking management operations.") - operation = None + logger.exception( + "Error reporting completed job; retaining result for retry." + ) - if operation is not None: - try: - self._send_support_bundle(operation) - except Exception: - logger.exception("Error executing management operation.") - return + # Support bundles precede new collections in either mode. An unresolved + # managed job must be reconciled before accepting management work. + idle = ( + self.job_running is None + and self._pending_managed_claim is None + and self._pending_managed_failure is None + ) + if idle and self._poll_management(): + return + if self.managed: + self._poll_managed_jobs() + return + + if self.job_running is None: try: available_job = self.check_jobs() if available_job: @@ -334,7 +532,7 @@ def _poll(self) -> None: logger.exception("Error checking for or starting jobs.") else: try: - self.client.jobs_current + _ = self.client.jobs_current except Exception: logger.exception("Error checking in-progress job.") @@ -348,12 +546,27 @@ def start(self) -> None: try: self.client.update_client_metadata() - except Exception as err: - logger.exception(f"Unable to update client metadata: {err}") + except Exception: + if self.managed: + logger.error("Unable to update client metadata.") + else: + 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; continuing support-bundle polling." + ) + except ManagedJobProcessingError: + logger.error( + "Managed scheduler poll failed; retrying next poll." + ) + else: + self._poll() time.sleep(self.interval) finally: self._shutdown() diff --git a/tests/test_aws_secrets_manager.py b/tests/test_aws_secrets_manager.py index 98b01e1..a6a8f7f 100644 --- a/tests/test_aws_secrets_manager.py +++ b/tests/test_aws_secrets_manager.py @@ -1,8 +1,13 @@ +import base64 +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, @@ -31,6 +36,38 @@ PROVIDER_MESSAGE = "provider message that must not be exposed" +@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, @@ -51,9 +88,7 @@ def get_secret_value(self, *, SecretId: str) -> Mapping[str, Any]: raise self.get_error return self.get_response or {} - def batch_get_secret_value( - self, *, SecretIdList: list[str] - ) -> Mapping[str, Any]: + 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": []} @@ -87,9 +122,7 @@ 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: +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) @@ -195,9 +228,10 @@ def test_get_secrets_continues_and_aggregates_typed_failures_safely( ) forbidden = set(secret_ids) | {SECRET_VALUE, PROVIDER_MESSAGE} - with caplog.at_level(logging.DEBUG, logger=LOGGER_NAME), pytest.raises( - SecretBatchError - ) as raised: + 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:]] @@ -220,7 +254,9 @@ def test_get_secrets_continues_and_aggregates_typed_failures_safely( if record.levelno == logging.ERROR ] assert len(error_records) == 2 - assert all(record.failure_category == "SecretBatchError" for record in error_records) + assert all( + record.failure_category == "SecretBatchError" for record in error_records + ) assert_safe_logs(caplog, forbidden) @@ -331,3 +367,204 @@ def test_successful_retrieval_logs_safe_attempt_and_result_records( 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} {SECRET_VALUE}", + } + 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_ID, + SECRET_VALUE, + PROVIDER_MESSAGE, + 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 + ) diff --git a/tests/test_bhe_job_scheduling.py b/tests/test_bhe_job_scheduling.py index 8f8a0c4..97e97eb 100644 --- a/tests/test_bhe_job_scheduling.py +++ b/tests/test_bhe_job_scheduling.py @@ -3,8 +3,10 @@ import hashlib import json import logging +import traceback from concurrent.futures import Future from concurrent.futures.process import BrokenProcessPool +from datetime import UTC, datetime, timedelta from pathlib import Path from types import SimpleNamespace from urllib.parse import urlsplit @@ -12,11 +14,16 @@ import pytest import requests +from botocore.exceptions import ClientError from fastapi import FastAPI, Request, Response from fastapi.testclient import TestClient from openhound.core.clients import bloodhound, bloodhound_enterprise -from openhound.core.clients.bhe_credentials import BHECredentials +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, @@ -33,6 +40,8 @@ from openhound.scheduler import service as scheduler_service from openhound.scheduler.service import ( ExtensionNotFoundError, + ManagedJobProcessingError, + ManagedRuntimeUnavailableError, Result, Service, _subprocess_collect, @@ -60,6 +69,16 @@ def mock_bloodhound_api(): 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.start_payload = None app.state.jobs_available_requests = 0 app.state.jobs_current_requests = 0 @@ -71,6 +90,7 @@ def mock_bloodhound_api(): 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 @@ -99,6 +119,18 @@ async def collector_job_queue(request: Request): app.state.collector_job_queue_requests.append(dict(request.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") @@ -107,6 +139,26 @@ async def start_job(body: dict): app.state.start_payload = body return load_json("job_start.json") + @app.post("/api/v2/collector-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) + 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 @@ -115,10 +167,16 @@ async def end_job(body: dict): @app.post("/api/v2/collector-job-queue/{job_id}/end") async def end_managed_job(job_id: str, body: dict): - app.state.collector_job_end_requests.append( - {"job_id": job_id, "body": body} + 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 ) - return Response(status_code=200) @app.post("/api/v2/ingest") async def ingest(request: Request): @@ -136,6 +194,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") @@ -264,6 +323,7 @@ def mock_request(method, url, **kwargs): token_key="test-key", token_id="test-id", collector_name="openhound-faker", + managed=False, ) @@ -331,12 +391,13 @@ def refresh(): credential_refresh=refresh, ) - client.request(method="GET", path="/test") + 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" @@ -570,10 +631,10 @@ def unavailable(*args, **kwargs): ) assert error.value.code == 503 - assert attempts == bloodhound_enterprise.MANAGED_JOB_END_MAX_RETRIES + 1 - assert delays == [ - bloodhound_enterprise.MANAGED_JOB_END_RETRY_DELAY_SECONDS - ] * bloodhound_enterprise.MANAGED_JOB_END_MAX_RETRIES + 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( @@ -748,7 +809,11 @@ def test_queue_http_error_is_logged_and_re_raised( 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() + secret_requests = [] monkeypatch.setattr(mock_service.executor, "submit", lambda *args: submitted) + monkeypatch.setattr( + mock_service.secrets_manager, "get_secret", secret_requests.append + ) mock_service._poll() @@ -756,6 +821,10 @@ def test_poll_starts_new_job(mock_service, mock_bloodhound_api, monkeypatch): assert mock_service.future is submitted 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): @@ -958,9 +1027,11 @@ def test_check_management_ignores_non_queued_operations( assert mock_service.check_management() is None +@pytest.mark.parametrize("managed", [False, True]) def test_poll_prioritizes_management_over_a_new_job( - mock_service, mock_bloodhound_api, monkeypatch + mock_service, mock_bloodhound_api, monkeypatch, managed ): + mock_service.managed = managed mock_bloodhound_api.app.state.management_operations = [_support_bundle_operation()] sent = [] monkeypatch.setattr(mock_service, "_send_support_bundle", sent.append) @@ -969,6 +1040,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( @@ -982,21 +1054,33 @@ 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 + assert any(record.exc_info is not None for record in caplog.records) +@pytest.mark.parametrize("managed", [False, True]) def test_poll_does_not_start_a_job_when_management_work_fails( - mock_service, mock_bloodhound_api, monkeypatch + mock_service, mock_bloodhound_api, monkeypatch, managed ): + mock_service.managed = managed mock_bloodhound_api.app.state.management_operations = [_support_bundle_operation()] def fail(operation): @@ -1007,6 +1091,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( @@ -1097,10 +1182,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") @@ -1115,7 +1202,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) @@ -1125,6 +1212,11 @@ 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( @@ -1152,7 +1244,7 @@ def request(): 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") @@ -1181,6 +1273,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): @@ -1205,7 +1299,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: @@ -1262,9 +1358,7 @@ def test_managed_refresh_rejects_invalid_id_and_recovers(mock_service, monkeypat ) secret = {"token_id": "invalid-id", "token_key": "new-key"} - provider = AWSBHECredentials( - "bhe", SimpleNamespace(get_secret=lambda _: secret) - ) + 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) @@ -1288,3 +1382,741 @@ def request(self, **kwargs): 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_parameter": True}, + 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": '{"credential":"private-credential-value"}'} + + 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 + + +def test_claim_request_is_bodyless_and_parses_data_job( + managed_service, mock_bloodhound_api, monkeypatch +): + 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 captured == [ + { + "method": "POST", + "path": f"/api/v2/collector-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_parameter": True} + assert prepared.job.status == "claimed" + assert prepared.credentials == {"credential": "private-credential-value"} + 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) + + +@pytest.mark.parametrize( + "document", + [ + "plain-secret", + '{"malformed":', + "{}", + '{"value":" "}', + '{"value":null}', + '{"value":{}}', + '{"value":[]}', + '{"value":{"nested":""}}', + '{"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 validation failed", + }, + } + ] + 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 validation failed" + ) + + +@pytest.mark.parametrize( + "error", + [ + ClientError( + { + "Error": { + "Code": "AccessDeniedException", + "Message": "private-credential-value", + } + }, + "GetSecretValue", + ), + RuntimeError("private-credential-value"), + ], + ids=["aws-permission-error", "unexpected-reader-error"], +) +def test_secret_retrieval_failure_is_redacted_and_reported_as_failed( + managed_service, mock_bloodhound_api, caplog, error +): + 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": "Credential validation failed", + } + assert managed_service.handoffs == [] + 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 + state.collector_job_claim_response["secret_key_id"] = None + 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 ( + 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 + with pytest.raises( + ManagedRuntimeUnavailableError, match="Managed collection runtime is required" + ): + service._poll() + + +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 = {"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() + ] == ["Managed scheduler poll failed; retrying next poll."] + 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.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/collector-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_parameter": True} + + 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_keeps_serving_support_bundles_without_collection_runtime( + managed_service, mock_bloodhound_api, monkeypatch, tmp_path, caplog +): + service = managed_service + state = mock_bloodhound_api.app.state + service.managed_job_runner = None + service.log_base_path = tmp_path + (tmp_path / "openhound.log").write_text("support log") + polls = [] + + def next_poll(seconds): + assert seconds == service.interval + polls.append(seconds) + if len(polls) == 1: + # Support work arrives after a poll found no collection runtime. + state.management_operations = [_support_bundle_operation()] + else: + 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) == state.management_available_requests == 2 + assert state.operation_started is True + assert state.artifact_completed is True + assert state.operation_end_payload["status"] == "succeeded" + 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; continuing support-bundle polling." + in caplog.text + ) + + +@pytest.mark.parametrize("pending", ["claim", "failure"]) +def test_managed_recovery_precedes_management_then_allows_support_work( + 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 == 1 + assert len(sent) == 1 + + +def test_support_bundle_tracebacks_keep_aws_auth_errors_redacted( + 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 "AWS Secrets Manager denied access to the secret" 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 forbidden in (secret_reference, provider_message): + assert forbidden not in caplog.text + assert forbidden not in formatted_error + assert all(forbidden not in repr(record.__dict__) for record in caplog.records) diff --git a/tests/test_log_handlers.py b/tests/test_log_handlers.py index cdb4c30..f89dae9 100644 --- a/tests/test_log_handlers.py +++ b/tests/test_log_handlers.py @@ -2,18 +2,66 @@ import logging from pathlib import Path +import dlt import pytest from openhound.core.logging import ( CustomLogger, + LogMode, OpenHoundJSONFormatter, OpenHoundTextFormatter, RotatingFileHandler, - logger_override, ) -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 +77,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 +96,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,13 +116,9 @@ 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") @@ -126,10 +170,11 @@ def test_file_formatter_selection(): ) -def test_managed_mode_writes_structured_file_logs(tmp_path, monkeypatch): +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.delenv("RUNTIME__LOG_FORMAT", raising=False) 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. @@ -184,10 +229,11 @@ def test_explicit_log_format_overrides_managed_default(tmp_path, monkeypatch): custom_logger.dlt_logger.handlers.clear() -def test_mode_change_refreshes_cached_file_handler(tmp_path, monkeypatch): +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)) - monkeypatch.delenv("RUNTIME__LOG_FORMAT", raising=False) 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") @@ -278,51 +324,33 @@ 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): +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.""" - # 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)) - - 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'" - ) - finally: - for handler in custom_logger._file_handlers.values(): - 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() + 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): From 1897469273e4cc10aabbb2b9d5aa93d882fd3614 Mon Sep 17 00:00:00 2001 From: Adrian Hannah Date: Fri, 9 Oct 2026 10:29:34 -0400 Subject: [PATCH 14/34] feat: add managed collector start and heartbeat requests --- .../core/clients/bloodhound_enterprise.py | 24 +++++++++ tests/test_bhe_job_scheduling.py | 50 +++++++++++++++++++ 2 files changed, 74 insertions(+) diff --git a/src/openhound/core/clients/bloodhound_enterprise.py b/src/openhound/core/clients/bloodhound_enterprise.py index a9c2c90..9b82f54 100644 --- a/src/openhound/core/clients/bloodhound_enterprise.py +++ b/src/openhound/core/clients/bloodhound_enterprise.py @@ -46,6 +46,8 @@ class ManagedJobOutcome(str, Enum): 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 @@ -221,6 +223,28 @@ 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/collector-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/collector-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, diff --git a/tests/test_bhe_job_scheduling.py b/tests/test_bhe_job_scheduling.py index 97e97eb..ef99be8 100644 --- a/tests/test_bhe_job_scheduling.py +++ b/tests/test_bhe_job_scheduling.py @@ -79,6 +79,8 @@ def mock_bloodhound_api(): "collector_jobs_available_with_job.json" )["data"]["jobs"][0] app.state.managed_events = [] + app.state.collector_job_start_requests = [] + app.state.collector_job_heartbeat_requests = [] app.state.start_payload = None app.state.jobs_available_requests = 0 app.state.jobs_current_requests = 0 @@ -178,6 +180,18 @@ async def end_managed_job(job_id: str, body: dict): "raw-provider-secret" if status != 200 else "", status_code=status ) + @app.post("/api/v2/collector-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())) + return Response(status_code=200) + + @app.post("/api/v2/collector-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() @@ -540,6 +554,42 @@ def test_end_managed_job_sends_managed_outcome( ] +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 = {} From d602658a32724de543c0c610483777522fa54fc5 Mon Sep 17 00:00:00 2001 From: Adrian Hannah Date: Fri, 9 Oct 2026 10:29:35 -0400 Subject: [PATCH 15/34] test: isolate integration collection output --- tests/test_bhe_job_scheduling.py | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/tests/test_bhe_job_scheduling.py b/tests/test_bhe_job_scheduling.py index ef99be8..172a16e 100644 --- a/tests/test_bhe_job_scheduling.py +++ b/tests/test_bhe_job_scheduling.py @@ -1035,8 +1035,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" ) From b85b594d1e76c379423b0ae721ee86ec5d254a91 Mon Sep 17 00:00:00 2001 From: Adrian Hannah Date: Fri, 9 Oct 2026 10:29:36 -0400 Subject: [PATCH 16/34] feat: run claimed managed jobs with five-minute heartbeats --- src/openhound/scheduler/service.py | 243 ++++++++++++++++++++------- tests/test_bhe_job_scheduling.py | 261 ++++++++++++++++++++++++++++- 2 files changed, 444 insertions(+), 60 deletions(-) diff --git a/src/openhound/scheduler/service.py b/src/openhound/scheduler/service.py index c3303c7..6f53584 100644 --- a/src/openhound/scheduler/service.py +++ b/src/openhound/scheduler/service.py @@ -1,5 +1,7 @@ +import json import logging import math +import os import signal import time from collections.abc import Callable @@ -35,12 +37,12 @@ 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 class ExtensionNotFoundError(Exception): """Raised when the configured collector extension cannot be found.""" - class ManagedRuntimeUnavailableError(RuntimeError): """Managed scheduling requires a configured collection runtime.""" @@ -83,14 +85,41 @@ def _usable_credential_value(value: Any) -> bool: return math.isfinite(value) return False - @dataclass class Result: results: dict - job_id: int + job_id: int | str + + +def _set_managed_source_environment( + collector, params: dict[str, Any], secrets: dict[str, Any] +) -> None: + """Provide one managed job's validated values to its DLT source.""" + source = collector.dlt_source + if source is None: + raise ValueError(f"Collector '{collector.name}' has no registered DLT source") + + section = source.section or source.__module__.rsplit(".", 1)[-1] + source_name = source.name or source.__name__ + prefix = f"SOURCES__{section.upper()}__{source_name.upper()}" + for name, value in params.items(): + if name.lower() == "credentials": + raise ValueError("Managed job parameters cannot contain credentials") + os.environ[f"{prefix}__{name.upper()}"] = ( + value if isinstance(value, str) else json.dumps(value) + ) + for name, value in secrets.items(): + os.environ[f"{prefix}__CREDENTIALS__{name.upper()}"] = ( + value if isinstance(value, str) else json.dumps(value) + ) -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] | None = None, +) -> Result: """A subprocess which runs the DLT pipeline for the specified collector. Loads the collector by name from Python entrypoints. @@ -109,9 +138,9 @@ def _subprocess_collect(collector_name: str, job_id: int) -> Result: available_collectors = CollectorManager.from_entrypoint() for collector in available_collectors.collectors: - if ( - collector.name == collector_name - ): # pyright: ignore[reportAttributeAccessIssue] + if collector.name == collector_name: # pyright: ignore[reportAttributeAccessIssue] + if params is not None or secrets is not None: + _set_managed_source_environment(collector, params or {}, secrets or {}) log_fields = { "collector_extension": collector.name, "collector_extension_version": ( @@ -187,8 +216,11 @@ def __init__( self._pending_managed_claim: str | None = None self._pending_managed_failure: str | None = None - # Stores the ID of currently running BHE job - self.job_running: int | None = None + # 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 @@ -203,11 +235,28 @@ 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.""" + """Finish the active job and keep checking in until its worker stops.""" logger.info("Collection service stopping.") + while self.future is not None and not self.future.done(): + self._poll_running_job() + time.sleep(self.interval) + if 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() self.executor.shutdown(wait=True, cancel_futures=True) logger.info("Collection service stopped.") + def _clear_running_job(self) -> None: + self.future = None + self.job_running = None + self.managed_job_id = None + self.next_heartbeat_at = None + def _reset_executor(self) -> None: """Tear down and recreate the process pool after it has entered a broken state.""" try: @@ -427,76 +476,155 @@ 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, params: dict[str, Any], secrets: dict[str, Any] + ) -> None: + """Start a validated managed job; the scheduler polls it until completion.""" + job_id = str(job_id) + self._launch_job(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] | 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") + + logger.info("Starting job %s with collector '%s'.", job_id, self.collector_name) + if managed: + self.client.start_collector_job(str(job_id)) + else: + assert isinstance(job_id, int) + 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 = str(job_id) if managed else None + self.next_heartbeat_at = ( + time.monotonic() + HEARTBEAT_INTERVAL if managed else None + ) try: 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, + dict(secrets or {}) if managed else None, ) 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) - Args: - future (Future[Result]): Future containing the result of the collection subprocess, including the job ID. - """ + def _send_managed_heartbeat(self, job_id: str) -> None: try: - result = future.result() + self.client.heartbeat_collector_job(job_id) + except Exception as error: # noqa: BLE001 - failed heartbeats must not stop collection + 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 _handle_completed_job(self, future: Future[Result]) -> None: + """Report the completed worker's outcome through the active job API.""" + 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 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__}, ) - - else: - logger.info(f"Job {result.job_id} completed successfully, notifying BHE.") - self.client.end_job( - JobStatus.COMPLETE, - f"Collector '{self.collector_name}' completed successfully", + status = JobStatus.FAILED + message = ( + f"Unexpected error while running '{self.collector_name}' collector" ) + else: + status = JobStatus.COMPLETE + message = f"Collector '{self.collector_name}' completed successfully" + self._end_job(status, message) + self._clear_running_job() - self.future = None - self.job_running = None + 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) + 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.""" @@ -531,10 +659,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""" diff --git a/tests/test_bhe_job_scheduling.py b/tests/test_bhe_job_scheduling.py index 172a16e..230b393 100644 --- a/tests/test_bhe_job_scheduling.py +++ b/tests/test_bhe_job_scheduling.py @@ -860,15 +860,22 @@ 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() secret_requests = [] - monkeypatch.setattr(mock_service.executor, "submit", lambda *args: submitted) 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 mock_service.future is submitted + assert submitted_args[0][1:] == ("openhound-faker", 123, None, None) 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 == [] @@ -2178,3 +2185,255 @@ def expired_token(self, **kwargs): assert forbidden not in caplog.text assert forbidden not in formatted_error assert all(forbidden not in repr(record.__dict__) for record in caplog.records) +def test_managed_worker_sets_dlt_environment_before_pipeline(monkeypatch, caplog): + collector = SimpleNamespace( + name="faker", + dlt_source=SimpleNamespace(section="source", name="random"), + package_version="1.0.0", + metadata=None, + ) + worker_environment = {} + observed_environment = {} + monkeypatch.setattr( + scheduler_service.CollectorManager, + "from_entrypoint", + lambda: SimpleNamespace(collectors=[collector]), + ) + monkeypatch.setattr( + scheduler_service, "os", SimpleNamespace(environ=worker_environment) + ) + + def capture_pipeline(extension): + observed_environment.update(worker_environment) + 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", + {"node_count": 4, "filters": ["one", "two"]}, + {"token": "secret-value"}, + ) + + assert result.job_id == "collector-job-123" + assert observed_environment == { + "SOURCES__SOURCE__RANDOM__NODE_COUNT": "4", + "SOURCES__SOURCE__RANDOM__FILTERS": '["one", "two"]', + "SOURCES__SOURCE__RANDOM__CREDENTIALS__TOKEN": "secret-value", + } + assert "secret-value" not in caplog.text + + +def test_managed_environment_requires_a_registered_dlt_source(): + collector = SimpleNamespace(name="example", dlt_source=None) + + with pytest.raises(ValueError, match="has no registered DLT source"): + scheduler_service._set_managed_source_environment(collector, {}, {}) + + +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) + + +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"}, + ) + 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 + + +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 + + +def test_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("claim lost", 404) + + 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 + + +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", + }, + } + ] + + +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 From 373e73d89191fbc209574bf33ac47e0be0fdf31e Mon Sep 17 00:00:00 2001 From: Adrian Hannah Date: Fri, 9 Oct 2026 10:29:36 -0400 Subject: [PATCH 17/34] feat: Review fixes - Explicitly register supported collector sources before injecting job configuration. - On heartbeat 404, kill and join the worker, clear state, and skip outcome reporting. --- src/openhound/scheduler/service.py | 45 ++++++++++- tests/test_bhe_job_scheduling.py | 118 ++++++++++++++++++++++++++++- 2 files changed, 156 insertions(+), 7 deletions(-) diff --git a/src/openhound/scheduler/service.py b/src/openhound/scheduler/service.py index 6f53584..bb9b0a1 100644 --- a/src/openhound/scheduler/service.py +++ b/src/openhound/scheduler/service.py @@ -9,6 +9,7 @@ from concurrent.futures.process import BrokenProcessPool from dataclasses import dataclass, field from datetime import UTC, datetime +from importlib import import_module from pathlib import Path from typing import Any @@ -37,12 +38,21 @@ 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 +HEARTBEAT_INTERVAL = ( + 300 # seconds; BHE expires managed jobs after 15 minutes without a heartbeat +) + +# 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.""" + class ManagedRuntimeUnavailableError(RuntimeError): """Managed scheduling requires a configured collection runtime.""" @@ -85,6 +95,7 @@ def _usable_credential_value(value: Any) -> bool: return math.isfinite(value) return False + @dataclass class Result: results: dict @@ -140,6 +151,13 @@ def _subprocess_collect( for collector in available_collectors.collectors: if collector.name == collector_name: # pyright: ignore[reportAttributeAccessIssue] if params is not None or secrets is not None: + if ( + collector.dlt_source is None + and collector.name in MANAGED_SOURCE_MODULES + ): + collector.dlt_source = import_module( + MANAGED_SOURCE_MODULES[collector.name] + ).source _set_managed_source_environment(collector, params or {}, secrets or {}) log_fields = { "collector_extension": collector.name, @@ -239,7 +257,8 @@ def _shutdown(self) -> None: logger.info("Collection service stopping.") while self.future is not None and not self.future.done(): self._poll_running_job() - time.sleep(self.interval) + if self.future is not None: + time.sleep(self.interval) if self.future is not None: try: self._handle_completed_job(self.future) @@ -562,7 +581,10 @@ def _end_job(self, status: JobStatus, message: str) -> None: def _send_managed_heartbeat(self, job_id: str) -> None: try: self.client.heartbeat_collector_job(job_id) - except Exception as error: # noqa: BLE001 - failed heartbeats must not stop collection + 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, @@ -577,6 +599,20 @@ def _send_managed_heartbeat(self, job_id: str) -> 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 + ) + # Python 3.13 has no public pool termination API, and Future.cancel() + # cannot stop running work. Join the killed worker before accepting jobs. + processes = list((self.executor._processes or {}).values()) + for process in processes: + process.kill() + self.executor.shutdown(wait=True, cancel_futures=True) + self._clear_running_job() + self.executor = ProcessPoolExecutor(max_workers=1, max_tasks_per_child=1) + def _handle_completed_job(self, future: Future[Result]) -> None: """Report the completed worker's outcome through the active job API.""" try: @@ -619,7 +655,8 @@ def _poll_running_job(self) -> None: and time.monotonic() >= self.next_heartbeat_at ): self._send_managed_heartbeat(self.managed_job_id) - self.next_heartbeat_at = time.monotonic() + HEARTBEAT_INTERVAL + if self.managed_job_id is not None: + self.next_heartbeat_at = time.monotonic() + HEARTBEAT_INTERVAL else: try: _ = self.client.jobs_current diff --git a/tests/test_bhe_job_scheduling.py b/tests/test_bhe_job_scheduling.py index 230b393..1d3661c 100644 --- a/tests/test_bhe_job_scheduling.py +++ b/tests/test_bhe_job_scheduling.py @@ -3,8 +3,11 @@ import hashlib import json import logging +import subprocess +import sys +import time import traceback -from concurrent.futures import Future +from concurrent.futures import Future, ProcessPoolExecutor from concurrent.futures.process import BrokenProcessPool from datetime import UTC, datetime, timedelta from pathlib import Path @@ -306,6 +309,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() @@ -2233,6 +2237,25 @@ def test_managed_environment_requires_a_registered_dlt_source(): scheduler_service._set_managed_source_environment(collector, {}, {}) +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", {"node_count": 4}, {} +) +assert result.results == {"node_count": 4}, result.results +""" + 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 @@ -2353,7 +2376,7 @@ def test_claimed_job_cannot_start_while_another_job_runs( assert mock_service.job_running == 123 -def test_managed_heartbeat_failures_keep_collection_running( +def test_transient_managed_heartbeat_failures_keep_collection_running( mock_service, mock_bloodhound_api, monkeypatch, caplog ): future = Future() @@ -2369,7 +2392,7 @@ def heartbeat(job_id): if len(attempts) == 1: raise requests.Timeout("no response") if len(attempts) == 2: - raise BloodHoundHTTPError("claim lost", 404) + raise BloodHoundHTTPError("temporary failure", 503) monkeypatch.setattr(mock_service.client, "heartbeat_collector_job", heartbeat) @@ -2390,6 +2413,95 @@ def heartbeat(job_id): assert "heartbeat succeeded" in caplog.text +@pytest.mark.parametrize("during_shutdown", [False, True], ids=["poll", "shutdown"]) +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 + + +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 ): From 60b9e2ad45d6689ed1cd7213e3f097f12df2dd0f Mon Sep 17 00:00:00 2001 From: Adrian Hannah Date: Fri, 9 Oct 2026 10:29:36 -0400 Subject: [PATCH 18/34] fix: ensure abort process can safely clean up jobs --- src/openhound/scheduler/service.py | 19 ++++++++++++------- 1 file changed, 12 insertions(+), 7 deletions(-) diff --git a/src/openhound/scheduler/service.py b/src/openhound/scheduler/service.py index bb9b0a1..06917c0 100644 --- a/src/openhound/scheduler/service.py +++ b/src/openhound/scheduler/service.py @@ -604,14 +604,19 @@ def _abort_managed_job(self, job_id: str) -> None: logger.warning( "Managed collector job %s claim lost; stopping collection.", job_id ) - # Python 3.13 has no public pool termination API, and Future.cancel() - # cannot stop running work. Join the killed worker before accepting jobs. processes = list((self.executor._processes or {}).values()) - for process in processes: - process.kill() - self.executor.shutdown(wait=True, cancel_futures=True) - self._clear_running_job() - self.executor = ProcessPoolExecutor(max_workers=1, max_tasks_per_child=1) + 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: + self._clear_running_job() + self.executor = ProcessPoolExecutor(max_workers=1, max_tasks_per_child=1) def _handle_completed_job(self, future: Future[Result]) -> None: """Report the completed worker's outcome through the active job API.""" From 69b9717430d7c5ef1cb2b53e89d698eaad561bd1 Mon Sep 17 00:00:00 2001 From: Adrian Hannah Date: Fri, 9 Oct 2026 10:32:58 -0400 Subject: [PATCH 19/34] feat: start claimed managed jobs --- src/openhound/scheduler/service.py | 7 ++++++- tests/test_bhe_job_scheduling.py | 5 +---- 2 files changed, 7 insertions(+), 5 deletions(-) diff --git a/src/openhound/scheduler/service.py b/src/openhound/scheduler/service.py index 06917c0..705ae9a 100644 --- a/src/openhound/scheduler/service.py +++ b/src/openhound/scheduler/service.py @@ -230,7 +230,7 @@ def __init__( # Managed collection settings and recovery state self.managed = is_managed() if managed is None else managed self.secrets_manager = secrets_manager or AWSSecretsManager() - self.managed_job_runner = managed_job_runner + self.managed_job_runner = managed_job_runner or self._run_prepared_managed_job self._pending_managed_claim: str | None = None self._pending_managed_failure: str | None = None @@ -417,6 +417,11 @@ def _require_managed_runtime(self) -> Callable[[PreparedManagedJob], None]: ) return self.managed_job_runner + def _run_prepared_managed_job(self, prepared: PreparedManagedJob) -> None: + self.run_claimed_job( + str(prepared.job.id), prepared.job.params, prepared.credentials + ) + def _poll_managed_jobs(self) -> None: runner = self._require_managed_runtime() try: diff --git a/tests/test_bhe_job_scheduling.py b/tests/test_bhe_job_scheduling.py index 1d3661c..c629c21 100644 --- a/tests/test_bhe_job_scheduling.py +++ b/tests/test_bhe_job_scheduling.py @@ -1836,10 +1836,7 @@ def test_mode_config_routes_scheduler_to_managed_dependency(mock_service, monkey collector_name="openhound-faker", ) assert service.managed is True - with pytest.raises( - ManagedRuntimeUnavailableError, match="Managed collection runtime is required" - ): - service._poll() + assert service.managed_job_runner == service._run_prepared_managed_job def test_runtime_failure_is_sanitized_and_not_reported_as_credential_failure( From cb084f07e77aed3b06c5b5c5683f90841aad8384 Mon Sep 17 00:00:00 2001 From: Adrian Hannah Date: Fri, 9 Oct 2026 10:33:12 -0400 Subject: [PATCH 20/34] fix: heartbeat active managed jobs --- src/openhound/scheduler/service.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/src/openhound/scheduler/service.py b/src/openhound/scheduler/service.py index 705ae9a..6ce89a5 100644 --- a/src/openhound/scheduler/service.py +++ b/src/openhound/scheduler/service.py @@ -695,7 +695,10 @@ def _poll(self) -> None: return if self.managed: - self._poll_managed_jobs() + if self.job_running is None: + self._poll_managed_jobs() + else: + self._poll_running_job() return if self.job_running is None: From e88926dae4b74525dd399e0513b97faab5595135 Mon Sep 17 00:00:00 2001 From: Adrian Hannah Date: Fri, 9 Oct 2026 10:33:26 -0400 Subject: [PATCH 21/34] fix: report managed job completion --- src/openhound/scheduler/service.py | 14 +++++--- tests/test_bhe_job_scheduling.py | 51 ++++++++++++++++++++++++++++++ 2 files changed, 60 insertions(+), 5 deletions(-) diff --git a/src/openhound/scheduler/service.py b/src/openhound/scheduler/service.py index 6ce89a5..df0d869 100644 --- a/src/openhound/scheduler/service.py +++ b/src/openhound/scheduler/service.py @@ -675,11 +675,15 @@ def _poll_running_job(self) -> None: def _poll(self) -> None: """Checks if jobs are completed and if a job should be run.""" - if not self.managed: - try: - if self.future is not None and self.future.done(): - self._handle_completed_job(self.future) - except Exception: + try: + if self.future is not None and self.future.done(): + self._handle_completed_job(self.future) + except Exception: + 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." ) diff --git a/tests/test_bhe_job_scheduling.py b/tests/test_bhe_job_scheduling.py index c629c21..a48eca7 100644 --- a/tests/test_bhe_job_scheduling.py +++ b/tests/test_bhe_job_scheduling.py @@ -1557,6 +1557,57 @@ def test_managed_poll_claims_retrieves_and_hands_off_authoritative_job( assert all(forbidden not in repr(r.__dict__) for r in caplog.records) +def test_managed_poll_starts_claimed_job_and_tracks_worker( + managed_service, mock_bloodhound_api, monkeypatch +): + service = managed_service + state = mock_bloodhound_api.app.state + 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_parameter": True}, + {"credential": "private-credential-value"}, + ) + ] + 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 == [(job_id, b"")] + + 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( "document", [ From 9528b659a37c32a564b52c4ea79818fa82d04860 Mon Sep 17 00:00:00 2001 From: Adrian Hannah Date: Fri, 9 Oct 2026 13:44:45 -0400 Subject: [PATCH 22/34] refactor: centralize managed job ID normalization --- src/openhound/scheduler/service.py | 22 +++++++++++--------- tests/test_bhe_job_scheduling.py | 32 ++++++++++++++++++++++++++++++ 2 files changed, 45 insertions(+), 9 deletions(-) diff --git a/src/openhound/scheduler/service.py b/src/openhound/scheduler/service.py index df0d869..ec33ab8 100644 --- a/src/openhound/scheduler/service.py +++ b/src/openhound/scheduler/service.py @@ -12,6 +12,7 @@ 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 @@ -418,9 +419,7 @@ def _require_managed_runtime(self) -> Callable[[PreparedManagedJob], None]: return self.managed_job_runner def _run_prepared_managed_job(self, prepared: PreparedManagedJob) -> None: - self.run_claimed_job( - str(prepared.job.id), prepared.job.params, prepared.credentials - ) + self.run_claimed_job(prepared.job.id, prepared.job.params, prepared.credentials) def _poll_managed_jobs(self) -> None: runner = self._require_managed_runtime() @@ -504,11 +503,12 @@ def _start_job(self, job: Job) -> None: self._launch_job(job.id) def run_claimed_job( - self, job_id: str, params: dict[str, Any], secrets: dict[str, Any] + self, job_id: str | UUID, params: dict[str, Any], secrets: dict[str, Any] ) -> None: """Start a validated managed job; the scheduler polls it until completion.""" - job_id = str(job_id) - self._launch_job(job_id, params=params, secrets=secrets, managed=True) + # 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) def _launch_job( self, @@ -523,16 +523,20 @@ def _launch_job( raise RuntimeError(f"Job {self.job_running} is already running") logger.info("Starting job %s with collector '%s'.", job_id, self.collector_name) + managed_job_id: str | None = None if managed: - self.client.start_collector_job(str(job_id)) + if not isinstance(job_id, str): + raise TypeError("Managed job IDs must be strings") + managed_job_id = job_id else: - assert isinstance(job_id, int) + 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 = str(job_id) if managed else None + self.managed_job_id = managed_job_id self.next_heartbeat_at = ( time.monotonic() + HEARTBEAT_INTERVAL if managed else None ) diff --git a/tests/test_bhe_job_scheduling.py b/tests/test_bhe_job_scheduling.py index a48eca7..ca0446b 100644 --- a/tests/test_bhe_job_scheduling.py +++ b/tests/test_bhe_job_scheduling.py @@ -2411,6 +2411,38 @@ def test_managed_start_failure_does_not_launch_worker(mock_service, monkeypatch) 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 ): From e7fc4e5dd580f3b373d520b65e730abb022ce38a Mon Sep 17 00:00:00 2001 From: Adrian Hannah Date: Fri, 9 Oct 2026 15:26:53 -0400 Subject: [PATCH 23/34] fix: retain managed claims until job start is confirmed --- src/openhound/scheduler/service.py | 36 ++++-- tests/test_bhe_job_scheduling.py | 192 +++++++++++++++++++++++++++++ 2 files changed, 220 insertions(+), 8 deletions(-) diff --git a/src/openhound/scheduler/service.py b/src/openhound/scheduler/service.py index ec33ab8..1f0adc3 100644 --- a/src/openhound/scheduler/service.py +++ b/src/openhound/scheduler/service.py @@ -233,6 +233,7 @@ def __init__( 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 # Stores the ID of the currently running BHE job. @@ -407,10 +408,14 @@ def _claim_and_validate_managed_job(self, job_id: str) -> PreparedManagedJob | N # Secret retrieval may have consumed the remaining lease time. self._validate_managed_claim(job, job_id) - self._pending_managed_claim = None 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( @@ -427,15 +432,22 @@ def _poll_managed_jobs(self) -> None: if self._pending_managed_failure is not None: self._report_managed_credential_failure(self._pending_managed_failure) return - 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) + 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" @@ -527,7 +539,14 @@ def _launch_job( 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") @@ -697,6 +716,7 @@ def _poll(self) -> None: 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 idle and self._poll_management(): diff --git a/tests/test_bhe_job_scheduling.py b/tests/test_bhe_job_scheduling.py index ca0446b..d9c86ee 100644 --- a/tests/test_bhe_job_scheduling.py +++ b/tests/test_bhe_job_scheduling.py @@ -83,6 +83,7 @@ def mock_bloodhound_api(): )["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 @@ -157,6 +158,11 @@ async def claim_managed_job(job_id: str, request: Request): 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 @@ -186,6 +192,23 @@ async def end_managed_job(job_id: str, body: dict): @app.post("/api/v2/collector-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/collector-job-queue/{job_id}/heartbeat") @@ -1608,6 +1631,175 @@ def submit(*args): 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/collector-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 == 1 + 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_parameter": True}, + {"credential": "private-credential-value"}, + ) + 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", [ From 6ace0fcc3456fcb3d12041f1df2fbb76a710cc07 Mon Sep 17 00:00:00 2001 From: Adrian Hannah Date: Fri, 16 Oct 2026 14:01:00 -0400 Subject: [PATCH 24/34] chore(managed): Disable heartbeats for managed mode Heartbeat endpoint on the BHE side has been pushed back, so we needed to disable heartbeats on our end. Added a constant to silently returns when OH attempts to send a heartbeat. Added monkeypatch to managed-mode unit tests to allow them to run as expected. ref: BED-10078 --- src/openhound/scheduler/service.py | 6 +++++ tests/test_bhe_job_scheduling.py | 40 +++++++++++++++++++++++------- 2 files changed, 37 insertions(+), 9 deletions(-) diff --git a/src/openhound/scheduler/service.py b/src/openhound/scheduler/service.py index 1f0adc3..455c9c2 100644 --- a/src/openhound/scheduler/service.py +++ b/src/openhound/scheduler/service.py @@ -42,6 +42,9 @@ 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 = { @@ -607,6 +610,9 @@ def _end_job(self, status: JobStatus, message: str) -> None: self.client.end_job(status, message) def _send_managed_heartbeat(self, job_id: str) -> None: + if _MANAGED_HEARTBEATS_DISABLED: + return + try: self.client.heartbeat_collector_job(job_id) except Exception as error: # noqa: BLE001 - transient heartbeat failures must not stop collection diff --git a/tests/test_bhe_job_scheduling.py b/tests/test_bhe_job_scheduling.py index d9c86ee..5365b28 100644 --- a/tests/test_bhe_job_scheduling.py +++ b/tests/test_bhe_job_scheduling.py @@ -588,9 +588,7 @@ def test_start_and_heartbeat_collector_job_use_bodyless_requests( 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_start_requests == [(job_id, b"")] assert mock_bloodhound_api.app.state.collector_job_heartbeat_requests == [ (job_id, b"") ] @@ -1618,7 +1616,7 @@ def submit(*args): service._poll() assert len(state.collector_job_claim_requests) == 1 assert len(state.collector_job_queue_requests) == 1 - assert state.collector_job_heartbeat_requests == [(job_id, b"")] + assert state.collector_job_heartbeat_requests == [] state.collector_job_queue_response = load_json( "collector_jobs_available_empty.json" @@ -1731,7 +1729,11 @@ def flaky_request(**kwargs): @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 + stateful_managed_service, + mock_bloodhound_api, + monkeypatch, + status, + lost_response_first, ): service = stateful_managed_service state = mock_bloodhound_api.app.state @@ -2429,6 +2431,8 @@ def expired_token(self, **kwargs): assert forbidden not in caplog.text assert forbidden not in formatted_error assert all(forbidden not in repr(record.__dict__) for record in caplog.records) + + def test_managed_worker_sets_dlt_environment_before_pipeline(monkeypatch, caplog): collector = SimpleNamespace( name="faker", @@ -2510,6 +2514,13 @@ def sleep(self, 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 ): @@ -2546,12 +2557,19 @@ def fail_unmanaged_checkin(self): "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)): + 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 + 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) @@ -2648,6 +2666,7 @@ def test_claimed_job_cannot_start_while_another_job_runs( 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 ): @@ -2686,6 +2705,7 @@ def heartbeat(job_id): @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 ): @@ -2735,6 +2755,7 @@ def lose_claim(job_id): 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 ): @@ -2797,6 +2818,7 @@ def test_managed_worker_failure_ends_job_failed( ] +@pytest.mark.usefixtures("enabled_managed_heartbeats") def test_shutdown_keeps_heartbeating_until_managed_worker_finishes( mock_service, mock_bloodhound_api, monkeypatch ): From 672d95cacfcdee23a25565fc89639b6c3259647d Mon Sep 17 00:00:00 2001 From: Stran Dutton <65475509+StranDutton@users.noreply.github.com> Date: Mon, 19 Oct 2026 10:41:17 -0500 Subject: [PATCH 25/34] BED-10083: Improve managed-mode AWS failure diagnostics (#100) * fix(managed): improve AWS credential failure diagnostics - BED-10083 * fix(managed): sanitize unexpected AWS runtime failures - BED-10083 * fix(managed): preserve safe AWS failure details - BED-10083 --- .../core/clients/aws_secrets_manager.py | 368 +++++++++++++++-- src/openhound/core/clients/bhe_credentials.py | 31 +- src/openhound/scheduler/service.py | 42 +- tests/test_aws_secrets_manager.py | 385 +++++++++++++++++- tests/test_bhe_credentials.py | 120 +++++- tests/test_bhe_job_scheduling.py | 167 ++++++-- 6 files changed, 1029 insertions(+), 84 deletions(-) diff --git a/src/openhound/core/clients/aws_secrets_manager.py b/src/openhound/core/clients/aws_secrets_manager.py index f600098..27ece94 100644 --- a/src/openhound/core/clients/aws_secrets_manager.py +++ b/src/openhound/core/clients/aws_secrets_manager.py @@ -1,5 +1,6 @@ import json import logging +import re from collections.abc import Iterable, Mapping from typing import Any, Protocol @@ -10,17 +11,32 @@ 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. @@ -66,18 +82,30 @@ def batch_get_secret_value( 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 + 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.""" @@ -88,7 +116,11 @@ class SecretBatchError(SecretRetrievalError): def __init__(self, failures: Mapping[str, SecretRetrievalError]) -> None: self.failures = dict(failures) - super().__init__("One or more AWS secret retrievals failed") + super().__init__( + "; ".join( + f"secret {secret_id}: {error}" for secret_id, error in failures.items() + ) + ) SecretValue = str | dict[str, Any] @@ -137,12 +169,20 @@ def _get_secrets_in_batches( _log_debug("batch_get_secret_value", "attempt", len(chunk), chunk_index + 1) try: response = client.batch_get_secret_value(SecretIdList=chunk) - except (ClientError, BotoCoreError) as error: - classified = _classify_aws_error(error) + 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) + chunk_values, chunk_failures = _parse_batch_response( + response, chunk, _client_region(client) + ) values.update(chunk_values) failures.update(chunk_failures) @@ -177,7 +217,13 @@ def _get_client(self) -> SecretsManagerClient: "secretsmanager", config=DEFAULT_AWS_CLIENT_CONFIG, ) - except BotoCoreError as error: + except ( + ClientError, + BotoCoreError, + FileNotFoundError, + PermissionError, + RuntimeError, + ) as error: raise _classify_aws_error(error) from None def _retrieve_secret( @@ -187,8 +233,16 @@ def _retrieve_secret( try: response = client.get_secret_value(SecretId=secret_id) value = _parse_secret_value(response) - except (ClientError, BotoCoreError) as error: - classified = _classify_aws_error(error) + 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: @@ -204,6 +258,10 @@ def _chunks(values: list[str], size: int) -> Iterable[list[str]]: 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): @@ -233,6 +291,7 @@ def _parse_secret_value(response: Mapping[str, Any]) -> SecretValue: 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] = {} @@ -255,7 +314,20 @@ def _parse_batch_response( 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(entry.get("ErrorCode")) + 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: @@ -267,37 +339,259 @@ def _parse_batch_response( def _classify_aws_error( - error: ClientError | BotoCoreError | str | None, + error: ClientError + | BotoCoreError + | FileNotFoundError + | PermissionError + | RuntimeError, + *, + secret_id: str | None = None, + region: str | None = None, ) -> SecretRetrievalError: - """Map a boto error or AWS error code to a public retrieval exception.""" + """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): + return _classify_runtime_error(error, context) + if isinstance(error, (FileNotFoundError, PermissionError)): + return _classify_file_error(error, context, location) if isinstance(error, ClientError): - error = error.response.get("Error", {}).get("Code") - elif isinstance( - error, - ( - ConfigNotFound, - CredentialRetrievalError, - InvalidConfigError, - NoCredentialsError, - NoRegionError, - PartialCredentialsError, - ProfileNotFound, - ), - ): + return _classify_client_error(error, context, location) + if isinstance(error, _PROVIDER_ERROR_TYPES): + return _classify_provider_error(error, context) + return _classify_transport_error(error, context, location) + + +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( - "AWS Secrets Manager client configuration is invalid" + 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, ) - if error == "ResourceNotFoundException": - return SecretNotFoundError("AWS secret was not found") - if error in { - "AccessDenied", - "AccessDeniedException", - "AuthorizationError", - "UnauthorizedOperation", - "UnrecognizedClientException", - }: - return SecretPermissionError("AWS Secrets Manager denied access to the secret") - return SecretRequestError("AWS Secrets Manager request could not be completed") + 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( @@ -339,4 +633,6 @@ def _log_failure( } if batch_number is not None: extra["batch_number"] = batch_number - logger.error("Secret retrieval failed", extra=extra) + extra.update(error.context) + # Managed credential callers report the cause at normal logging levels. + logger.debug("AWS %s failed: %s", operation, error, extra=extra) diff --git a/src/openhound/core/clients/bhe_credentials.py b/src/openhound/core/clients/bhe_credentials.py index 37b4d3b..73dcfaa 100644 --- a/src/openhound/core/clients/bhe_credentials.py +++ b/src/openhound/core/clients/bhe_credentials.py @@ -6,7 +6,7 @@ from typing import Protocol from uuid import UUID -from .aws_secrets_manager import AWSSecretsManager, SecretValue +from .aws_secrets_manager import AWSSecretsManager, SecretRetrievalError, SecretValue logger = logging.getLogger(__name__) @@ -45,7 +45,23 @@ def refresh(self) -> BHECredentials: 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 @@ -78,12 +94,19 @@ def poll_until_available(self, interval: int) -> BHECredentials: while True: try: - return self.refresh() - except Exception: + 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( - "BHE credentials are unavailable; retrying in %s seconds.", + "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/scheduler/service.py b/src/openhound/scheduler/service.py index 455c9c2..24a7dc2 100644 --- a/src/openhound/scheduler/service.py +++ b/src/openhound/scheduler/service.py @@ -17,7 +17,10 @@ import openhound import openhound.core.logging as openhound_logging from openhound.config import is_managed -from openhound.core.clients.aws_secrets_manager import AWSSecretsManager +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 ( @@ -64,6 +67,10 @@ class ManagedRuntimeUnavailableError(RuntimeError): 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"): + super().__init__(message) + self.summary = summary + class ManagedJobProcessingError(RuntimeError): """A sanitized managed polling failure that can be retried next poll.""" @@ -238,6 +245,7 @@ def __init__( 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. @@ -332,14 +340,16 @@ def _report_managed_credential_failure(self, job_id: str) -> None: self.client.end_managed_job( job_id, ManagedJobOutcome.FAILED, - failure_message="Credential validation failed", + failure_message=self._pending_managed_failure_summary, ) self._pending_managed_failure = None def _load_managed_credentials(self, job: CollectorJob) -> dict[str, Any]: """Retrieve and validate credentials, exposing only a safe failure message.""" if job.secret_key_id is None or not job.secret_key_id.strip(): - raise CredentialValidationError("Credential validation failed") + raise CredentialValidationError( + "Managed job has no credential secret identifier" + ) logger.debug( "Retrieving managed collector credentials.", @@ -353,9 +363,16 @@ def _load_managed_credentials(self, job: CollectorJob) -> dict[str, Any]: credentials = self.secrets_manager.get_secret(job.secret_key_id) if isinstance(credentials, dict) and _usable_credential_value(credentials): return credentials - except Exception: # noqa: BLE001 - redact secret reader and validation errors - raise CredentialValidationError("Credential validation failed") from None - raise CredentialValidationError("Credential validation failed") + 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 non-empty JSON object with usable values" + ) def _validate_managed_claim(self, job: CollectorJob, job_id: str) -> None: """Check the job identity and lease confirmed by BHE's claim response.""" @@ -403,9 +420,15 @@ def _claim_and_validate_managed_job(self, job_id: str) -> PreparedManagedJob | N logger.info("Job successfully picked up", extra=context) try: credentials = self._load_managed_credentials(job) - except CredentialValidationError: - logger.error("Credential validation failed", extra=context) + except CredentialValidationError as error: + logger.error( + "Cannot prepare managed job %s: %s. Reporting this job failed.", + job_id, + str(error).rstrip("."), + extra={**context, "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 @@ -772,7 +795,8 @@ def start(self) -> None: ) except ManagedJobProcessingError: logger.error( - "Managed scheduler poll failed; retrying next poll." + "Managed scheduler poll failed; retrying next poll in %s seconds.", + self.interval, ) else: self._poll() diff --git a/tests/test_aws_secrets_manager.py b/tests/test_aws_secrets_manager.py index a6a8f7f..85f5d29 100644 --- a/tests/test_aws_secrets_manager.py +++ b/tests/test_aws_secrets_manager.py @@ -1,4 +1,5 @@ import base64 +import errno import json import logging import traceback @@ -14,9 +15,12 @@ ConfigNotFound, CredentialRetrievalError, EndpointConnectionError, + HTTPClientError, InvalidConfigError, + MetadataRetrievalError, NoCredentialsError, NoRegionError, + ParamValidationError, PartialCredentialsError, ProfileNotFound, ) @@ -31,9 +35,11 @@ ) LOGGER_NAME = "openhound.core.clients.aws_secrets_manager" -SECRET_ID = "secret-id-that-must-not-be-logged" +SECRET_ID = "openhound/bhe" SECRET_VALUE = "secret-value-that-must-not-be-logged" -PROVIDER_MESSAGE = "provider message that must not be exposed" +PROVIDER_MESSAGE = ( + "role arn:aws:iam::123456789012:role/openhound cannot read the secret" +) @pytest.fixture(autouse=True) @@ -226,7 +232,7 @@ def test_get_secrets_continues_and_aggregates_typed_failures_safely( ), ] ) - forbidden = set(secret_ids) | {SECRET_VALUE, PROVIDER_MESSAGE} + forbidden = {SECRET_VALUE} with ( caplog.at_level(logging.DEBUG, logger=LOGGER_NAME), @@ -251,9 +257,12 @@ def test_get_secrets_continues_and_aggregates_typed_failures_safely( error_records = [ cast(Any, record) for record in caplog.records - if record.levelno == logging.ERROR + 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 ) @@ -286,7 +295,11 @@ def test_get_secret_classifies_aws_errors( with pytest.raises(expected_type) as raised: AWSSecretsManager(client).get_secret(SECRET_ID) - assert_safe_failure(raised.value, {SECRET_ID, PROVIDER_MESSAGE}) + 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( @@ -314,8 +327,15 @@ def fail_client(*_args: object, **_kwargs: object) -> NoReturn: "openhound.core.clients.aws_secrets_manager.boto3.client", fail_client ) - with pytest.raises(AWSConfigurationError): + 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( @@ -437,7 +457,7 @@ def test_real_sdk_http_pipeline_does_not_expose_secrets( if response_kind == "provider_error": response_body = { "__type": "AccessDeniedException", - "message": f"{PROVIDER_MESSAGE}: {SECRET_ID} {SECRET_VALUE}", + "message": f"{PROVIDER_MESSAGE}: {SECRET_ID}", } status = 400 else: @@ -449,9 +469,7 @@ def test_real_sdk_http_pipeline_does_not_expose_secrets( status = 200 raw_body = json.dumps(response_body).encode() forbidden = { - SECRET_ID, SECRET_VALUE, - PROVIDER_MESSAGE, encoded_document, raw_body.decode(), } @@ -568,3 +586,352 @@ def send(request): 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 index c3d27c3..a82d59c 100644 --- a/tests/test_bhe_credentials.py +++ b/tests/test_bhe_credentials.py @@ -18,11 +18,15 @@ def get_secret(self, secret_id): def test_aws_bhe_credentials_loads_token_pair_from_json_secret(): - manager = FakeSecretsManager({"token_id": "12345678-1234-1234-1234-123456789abc", "token_key": "key"}) + 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 result == BHECredentials( + token_id="12345678-1234-1234-1234-123456789abc", token_key="key" + ) assert manager.requested == ["production/bhe"] @@ -42,3 +46,115 @@ def test_aws_bhe_credentials_loads_token_pair_from_json_secret(): 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 5365b28..40b469f 100644 --- a/tests/test_bhe_job_scheduling.py +++ b/tests/test_bhe_job_scheduling.py @@ -1838,7 +1838,7 @@ def retrieve(*, SecretId): "job_id": state.collector_job_claim_response["id"], "body": { "outcome": "failed", - "failure_message": "Credential validation failed", + "failure_message": "Credential document is invalid", }, } ] @@ -1862,28 +1862,46 @@ def test_missing_managed_secret_reference_fails_without_local_fallback( assert managed_service.secret_requests == managed_service.handoffs == [] assert ( state.collector_job_end_requests[0]["body"]["failure_message"] - == "Credential validation failed" + == "Credential document is invalid" ) @pytest.mark.parametrize( - "error", + "error,summary", [ - ClientError( - { - "Error": { - "Code": "AccessDeniedException", - "Message": "private-credential-value", - } - }, - "GetSecretValue", + ( + 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", ), - RuntimeError("private-credential-value"), ], - ids=["aws-permission-error", "unexpected-reader-error"], + ids=["aws-permission-error", "aws-missing-secret", "unexpected-reader-error"], ) -def test_secret_retrieval_failure_is_redacted_and_reported_as_failed( - managed_service, mock_bloodhound_api, caplog, 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 @@ -1900,9 +1918,18 @@ def retrieve(*, SecretId): assert state.managed_events == ["claim", "retrieve", "end"] assert state.collector_job_end_requests[0]["body"] == { "outcome": "failed", - "failure_message": "Credential validation 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 @@ -1997,7 +2024,14 @@ def test_exhausted_credential_failure_end_is_retained_and_retried_before_discove managed_service, mock_bloodhound_api, monkeypatch, caplog, resolved_status ): state = mock_bloodhound_api.app.state - state.collector_job_claim_response["secret_key_id"] = None + 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 = [] @@ -2024,6 +2058,10 @@ def test_exhausted_credential_failure_end_is_retained_and_retried_before_discove 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) @@ -2153,7 +2191,9 @@ def next_poll(seconds): assert managed_service.handoffs == [] assert [ r.getMessage() for r in caplog.records if "poll failed" in r.getMessage() - ] == ["Managed scheduler poll failed; retrying next poll."] + ] == [ + 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 @@ -2162,6 +2202,85 @@ def next_poll(seconds): 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"), + ("metadata", "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 == "metadata": + managed_service.exit = True + else: + monkeypatch.setattr( + managed_service.client, "update_client_metadata", lambda: None + ) + 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 == [] + if stage == "metadata": + assert "Unable to update client metadata" in caplog.text + else: + 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 @@ -2398,7 +2517,7 @@ def test_managed_recovery_precedes_management_then_allows_support_work( assert len(sent) == 1 -def test_support_bundle_tracebacks_keep_aws_auth_errors_redacted( +def test_support_bundle_tracebacks_preserve_aws_auth_context( mock_service, monkeypatch, caplog ): secret_reference = "private-bootstrap-secret-reference" @@ -2424,13 +2543,13 @@ def expired_token(self, **kwargs): with caplog.at_level(logging.DEBUG), pytest.raises(SecretPermissionError) as raised: mock_service._send_support_bundle(operation) - assert "AWS Secrets Manager denied access to the secret" in caplog.text + 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 forbidden in (secret_reference, provider_message): - assert forbidden not in caplog.text - assert forbidden not in formatted_error - assert all(forbidden not in repr(record.__dict__) for record in caplog.records) + for context in (secret_reference, provider_message): + assert context in caplog.text + assert context in formatted_error def test_managed_worker_sets_dlt_environment_before_pipeline(monkeypatch, caplog): From 0ae0fc35424c1a7f8ce5508bcff3f393c85fbf85 Mon Sep 17 00:00:00 2001 From: Stran Dutton <65475509+StranDutton@users.noreply.github.com> Date: Mon, 19 Oct 2026 10:46:27 -0500 Subject: [PATCH 26/34] fix(managed): request jobs in priority order (BED-10130) (#101) --- src/openhound/core/clients/bloodhound_enterprise.py | 2 +- tests/test_bhe_job_scheduling.py | 10 ++++++++-- 2 files changed, 9 insertions(+), 3 deletions(-) diff --git a/src/openhound/core/clients/bloodhound_enterprise.py b/src/openhound/core/clients/bloodhound_enterprise.py index 9b82f54..7f1622d 100644 --- a/src/openhound/core/clients/bloodhound_enterprise.py +++ b/src/openhound/core/clients/bloodhound_enterprise.py @@ -137,7 +137,7 @@ def available_collector_jobs(self, job_key: str) -> CollectorJobsAvailable: encoded_job_key = quote(f"eq:{job_key}", safe="") path = ( "/api/v2/collector-job-queue/available?limit=1&job_key=" - f"{encoded_job_key}" + f"{encoded_job_key}&sort_by=-priority&sort_by=created_at" ) logger.debug( "Polling managed collector job queue.", diff --git a/tests/test_bhe_job_scheduling.py b/tests/test_bhe_job_scheduling.py index 40b469f..a76cdf1 100644 --- a/tests/test_bhe_job_scheduling.py +++ b/tests/test_bhe_job_scheduling.py @@ -122,7 +122,9 @@ async def jobs_current(): @app.get("/api/v2/collector-job-queue/available") async def collector_job_queue(request: Request): - app.state.collector_job_queue_requests.append(dict(request.query_params)) + 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: @@ -795,7 +797,11 @@ def test_managed_collector_job_check_requests_first_queue_job( 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"} + { + "limit": "1", + "job_key": "eq:openhound-faker", + "sort_by": ["-priority", "created_at"], + } ] assert mock_bloodhound_api.app.state.jobs_available_requests == 0 From debf3ca384450f1505d13c8385b63a4abd162b16 Mon Sep 17 00:00:00 2001 From: Stran Dutton <65475509+StranDutton@users.noreply.github.com> Date: Mon, 19 Oct 2026 14:43:42 -0500 Subject: [PATCH 27/34] fix(managed): use 'collie' job queue API paths (BED-10135) (#103) --- AGENTS.md | 23 +++++++++++++++++++ .../core/clients/bloodhound_enterprise.py | 10 ++++---- tests/test_bhe_job_scheduling.py | 18 +++++++-------- 3 files changed, 37 insertions(+), 14 deletions(-) create mode 100644 AGENTS.md diff --git a/AGENTS.md b/AGENTS.md new file mode 100644 index 0000000..c03e646 --- /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/src/openhound/core/clients/bloodhound_enterprise.py b/src/openhound/core/clients/bloodhound_enterprise.py index 7f1622d..7989937 100644 --- a/src/openhound/core/clients/bloodhound_enterprise.py +++ b/src/openhound/core/clients/bloodhound_enterprise.py @@ -136,7 +136,7 @@ def available_collector_jobs(self, job_key: str) -> CollectorJobsAvailable: """ encoded_job_key = quote(f"eq:{job_key}", safe="") path = ( - "/api/v2/collector-job-queue/available?limit=1&job_key=" + "/api/v2/collie-job-queue/available?limit=1&job_key=" f"{encoded_job_key}&sort_by=-priority&sort_by=created_at" ) logger.debug( @@ -169,7 +169,7 @@ def claim_managed_job(self, job_id: str) -> CollectorJob: 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/collector-job-queue/{job_id}/claim" + path = f"/api/v2/collie-job-queue/{job_id}/claim" for attempt in range(1, MANAGED_JOB_CLAIM_MAX_ATTEMPTS + 1): context = { "endpoint": path, @@ -227,7 +227,7 @@ 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/collector-job-queue/{job_id}/start", + path=f"/api/v2/collie-job-queue/{job_id}/start", timeout=( MANAGED_JOB_LIFECYCLE_CONNECT_TIMEOUT_SECONDS, MANAGED_JOB_LIFECYCLE_READ_TIMEOUT_SECONDS, @@ -238,7 +238,7 @@ 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/collector-job-queue/{job_id}/heartbeat", + path=f"/api/v2/collie-job-queue/{job_id}/heartbeat", timeout=( MANAGED_JOB_LIFECYCLE_CONNECT_TIMEOUT_SECONDS, MANAGED_JOB_LIFECYCLE_READ_TIMEOUT_SECONDS, @@ -253,7 +253,7 @@ def end_managed_job( metadata: dict[str, Any] | None = None, ) -> None: """End a claimed managed collector job.""" - path = f"/api/v2/collector-job-queue/{job_id}/end" + 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 diff --git a/tests/test_bhe_job_scheduling.py b/tests/test_bhe_job_scheduling.py index a76cdf1..b22bfc8 100644 --- a/tests/test_bhe_job_scheduling.py +++ b/tests/test_bhe_job_scheduling.py @@ -120,7 +120,7 @@ async def jobs_current(): app.state.jobs_current_requests += 1 return Response(status_code=404) - @app.get("/api/v2/collector-job-queue/available") + @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") @@ -147,7 +147,7 @@ async def start_job(body: dict): app.state.start_payload = body return load_json("job_start.json") - @app.post("/api/v2/collector-job-queue/{job_id}/claim") + @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( @@ -178,7 +178,7 @@ async def end_job(body: dict): app.state.end_payload = body return load_json("job_end.json") - @app.post("/api/v2/collector-job-queue/{job_id}/end") + @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}) @@ -191,7 +191,7 @@ async def end_managed_job(job_id: str, body: dict): "raw-provider-secret" if status != 200 else "", status_code=status ) - @app.post("/api/v2/collector-job-queue/{job_id}/start") + @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: @@ -213,7 +213,7 @@ async def start_collector_job(job_id: str, request: Request): ).isoformat() return Response(status_code=200) - @app.post("/api/v2/collector-job-queue/{job_id}/heartbeat") + @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()) @@ -683,7 +683,7 @@ def flaky_request(*args, **kwargs): assert any( record.levelno == logging.DEBUG and getattr(record, "endpoint", None) - == "/api/v2/collector-job-queue/collector-job-123/end" + == "/api/v2/collie-job-queue/collector-job-123/end" for record in caplog.records ) @@ -1540,7 +1540,7 @@ def request(**kwargs): assert captured == [ { "method": "POST", - "path": f"/api/v2/collector-job-queue/{job_id}/claim", + "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, @@ -1696,7 +1696,7 @@ def flaky_request(**kwargs): # Recovery must repeat start, not claim: BHE rejects re-claim of a running job. service._poll() - assert starts == [f"/api/v2/collector-job-queue/{job_id}/start"] * 4 + 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 == 1 @@ -2320,7 +2320,7 @@ def test_committed_claim_with_lost_responses_is_reconciled_before_discovery( # 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/collector-job-queue/available").json() + 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" ] From 4e2c6451514a73dd91b721bef516028cea9eb67f Mon Sep 17 00:00:00 2001 From: Justin LaRose Date: Tue, 20 Oct 2026 15:07:09 -0400 Subject: [PATCH 28/34] BED-9517 - Ingest path for Managed OpenHound (#102) * adding new ingest path and associated tests. BED-9517 * adding the new endpoint name for BHE. BED-9517 --- src/openhound/core/clients/bloodhound.py | 18 +- .../core/clients/bloodhound_enterprise.py | 34 +- src/openhound/scheduler/dataflow.py | 16 +- src/openhound/scheduler/service.py | 133 +++- tests/test_bhe_job_scheduling.py | 690 ++++++++++++++++++ tests/test_collector_job_ingest.py | 114 +++ tests/test_scheduler_dataflow.py | 59 ++ 7 files changed, 1047 insertions(+), 17 deletions(-) create mode 100644 tests/test_collector_job_ingest.py create mode 100644 tests/test_scheduler_dataflow.py diff --git a/src/openhound/core/clients/bloodhound.py b/src/openhound/core/clients/bloodhound.py index c33427b..e96f529 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 7989937..f9c9e90 100644 --- a/src/openhound/core/clients/bloodhound_enterprise.py +++ b/src/openhound/core/clients/bloodhound_enterprise.py @@ -9,8 +9,9 @@ from collections.abc import Callable from enum import Enum from pathlib import Path -from typing import Any, TypeVar +from typing import Any, BinaryIO, TypeVar from urllib.parse import quote +from uuid import UUID import requests @@ -87,11 +88,14 @@ 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, ): """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, @@ -104,6 +108,8 @@ def request( 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 @@ -350,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( diff --git a/src/openhound/scheduler/dataflow.py b/src/openhound/scheduler/dataflow.py index 65bc185..252abe6 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 24a7dc2..44e95ea 100644 --- a/src/openhound/scheduler/service.py +++ b/src/openhound/scheduler/service.py @@ -2,7 +2,9 @@ import logging import math import os +import shutil import signal +import tempfile import time from collections.abc import Callable from concurrent.futures import Future, ProcessPoolExecutor @@ -111,6 +113,10 @@ def _usable_credential_value(value: Any) -> bool: class Result: results: dict 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 _set_managed_source_environment( @@ -141,6 +147,7 @@ def _subprocess_collect( job_id: int | str, params: dict[str, Any] | None = None, secrets: dict[str, Any] | None = None, + managed: bool = False, ) -> Result: """A subprocess which runs the DLT pipeline for the specified collector. @@ -149,6 +156,7 @@ def _subprocess_collect( 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. @@ -189,13 +197,30 @@ def _subprocess_collect( 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.") @@ -266,13 +291,15 @@ def _exit_handler(self, sig: int, frame): logger.warning(f"Received signal {sig}, shutting down gracefully.") def _shutdown(self) -> None: - """Finish the active job and keep checking in until its worker stops.""" + """Wait for the worker and finish any pending managed uploads.""" logger.info("Collection service stopping.") 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.future is not None: + 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: @@ -283,6 +310,22 @@ def _shutdown(self) -> None: self.executor.shutdown(wait=True, cancel_futures=True) 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.future = None self.job_running = None @@ -592,6 +635,7 @@ def _launch_job( job_id, dict(params or {}) if managed else None, dict(secrets or {}) if managed else None, + *((True,) if managed else ()), ) except BrokenProcessPool: logger.exception("Failed to submit job %s: process pool is broken.", job_id) @@ -661,6 +705,7 @@ def _abort_managed_job(self, job_id: str) -> None: 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: @@ -672,11 +717,21 @@ def _abort_managed_job(self, job_id: str) -> None: 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 = ProcessPoolExecutor(max_workers=1, max_tasks_per_child=1) def _handle_completed_job(self, future: Future[Result]) -> None: """Report the completed worker's outcome through the active job API.""" + result: Result | None = None try: result = future.result() logger.info(f"Job {result.job_id} completed successfully, notifying BHE.") @@ -705,11 +760,55 @@ def _handle_completed_job(self, future: Future[Result]) -> None: f"Unexpected error while running '{self.collector_name}' collector" ) else: - status = JobStatus.COMPLETE - message = f"Collector '{self.collector_name}' completed successfully" + 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() + 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 ( @@ -725,6 +824,28 @@ def _poll_running_job(self) -> None: except Exception: logger.exception("Error checking in-progress job.") + def start_claimed_managed_job(self, job: CollectorJob) -> None: + """Run a managed job after the queue caller has claimed it. + + Claiming and renewing the job belongs to the queue scheduling flow. + """ + if not self.managed: + raise ValueError("Managed collections are not enabled") + if self.job_running is not None or self.future is not None: + raise ValueError("A collector job is already running") + if job.status != "running" or job.claimed_by is None: + raise ValueError("The managed collector job must be claimed and running") + self.job_running = str(job.id) + self.managed_job_id = str(job.id) + self.next_heartbeat_at = time.monotonic() + HEARTBEAT_INTERVAL + try: + self.future = self.executor.submit( + _subprocess_collect, self.collector_name, str(job.id), None, None, True + ) + except BrokenProcessPool as error: + self.future = Future() + self.future.set_exception(error) + def _poll(self) -> None: """Checks if jobs are completed and if a job should be run.""" try: diff --git a/tests/test_bhe_job_scheduling.py b/tests/test_bhe_job_scheduling.py index b22bfc8..d6fed0b 100644 --- a/tests/test_bhe_job_scheduling.py +++ b/tests/test_bhe_job_scheduling.py @@ -107,6 +107,9 @@ 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(): @@ -229,6 +232,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 @@ -345,6 +358,9 @@ def shutdown(self, *args, **kwargs): return None def mock_request(method, url, **kwargs): + 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: @@ -1094,6 +1110,18 @@ def test_scheduler_ingest_opengraph( 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: @@ -1613,6 +1641,7 @@ def submit(*args): job_id, {"claimed_parameter": True}, {"credential": "private-credential-value"}, + True, ) ] assert service.job_running == service.managed_job_id == job_id @@ -1711,6 +1740,7 @@ def flaky_request(**kwargs): job_id, {"claimed_parameter": True}, {"credential": "private-credential-value"}, + True, ) assert service.future is future for forbidden in ("raw-provider-secret", "private-credential-value"): @@ -2705,6 +2735,7 @@ def fail_unmanaged_checkin(self): "collector-job-123", {"node_count": 4}, {"token": "secret-value"}, + True, ) assert mock_bloodhound_api.app.state.collector_job_start_requests == [ ("collector-job-123", b"") @@ -2968,3 +2999,662 @@ def finish_worker(now): {"job_id": "collector-job-123", "body": {"outcome": "succeeded"}} ] assert mock_service.future is None + + +@pytest.mark.parametrize("failure_phase", ["upload", "success_report", "failure_report"]) +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() + + +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" + + +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() + + +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 +): + 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_start_claimed_managed_job_passes_uuid_and_staging_mode_to_worker(mock_service): + job = CollectorJob.model_validate( + load_json("collector_jobs_available_with_job.json")["data"]["jobs"][0] + ) + job.status = "running" + job.claimed_by = UUID("c894ea20-8c11-49bd-83f8-5550c8255148") + mock_service.managed = True + mock_service.start_claimed_managed_job(job) + args, _kwargs, future = mock_service.executor.submitted[0] + assert args == ( + _subprocess_collect, "openhound-faker", str(job.id), None, None, True + ) + assert mock_service.future is future + assert mock_service.job_running == str(job.id) + + +@pytest.mark.parametrize("managed", [False, True]) +def test_start_claimed_managed_job_rejects_unclaimed_jobs(mock_service, managed): + job = CollectorJob.model_validate( + load_json("collector_jobs_available_with_job.json")["data"]["jobs"][0] + ) + mock_service.managed = managed + with pytest.raises(ValueError): + mock_service.start_claimed_managed_job(job) + assert mock_service.executor.submitted == [] + assert mock_service.job_running is None + + +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(RuntimeError, match="conversion failed"): + _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 0000000..527e177 --- /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_scheduler_dataflow.py b/tests/test_scheduler_dataflow.py new file mode 100644 index 0000000..3df45af --- /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)] From 5f50edf6e954fb39c86e605ed2d38661aa6f5bb2 Mon Sep 17 00:00:00 2001 From: Stran Dutton <65475509+StranDutton@users.noreply.github.com> Date: Tue, 20 Oct 2026 16:11:03 -0500 Subject: [PATCH 29/34] fix(managed): align profile and schedule IDs with BHE UUIDs - BED-10179 (#104) --- src/openhound/core/clients/models/jobs.py | 4 +- tests/test_bhe_job_scheduling.py | 175 +++++++++++++++--- .../collector_jobs_available_with_jobs.json | 4 +- 3 files changed, 151 insertions(+), 32 deletions(-) diff --git a/src/openhound/core/clients/models/jobs.py b/src/openhound/core/clients/models/jobs.py index 8ecd5a5..fa301da 100644 --- a/src/openhound/core/clients/models/jobs.py +++ b/src/openhound/core/clients/models/jobs.py @@ -58,8 +58,8 @@ class JobsEnd(BaseModel): class CollectorJob(BaseModel): id: UUID - job_schedule_id: int | None - job_profile_id: int | None + job_schedule_id: UUID | None + job_profile_id: UUID | None job_type_id: int job_key: str params_version: str diff --git a/tests/test_bhe_job_scheduling.py b/tests/test_bhe_job_scheduling.py index d6fed0b..f1fe226 100644 --- a/tests/test_bhe_job_scheduling.py +++ b/tests/test_bhe_job_scheduling.py @@ -791,6 +791,7 @@ def test_jobs_starts_new_job(mock_service, mock_bloodhound_api): 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 @@ -822,9 +823,13 @@ def test_managed_collector_job_check_requests_first_queue_job( assert mock_bloodhound_api.app.state.jobs_available_requests == 0 -def test_collector_job_queue_response_parses_contract_fields(): +def test_collector_job_queue_response_parses_contract_fields( + mock_service, mock_bloodhound_api +): payload = load_json("collector_jobs_available_with_jobs.json") - response = CollectorJobsAvailable.model_validate(payload) + 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 @@ -835,6 +840,7 @@ def test_collector_job_queue_response_parses_contract_fields(): 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.params_version == "v2" assert job.params == { @@ -857,8 +863,20 @@ def test_collector_job_queue_response_parses_contract_fields(): 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.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", "scope_client_id", "claimed_by"]) +@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") @@ -868,7 +886,9 @@ def test_collector_job_queue_response_rejects_invalid_uuids(field): CollectorJobsAvailable.model_validate(payload) -@pytest.mark.parametrize("field", ["scope_client_id", "claimed_by"]) +@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") @@ -878,6 +898,35 @@ def test_collector_job_queue_response_requires_nullable_uuid_fields(field): 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 ): @@ -921,8 +970,10 @@ def submit(*args): 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 == [] @@ -1515,9 +1566,9 @@ def managed_service(mock_service, mock_bloodhound_api): "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" - ) + 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", @@ -1548,9 +1599,20 @@ def start_handoff(prepared): 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 + 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 @@ -1565,6 +1627,11 @@ def request(**kwargs): 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 captured == [ { "method": "POST", @@ -1612,11 +1679,22 @@ def test_managed_poll_claims_retrieves_and_hands_off_authoritative_job( assert all(forbidden not in repr(r.__dict__) for r in caplog.records) +@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 + 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 = [] @@ -3001,7 +3079,10 @@ def finish_worker(now): assert mock_service.future is None -@pytest.mark.parametrize("failure_phase", ["upload", "success_report", "failure_report"]) +@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 ): @@ -3069,7 +3150,9 @@ def check_retry_state(now): 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.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 @@ -3115,7 +3198,9 @@ def end(job, outcome, **kwargs): mock_service._poll() assert uploads == [b"file-0"] - assert len(ends) == 1, "Missing staged data must fail the job instead of retrying forever" + 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 @@ -3124,6 +3209,7 @@ def end(job, outcome, **kwargs): 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 ): @@ -3163,7 +3249,8 @@ def lose_claim(job): 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", + mock_service.client, + "end_managed_job", lambda *args, **kwargs: ends.append((args, kwargs)), ) @@ -3177,9 +3264,12 @@ def lose_claim(job): 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" + 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 ): @@ -3231,9 +3321,7 @@ def check_retry_state(now): 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 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"}} @@ -3245,6 +3333,7 @@ def check_retry_state(now): 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 ): @@ -3254,7 +3343,12 @@ def test_managed_shutdown_cleans_staged_output_when_heartbeat_loses_claim( 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]) + 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 @@ -3270,7 +3364,9 @@ 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) + monkeypatch.setattr( + mock_service.client, "upload_collector_job_data", unexpected_upload + ) mock_service._shutdown() assert mock_service.future is None @@ -3377,7 +3473,9 @@ def end(job, outcome, **kwargs): mock_service._handle_completed_job(future) assert calls == [ - b"file-0", b"file-1", b"file-1", + b"file-0", + b"file-1", + b"file-1", (str(job_id), ManagedJobOutcome.SUCCEEDED), ] assert mock_service.future is None @@ -3485,7 +3583,9 @@ def test_managed_end_retry_does_not_upload_accepted_files_again( 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]) + result = Result( + results={}, job_id=job_id, output_path=output, data_files=[data_file] + ) future = Future() future.set_result(result) mock_service.managed = True @@ -3495,7 +3595,8 @@ def test_managed_end_retry_does_not_upload_accepted_files_again( uploads = [] ends = [] monkeypatch.setattr( - mock_service.client, "upload_collector_job_data", + mock_service.client, + "upload_collector_job_data", lambda job, data: uploads.append(data.read()), ) @@ -3527,14 +3628,17 @@ def test_managed_collection_failure_ends_failed_without_upload( mock_service.future = future ends = [] monkeypatch.setattr( - mock_service.client, "end_managed_job", + 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) + 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" @@ -3543,8 +3647,15 @@ def unexpected_upload(*args, **kwargs): def test_managed_collection_stages_then_uploads_before_job_end( - mock_service, mock_bloodhound_api + 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 @@ -3581,7 +3692,12 @@ def test_start_claimed_managed_job_passes_uuid_and_staging_mode_to_worker(mock_s mock_service.start_claimed_managed_job(job) args, _kwargs, future = mock_service.executor.submitted[0] assert args == ( - _subprocess_collect, "openhound-faker", str(job.id), None, None, True + _subprocess_collect, + "openhound-faker", + str(job.id), + None, + None, + True, ) assert mock_service.future is future assert mock_service.job_running == str(job.id) @@ -3602,10 +3718,13 @@ def test_start_claimed_managed_job_rejects_unclaimed_jobs(mock_service, managed) 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)) + 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", + scheduler_service.CollectorManager, + "from_entrypoint", lambda: SimpleNamespace(collectors=[collector]), ) 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 index 7b68a33..737b711 100644 --- a/tests/test_data/api/jobs/collector_jobs_available_with_jobs.json +++ b/tests/test_data/api/jobs/collector_jobs_available_with_jobs.json @@ -33,8 +33,8 @@ }, { "id": "22222222-2222-2222-2222-222222222222", - "job_schedule_id": 42, - "job_profile_id": 84, + "job_schedule_id": "44444444-4444-4444-4444-444444444444", + "job_profile_id": "55555555-5555-5555-5555-555555555555", "job_type_id": 8, "job_key": "openhound-faker", "params_version": "v3", From a9d36ce5e31c87c92f866923931ce3a928201cd5 Mon Sep 17 00:00:00 2001 From: Justin LaRose Date: Thu, 22 Oct 2026 13:16:15 -0400 Subject: [PATCH 30/34] chore: cleanup some unused code and stop OH from making calls that won't succeed when in managed mode. BED-100233 (#107) --- src/openhound/scheduler/service.py | 54 ++++------- tests/test_bhe_job_scheduling.py | 144 +++++++++++++---------------- 2 files changed, 79 insertions(+), 119 deletions(-) diff --git a/src/openhound/scheduler/service.py b/src/openhound/scheduler/service.py index 44e95ea..36e2a67 100644 --- a/src/openhound/scheduler/service.py +++ b/src/openhound/scheduler/service.py @@ -762,10 +762,11 @@ def _handle_completed_job(self, future: Future[Result]) -> None: 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) + 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: @@ -777,7 +778,9 @@ def _handle_completed_job(self, future: Future[Result]) -> None: break try: with data: - self.client.upload_collector_job_data(self.managed_job_id, 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) @@ -824,28 +827,6 @@ def _poll_running_job(self) -> None: except Exception: logger.exception("Error checking in-progress job.") - def start_claimed_managed_job(self, job: CollectorJob) -> None: - """Run a managed job after the queue caller has claimed it. - - Claiming and renewing the job belongs to the queue scheduling flow. - """ - if not self.managed: - raise ValueError("Managed collections are not enabled") - if self.job_running is not None or self.future is not None: - raise ValueError("A collector job is already running") - if job.status != "running" or job.claimed_by is None: - raise ValueError("The managed collector job must be claimed and running") - self.job_running = str(job.id) - self.managed_job_id = str(job.id) - self.next_heartbeat_at = time.monotonic() + HEARTBEAT_INTERVAL - try: - self.future = self.executor.submit( - _subprocess_collect, self.collector_name, str(job.id), None, None, True - ) - except BrokenProcessPool as error: - self.future = Future() - self.future.set_exception(error) - def _poll(self) -> None: """Checks if jobs are completed and if a job should be run.""" try: @@ -861,15 +842,14 @@ def _poll(self) -> None: "Error reporting completed job; retaining result for retry." ) - # Support bundles precede new collections in either mode. An unresolved - # managed job must be reconciled before accepting management work. + # 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 idle and self._poll_management(): + if not self.managed and idle and self._poll_management(): return if self.managed: @@ -896,13 +876,11 @@ 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: - if self.managed: - logger.error("Unable to update client metadata.") - else: + # 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: @@ -912,7 +890,7 @@ def start(self) -> None: self._poll() except ManagedRuntimeUnavailableError: logger.error( - "Managed collection runtime is unavailable; continuing support-bundle polling." + "Managed collection runtime is unavailable; retrying next poll." ) except ManagedJobProcessingError: logger.error( diff --git a/tests/test_bhe_job_scheduling.py b/tests/test_bhe_job_scheduling.py index f1fe226..fd7dbfe 100644 --- a/tests/test_bhe_job_scheduling.py +++ b/tests/test_bhe_job_scheduling.py @@ -405,6 +405,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 = {} @@ -1202,11 +1232,9 @@ def test_check_management_ignores_non_queued_operations( assert mock_service.check_management() is None -@pytest.mark.parametrize("managed", [False, True]) def test_poll_prioritizes_management_over_a_new_job( - mock_service, mock_bloodhound_api, monkeypatch, managed + mock_service, mock_bloodhound_api, monkeypatch ): - mock_service.managed = managed mock_bloodhound_api.app.state.management_operations = [_support_bundle_operation()] sent = [] monkeypatch.setattr(mock_service, "_send_support_bundle", sent.append) @@ -1247,15 +1275,15 @@ def fail_management(): assert state.managed_events == ["claim", "retrieve", "start_handoff"] else: assert state.job_started is True - assert "management unavailable" in caplog.text - assert any(record.exc_info is not None for record in caplog.records) + 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 + ) -@pytest.mark.parametrize("managed", [False, True]) def test_poll_does_not_start_a_job_when_management_work_fails( - mock_service, mock_bloodhound_api, monkeypatch, managed + mock_service, mock_bloodhound_api, monkeypatch ): - mock_service.managed = managed mock_bloodhound_api.app.state.management_operations = [_support_bundle_operation()] def fail(operation): @@ -1566,9 +1594,9 @@ def managed_service(mock_service, mock_bloodhound_api): "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" + 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", @@ -1806,7 +1834,7 @@ def flaky_request(**kwargs): 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 == 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 @@ -2320,7 +2348,7 @@ def next_poll(seconds): "stage,failure,expected", [ ("poll", "region", "AWS_DEFAULT_REGION"), - ("metadata", "provider", "custom-process"), + ("poll", "provider", "custom-process"), ("end", "document", "non-empty token_id"), ], ) @@ -2359,12 +2387,6 @@ def next_poll(seconds): 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 == "metadata": - managed_service.exit = True - else: - monkeypatch.setattr( - managed_service.client, "update_client_metadata", lambda: None - ) if stage == "end": managed_service._pending_managed_failure = ( "11111111-1111-1111-1111-111111111111" @@ -2385,12 +2407,7 @@ def next_poll(seconds): 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 == [] - if stage == "metadata": - assert "Unable to update client metadata" in caplog.text - else: - assert ( - f"retrying next poll in {managed_service.interval} seconds" in caplog.text - ) + 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 @@ -2561,23 +2578,19 @@ def test_malformed_committed_claim_response_keeps_id_until_reconciled( assert str(service.handoffs[0].job.id) == job_id -def test_managed_scheduler_keeps_serving_support_bundles_without_collection_runtime( - managed_service, mock_bloodhound_api, monkeypatch, tmp_path, caplog +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 - service.log_base_path = tmp_path - (tmp_path / "openhound.log").write_text("support log") + state.management_operations = [_support_bundle_operation()] polls = [] def next_poll(seconds): assert seconds == service.interval polls.append(seconds) - if len(polls) == 1: - # Support work arrives after a poll found no collection runtime. - state.management_operations = [_support_bundle_operation()] - else: + if len(polls) == 2: service.exit = True monkeypatch.setattr(scheduler_service.time, "sleep", next_poll) @@ -2585,10 +2598,12 @@ def next_poll(seconds): with caplog.at_level(logging.DEBUG): service.start() - assert len(polls) == state.management_available_requests == 2 - assert state.operation_started is True - assert state.artifact_completed is True - assert state.operation_end_payload["status"] == "succeeded" + 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 == [] ) @@ -2596,13 +2611,12 @@ def next_poll(seconds): assert state.job_started is False assert service.handoffs == [] assert ( - "Managed collection runtime is unavailable; continuing support-bundle polling." - in caplog.text + "Managed collection runtime is unavailable; retrying next poll." in caplog.text ) @pytest.mark.parametrize("pending", ["claim", "failure"]) -def test_managed_recovery_precedes_management_then_allows_support_work( +def test_managed_recovery_continues_without_legacy_management( managed_service, mock_bloodhound_api, monkeypatch, pending ): service = managed_service @@ -2627,8 +2641,9 @@ def test_managed_recovery_precedes_management_then_allows_support_work( ) service._poll() - assert state.management_available_requests == 1 - assert len(sent) == 1 + 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( @@ -3198,9 +3213,9 @@ def end(job, outcome, **kwargs): mock_service._poll() assert uploads == [b"file-0"] - assert ( - len(ends) == 1 - ), "Missing staged data must fail the job instead of retrying forever" + 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 @@ -3264,9 +3279,9 @@ def lose_claim(job): 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" + assert not output.exists(), ( + "Claim loss must also remove partially uploaded staged data" + ) @pytest.mark.usefixtures("enabled_managed_heartbeats") @@ -3682,39 +3697,6 @@ def test_managed_collection_stages_then_uploads_before_job_end( assert not result.output_path.exists() -def test_start_claimed_managed_job_passes_uuid_and_staging_mode_to_worker(mock_service): - job = CollectorJob.model_validate( - load_json("collector_jobs_available_with_job.json")["data"]["jobs"][0] - ) - job.status = "running" - job.claimed_by = UUID("c894ea20-8c11-49bd-83f8-5550c8255148") - mock_service.managed = True - mock_service.start_claimed_managed_job(job) - args, _kwargs, future = mock_service.executor.submitted[0] - assert args == ( - _subprocess_collect, - "openhound-faker", - str(job.id), - None, - None, - True, - ) - assert mock_service.future is future - assert mock_service.job_running == str(job.id) - - -@pytest.mark.parametrize("managed", [False, True]) -def test_start_claimed_managed_job_rejects_unclaimed_jobs(mock_service, managed): - job = CollectorJob.model_validate( - load_json("collector_jobs_available_with_job.json")["data"]["jobs"][0] - ) - mock_service.managed = managed - with pytest.raises(ValueError): - mock_service.start_claimed_managed_job(job) - assert mock_service.executor.submitted == [] - assert mock_service.job_running is None - - def test_failed_managed_conversion_cleans_up_staged_files(monkeypatch, tmp_path): output = tmp_path / "output" output.mkdir() From 2a797909621088772829960f07fae048216e1d60 Mon Sep 17 00:00:00 2001 From: Stran Dutton <65475509+StranDutton@users.noreply.github.com> Date: Thu, 22 Oct 2026 12:42:52 -0500 Subject: [PATCH 31/34] feat(managed): resolve nested job configuration - BED-10080 (#105) Resolve job params and retrieved secrets through a job-scoped DLT provider before collector loading. Align the job model with BHE's params_version-to-version rename (BED-10069, BED-10180). --- src/openhound/core/clients/models/jobs.py | 2 +- src/openhound/scheduler/configuration.py | 71 ++++ src/openhound/scheduler/service.py | 87 +++-- tests/test_bhe_job_scheduling.py | 272 ++++++++++++-- .../collector_jobs_available_with_job.json | 14 +- .../collector_jobs_available_with_jobs.json | 18 +- .../test_data/managed_aws_configuration.json | 21 ++ tests/test_data/managed_aws_secrets.json | 11 + tests/test_managed_configuration.py | 331 ++++++++++++++++++ 9 files changed, 746 insertions(+), 81 deletions(-) create mode 100644 src/openhound/scheduler/configuration.py create mode 100644 tests/test_data/managed_aws_configuration.json create mode 100644 tests/test_data/managed_aws_secrets.json create mode 100644 tests/test_managed_configuration.py diff --git a/src/openhound/core/clients/models/jobs.py b/src/openhound/core/clients/models/jobs.py index fa301da..5244ca3 100644 --- a/src/openhound/core/clients/models/jobs.py +++ b/src/openhound/core/clients/models/jobs.py @@ -62,7 +62,7 @@ class CollectorJob(BaseModel): job_profile_id: UUID | None job_type_id: int job_key: str - params_version: str + version: str params: dict[str, Any] scope_client_id: UUID | None secret_key_id: str | None diff --git a/src/openhound/scheduler/configuration.py b/src/openhound/scheduler/configuration.py new file mode 100644 index 0000000..912195e --- /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/service.py b/src/openhound/scheduler/service.py index 36e2a67..2391940 100644 --- a/src/openhound/scheduler/service.py +++ b/src/openhound/scheduler/service.py @@ -1,7 +1,5 @@ -import json import logging import math -import os import shutil import signal import tempfile @@ -9,6 +7,7 @@ from collections.abc import Callable from concurrent.futures import Future, ProcessPoolExecutor from concurrent.futures.process import BrokenProcessPool +from contextlib import nullcontext from dataclasses import dataclass, field from datetime import UTC, datetime from importlib import import_module @@ -40,6 +39,10 @@ from openhound.core.manager import CollectorManager from openhound.core.support_bundle import create_support_bundle from openhound.scheduler import dataflow +from openhound.scheduler.configuration import ( + ManagedConfigurationError, + managed_job_configuration, +) logger = logging.getLogger(__name__) @@ -119,29 +122,6 @@ class Result: upload_failure: str | None = None -def _set_managed_source_environment( - collector, params: dict[str, Any], secrets: dict[str, Any] -) -> None: - """Provide one managed job's validated values to its DLT source.""" - source = collector.dlt_source - if source is None: - raise ValueError(f"Collector '{collector.name}' has no registered DLT source") - - section = source.section or source.__module__.rsplit(".", 1)[-1] - source_name = source.name or source.__name__ - prefix = f"SOURCES__{section.upper()}__{source_name.upper()}" - for name, value in params.items(): - if name.lower() == "credentials": - raise ValueError("Managed job parameters cannot contain credentials") - os.environ[f"{prefix}__{name.upper()}"] = ( - value if isinstance(value, str) else json.dumps(value) - ) - for name, value in secrets.items(): - os.environ[f"{prefix}__CREDENTIALS__{name.upper()}"] = ( - value if isinstance(value, str) else json.dumps(value) - ) - - def _subprocess_collect( collector_name: str, job_id: int | str, @@ -165,19 +145,35 @@ def _subprocess_collect( ExtensionNotFoundError: If the collector cannot be found via entrypoints. """ signal.signal(signal.SIGINT, signal.SIG_IGN) + 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: + 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 params is not None or secrets is not None: - if ( - collector.dlt_source is None - and collector.name in MANAGED_SOURCE_MODULES - ): - collector.dlt_source = import_module( - MANAGED_SOURCE_MODULES[collector.name] - ).source - _set_managed_source_environment(collector, params or {}, secrets or {}) + 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": ( @@ -318,7 +314,8 @@ def _finish_managed_job_on_shutdown(self) -> None: return try: self._handle_completed_job(self.future) - except Exception: # noqa: BLE001 - retain managed work and redact sensitive errors + except Exception: # noqa: BLE001 + # Retain managed work and redact sensitive errors. logger.error( "Error reporting completed managed job during shutdown; " "retaining result for retry." @@ -408,7 +405,8 @@ def _load_managed_credentials(self, job: CollectorJob) -> dict[str, Any]: return credentials except SecretRetrievalError as error: raise CredentialValidationError(str(error), error.failure_summary) from None - except Exception as error: # noqa: BLE001 - unexpected readers can expose payloads + 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", @@ -682,7 +680,8 @@ def _send_managed_heartbeat(self, job_id: str) -> None: try: self.client.heartbeat_collector_job(job_id) - except Exception as error: # noqa: BLE001 - transient heartbeat failures must not stop collection + 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 @@ -741,6 +740,18 @@ def _handle_completed_job(self, future: Future[Result]) -> None: ) 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." diff --git a/tests/test_bhe_job_scheduling.py b/tests/test_bhe_job_scheduling.py index fd7dbfe..2ca23ab 100644 --- a/tests/test_bhe_job_scheduling.py +++ b/tests/test_bhe_job_scheduling.py @@ -15,6 +15,7 @@ from urllib.parse import urlsplit from uuid import UUID +import dlt import pytest import requests from botocore.exceptions import ClientError @@ -41,6 +42,7 @@ ) from openhound.core.models.graph import Graph from openhound.scheduler import service as scheduler_service +from openhound.scheduler.configuration import ManagedConfigurationError from openhound.scheduler.service import ( ExtensionNotFoundError, ManagedJobProcessingError, @@ -52,6 +54,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: @@ -872,11 +880,17 @@ def test_collector_job_queue_response_parses_contract_fields( assert job.job_type_id == 7 assert isinstance(job.job_type_id, int) assert job.job_key == "openhound-faker" - assert job.params_version == "v2" + assert job.version == "v2" assert job.params == { - "domains": ["example.com", "example.org"], - "include_deleted": False, - "nested": {"depth": 2}, + "sources": { + "source": { + "random": { + "domains": ["example.com", "example.org"], + "include_deleted": False, + "nested": {"depth": 2}, + } + } + } } assert job.scope_client_id is None assert job.secret_key_id is None @@ -1601,7 +1615,7 @@ def managed_service(mock_service, mock_bloodhound_api): state.collector_job_claim_response.update( secret_key_id="private-secret-reference", status="claimed", - params={"claimed_parameter": True}, + params=CLAIMED_PARAMS, claimed_by="33333333-3333-3333-3333-333333333333", claimed_at=claimed_at.isoformat(), claim_expires_at=(claimed_at + timedelta(minutes=15)).isoformat(), @@ -1613,7 +1627,7 @@ def managed_service(mock_service, mock_bloodhound_api): def get_secret_value(*, SecretId): state.managed_events.append("retrieve") mock_service.secret_requests.append(SecretId) - return {"SecretString": '{"credential":"private-credential-value"}'} + return {"SecretString": json.dumps(COLLECTOR_SECRETS)} mock_service.secrets_manager = AWSSecretsManager( SimpleNamespace(get_secret_value=get_secret_value) @@ -1660,6 +1674,8 @@ def request(**kwargs): 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", @@ -1685,9 +1701,9 @@ def test_managed_poll_claims_retrieves_and_hands_off_authoritative_job( 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_parameter": True} + assert prepared.job.params == CLAIMED_PARAMS assert prepared.job.status == "claimed" - assert prepared.credentials == {"credential": "private-credential-value"} + 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 == [] @@ -1707,6 +1723,74 @@ def test_managed_poll_claims_retrieves_and_hands_off_authoritative_job( 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) + + 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 + + @pytest.mark.parametrize( "association_index", [0, 1], ids=["null-associations", "uuid-associations"] ) @@ -1745,8 +1829,8 @@ def submit(*args): scheduler_service._subprocess_collect, "openhound-faker", job_id, - {"claimed_parameter": True}, - {"credential": "private-credential-value"}, + CLAIMED_PARAMS, + COLLECTOR_SECRETS, True, ) ] @@ -1844,8 +1928,8 @@ def flaky_request(**kwargs): assert args[1:] == ( "openhound-faker", job_id, - {"claimed_parameter": True}, - {"credential": "private-credential-value"}, + CLAIMED_PARAMS, + COLLECTOR_SECRETS, True, ) assert service.future is future @@ -2285,7 +2369,11 @@ def fail_start(prepared): def test_valid_generic_nested_json_is_handed_off_without_field_schema(managed_service): - credentials = {"arbitrary_field": {"nested": ["value", False, 0, 1.5]}} + 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)} @@ -2456,7 +2544,7 @@ def test_committed_claim_with_lost_responses_is_reconciled_before_discovery( assert service.secret_requests == ["private-secret-reference"] (prepared,) = service.handoffs assert str(prepared.job.id) == job_id - assert prepared.job.params == {"claimed_parameter": True} + 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 @@ -2681,26 +2769,28 @@ def expired_token(self, **kwargs): assert context in formatted_error -def test_managed_worker_sets_dlt_environment_before_pipeline(monkeypatch, caplog): +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, ) - worker_environment = {} - observed_environment = {} + observed_configuration = {} monkeypatch.setattr( scheduler_service.CollectorManager, "from_entrypoint", lambda: SimpleNamespace(collectors=[collector]), ) - monkeypatch.setattr( - scheduler_service, "os", SimpleNamespace(environ=worker_environment) - ) def capture_pipeline(extension): - observed_environment.update(worker_environment) + 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) @@ -2709,24 +2799,135 @@ def capture_pipeline(extension): result = _subprocess_collect( "faker", "collector-job-123", - {"node_count": 4, "filters": ["one", "two"]}, - {"token": "secret-value"}, + { + "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_environment == { - "SOURCES__SOURCE__RANDOM__NODE_COUNT": "4", - "SOURCES__SOURCE__RANDOM__FILTERS": '["one", "two"]', - "SOURCES__SOURCE__RANDOM__CREDENTIALS__TOKEN": "secret-value", + assert observed_configuration == { + "node_count": 4, + "filters": ["one", "two"], + "token": "secret-value", } assert "secret-value" not in caplog.text -def test_managed_environment_requires_a_registered_dlt_source(): - collector = SimpleNamespace(name="example", dlt_source=None) +@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) - with pytest.raises(ValueError, match="has no registered DLT source"): - scheduler_service._set_managed_source_environment(collector, {}, {}) + (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 + ] + 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): @@ -2739,9 +2940,16 @@ def inspect_collection(extension): dataflow.pipeline = inspect_collection result = _subprocess_collect( - "faker", "e0d43dbd-2cdf-4d68-9a81-2f95cd58a92f", {"node_count": 4}, {} + "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 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 index 7e59a5c..3e66f3b 100644 --- a/tests/test_data/api/jobs/collector_jobs_available_with_job.json +++ b/tests/test_data/api/jobs/collector_jobs_available_with_job.json @@ -10,11 +10,17 @@ "job_profile_id": null, "job_type_id": 7, "job_key": "openhound-faker", - "params_version": "v2", + "version": "v2", "params": { - "domains": ["example.com", "example.org"], - "include_deleted": false, - "nested": {"depth": 2} + "sources": { + "source": { + "random": { + "domains": ["example.com", "example.org"], + "include_deleted": false, + "nested": {"depth": 2} + } + } + } }, "scope_client_id": null, "secret_key_id": null, 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 index 737b711..c51be26 100644 --- a/tests/test_data/api/jobs/collector_jobs_available_with_jobs.json +++ b/tests/test_data/api/jobs/collector_jobs_available_with_jobs.json @@ -10,11 +10,17 @@ "job_profile_id": null, "job_type_id": 7, "job_key": "openhound-faker", - "params_version": "v2", + "version": "v2", "params": { - "domains": ["example.com", "example.org"], - "include_deleted": false, - "nested": {"depth": 2} + "sources": { + "source": { + "random": { + "domains": ["example.com", "example.org"], + "include_deleted": false, + "nested": {"depth": 2} + } + } + } }, "scope_client_id": null, "secret_key_id": null, @@ -37,8 +43,8 @@ "job_profile_id": "55555555-5555-5555-5555-555555555555", "job_type_id": 8, "job_key": "openhound-faker", - "params_version": "v3", - "params": {"mode": "full"}, + "version": "v3", + "params": {"sources": {"source": {"random": {"mode": "full"}}}}, "scope_client_id": "33333333-3333-3333-3333-333333333333", "secret_key_id": "secret-1", "priority": 5, diff --git a/tests/test_data/managed_aws_configuration.json b/tests/test_data/managed_aws_configuration.json new file mode 100644 index 0000000..411b72b --- /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 0000000..20a8ec0 --- /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_managed_configuration.py b/tests/test_managed_configuration.py new file mode 100644 index 0000000..3ec429f --- /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"] From 13ab19a4b850a2c0c4a73e6fdac517b3d2a19ad4 Mon Sep 17 00:00:00 2001 From: Stran Dutton <65475509+StranDutton@users.noreply.github.com> Date: Thu, 22 Oct 2026 15:06:44 -0500 Subject: [PATCH 32/34] feat(managed): add job-scoped AWS profile adapter - BED-10231 (#106) * feat(managed): add job-scoped AWS profile adapter - BED-10231 * refactor(managed): centralize AWS preparation and preserve redacted diagnostics - BED-10231 * fix(logging): format original arguments before redaction - BED-10231 * fix(managed): preserve safe AWS preparation diagnostics - BED-10231 * fix(managed): log safe AWS worker failure diagnostics - BED-10231 --- .github/workflows/test.yml | 4 + pyproject.toml | 1 + .../core/clients/aws_secrets_manager.py | 138 +++- src/openhound/core/clients/models/jobs.py | 1 + src/openhound/core/logging.py | 92 ++- src/openhound/scheduler/aws_profile.py | 432 ++++++++++++ src/openhound/scheduler/service.py | 112 ++- tests/conftest.py | 62 ++ tests/test_aws_profile_adapter.py | 661 ++++++++++++++++++ tests/test_bhe_job_scheduling.py | 579 +++++++++++++++ .../collector_jobs_available_with_job.json | 1 + .../collector_jobs_available_with_jobs.json | 2 + tests/test_log_handlers.py | 107 +++ uv.lock | 2 + 14 files changed, 2162 insertions(+), 32 deletions(-) create mode 100644 src/openhound/scheduler/aws_profile.py create mode 100644 tests/test_aws_profile_adapter.py diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index bdc53f0..37a727a 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -46,6 +46,10 @@ jobs: 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/pyproject.toml b/pyproject.toml index 79ee26d..db615b6 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -20,6 +20,7 @@ dependencies = [ "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/core/clients/aws_secrets_manager.py b/src/openhound/core/clients/aws_secrets_manager.py index 27ece94..7f1f7a5 100644 --- a/src/openhound/core/clients/aws_secrets_manager.py +++ b/src/openhound/core/clients/aws_secrets_manager.py @@ -87,6 +87,8 @@ class SecretRetrievalError(Exception): 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): @@ -358,14 +360,131 @@ def _classify_aws_error( context["region"] = region location += f" in {region}" if isinstance(error, RuntimeError): - return _classify_runtime_error(error, context) - if isinstance(error, (FileNotFoundError, PermissionError)): - return _classify_file_error(error, context, location) + 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): - return _classify_client_error(error, context, location) - if isinstance(error, _PROVIDER_ERROR_TYPES): - return _classify_provider_error(error, context) - return _classify_transport_error(error, context, location) + 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( @@ -633,6 +752,7 @@ def _log_failure( } if batch_number is not None: extra["batch_number"] = batch_number - extra.update(error.context) + 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, error, extra=extra) + logger.debug("AWS %s failed: %s", operation, message, extra=extra) diff --git a/src/openhound/core/clients/models/jobs.py b/src/openhound/core/clients/models/jobs.py index 5244ca3..2da85e7 100644 --- a/src/openhound/core/clients/models/jobs.py +++ b/src/openhound/core/clients/models/jobs.py @@ -65,6 +65,7 @@ class CollectorJob(BaseModel): 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 diff --git a/src/openhound/core/logging.py b/src/openhound/core/logging.py index 7e932db..2a6ae8d 100644 --- a/src/openhound/core/logging.py +++ b/src/openhound/core/logging.py @@ -8,6 +8,9 @@ import re import sys import time +import traceback +from collections.abc import Iterator +from contextlib import contextmanager from enum import Enum from logging.handlers import TimedRotatingFileHandler from pathlib import Path @@ -15,10 +18,81 @@ import dlt 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 = values + 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] = [] + + +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"] VALID_FORMATS = ["json", "text"] @@ -147,7 +221,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(): @@ -195,7 +271,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 @@ -214,6 +292,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 @@ -411,7 +493,7 @@ def _get_file_handler(self, file_path: Path) -> "RotatingFileHandler": # 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 handler + 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""" @@ -422,7 +504,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)) @@ -442,7 +524,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 0000000..2434517 --- /dev/null +++ b/src/openhound/scheduler/aws_profile.py @@ -0,0 +1,432 @@ +"""Adapt BHE's raw IAM access-key pair to the AWS extension's named profile. + +Only the worker installs AWS/DLT configuration or writes credential material. +The scheduler owns the workspace and its lifetime, including killed workers. +""" + +import logging +import os +import re +import shutil +import stat +import tempfile +import traceback +from collections.abc import Callable, Iterator +from configparser import ConfigParser +from contextlib import contextmanager +from copy import deepcopy +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any +from uuid import uuid4 + +from botocore.exceptions import BotoCoreError, ClientError +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__) + +_CREDENTIAL_SETTINGS = { + "region_name", + "endpoint_url", + "org_role_name", + "org_role_session_name", + "include_management_account", + "source_account", + "management_account_id", + "management_role_name", +} + + +class AWSAccessKeyError(ManagedConfigurationError): + """An AWS preparation error containing no supplied credential values.""" + + def __init__( + self, + message: str, + summary: str = "Managed AWS credentials or authentication settings are invalid", + **context: Any, + ) -> None: + super().__init__(message) + self.summary = summary + self.context = context + + +def validate_access_key_id(value: Any) -> None: + # ASIA keys are temporary and require a session token, absent in BHE's beta + # contract. Do not strip input: INI/SDK whitespace normalization would change it. + if not isinstance(value, str) or re.fullmatch(r"AKIA[A-Z0-9]{16}", value) is None: + raise AWSAccessKeyError( + "Managed AWS job requires a long-lived IAM Access Key ID in external_secret_key_id" + ) + + +def validate_secret_reference(value: Any) -> str: + if ( + not isinstance(value, str) + or re.fullmatch(r"[A-Za-z0-9/_+=.@:-]{1,2048}", value) is None + ): + 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: + validate_access_key_id(self.access_key_id) + if ( + not isinstance(self.secret_access_key, str) + or re.fullmatch(r"[A-Za-z0-9/+=]{40}", self.secret_access_key) is None + ): + raise AWSAccessKeyError( + "Managed AWS credential secret must be the raw 40-character Secret Access Key" + ) + + +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 valid 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: + validate_access_key_id(job.external_secret_key_id) + 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 raw-key mode and copy its full nested settings before /start.""" + 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) + _validate_credential_sections(document) + 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) or credentials.keys() - _CREDENTIAL_SETTINGS: + raise AWSAccessKeyError( + "Managed AWS access-key jobs cannot select another profile or authentication mode" + ) + if any( + key in node + for key in ("profile_name", "profile_names", "auth_mode", "authentication_mode") + ): + raise AWSAccessKeyError( + "Managed AWS access-key jobs contain conflicting authentication settings" + ) + # An incomplete organization configuration must not fall through to the single + # profile model and silently ignore account/role settings. + organization = _CREDENTIAL_SETTINGS - {"region_name", "endpoint_url"} + if credentials.keys() & organization and ( + not isinstance(credentials.get("org_role_name"), str) + or not credentials["org_role_name"].strip() + ): + raise AWSAccessKeyError( + "Managed AWS organization settings require org_role_name" + ) + return document + + +def _validate_credential_sections( + document: dict[str, Any], path: tuple[str, ...] = () +) -> None: + # DLT accepts compact, less-specific and pipeline-prefixed layouts too. + # Keep every source-credential setting in the canonical raw-key section so + # another union alternative cannot win during pipeline-scoped resolution. + for key, value in document.items(): + if ( + key == "credentials" + and path != ("sources", "source", "aws") + and "destination" not in path + ): + raise AWSAccessKeyError( + "Managed AWS credential settings must use sources.source.aws.credentials" + ) + if isinstance(value, dict): + _validate_credential_sections(value, (*path, key)) + + +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]: + if not self.provider.supports_sections and key in _CREDENTIAL_SETTINGS | { + "profile_name", + "profile_names", + "credentials", + }: + 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, str]]) -> None: + config = ConfigParser(interpolation=None) + config.read_dict(sections) + # 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: + config.write(stream) + + +@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") + + +@contextmanager +def managed_aws_profile( + params: dict[str, Any], credentials: AWSAccessKeyCredentials, directory: Path +) -> Iterator[str]: + """Install one profile before loading collectors; restore all worker state.""" + document = aws_profile_params(params) + profile = f"openhound-job-{uuid4().hex}" + document["sources"]["source"]["aws"]["credentials"]["profile_name"] = profile + 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"], + { + 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}": {}}) + 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/service.py b/src/openhound/scheduler/service.py index 2391940..2f61ce6 100644 --- a/src/openhound/scheduler/service.py +++ b/src/openhound/scheduler/service.py @@ -39,6 +39,15 @@ 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, +) from openhound.scheduler.configuration import ( ManagedConfigurationError, managed_job_configuration, @@ -72,9 +81,15 @@ class ManagedRuntimeUnavailableError(RuntimeError): 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"): + 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): @@ -85,11 +100,11 @@ class ManagedJobProcessingError(RuntimeError): class PreparedManagedJob: """Validated job and credentials for the managed collection runtime. - Keep credentials in memory; never serialize or log them. + Keep credentials in the in-memory worker handoff; never persist or log them. """ job: CollectorJob = field(repr=False) - credentials: dict[str, Any] = field(repr=False) + credentials: dict[str, Any] | AWSAccessKeyCredentials = field(repr=False) def _usable_credential_value(value: Any) -> bool: @@ -126,8 +141,9 @@ def _subprocess_collect( collector_name: str, job_id: int | str, params: dict[str, Any] | None = None, - secrets: 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. @@ -145,6 +161,13 @@ def _subprocess_collect( ExtensionNotFoundError: If the collector cannot be found via entrypoints. """ signal.signal(signal.SIGINT, signal.SIG_IGN) + 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): + 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 {}) @@ -163,9 +186,7 @@ def _collect_in_worker( available_collectors = CollectorManager.from_entrypoint() for collector in available_collectors.collectors: - if ( - collector.name == collector_name - ): # pyright: ignore[reportAttributeAccessIssue] + if collector.name == collector_name: # pyright: ignore[reportAttributeAccessIssue] if ( configured and collector.dlt_source is None @@ -261,6 +282,9 @@ def __init__( # 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 @@ -324,11 +348,17 @@ def _finish_managed_job_on_shutdown(self) -> 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: @@ -384,13 +414,10 @@ def _report_managed_credential_failure(self, job_id: str) -> None: ) self._pending_managed_failure = None - def _load_managed_credentials(self, job: CollectorJob) -> dict[str, Any]: + def _load_managed_credentials( + self, job: CollectorJob + ) -> dict[str, Any] | AWSAccessKeyCredentials: """Retrieve and validate credentials, exposing only a safe failure message.""" - if job.secret_key_id is None or not job.secret_key_id.strip(): - raise CredentialValidationError( - "Managed job has no credential secret identifier" - ) - logger.debug( "Retrieving managed collector credentials.", extra={ @@ -400,9 +427,26 @@ def _load_managed_credentials(self, job: CollectorJob) -> dict[str, Any]: }, ) try: - credentials = self.secrets_manager.get_secret(job.secret_key_id) + 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) and _usable_credential_value(credentials): return credentials + except AWSAccessKeyError as error: + raise CredentialValidationError( + str(error), error.summary, **error.context + ) from None + except CredentialValidationError: + raise except SecretRetrievalError as error: raise CredentialValidationError(str(error), error.failure_summary) from None except Exception as error: # noqa: BLE001 @@ -466,7 +510,13 @@ def _claim_and_validate_managed_job(self, job_id: str) -> PreparedManagedJob | N "Cannot prepare managed job %s: %s. Reporting this job failed.", job_id, str(error).rstrip("."), - extra={**context, "secret_id": job.secret_key_id}, + 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 @@ -582,7 +632,10 @@ def _start_job(self, job: Job) -> None: self._launch_job(job.id) def run_claimed_job( - self, job_id: str | UUID, params: dict[str, Any], secrets: dict[str, Any] + 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 @@ -594,13 +647,17 @@ def _launch_job( job_id: int | str, *, params: dict[str, Any] | None = None, - secrets: 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: @@ -627,13 +684,29 @@ def _launch_job( 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, dict(params or {}) if managed else None, - dict(secrets 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("Failed to submit job %s: process pool is broken.", job_id) @@ -730,6 +803,9 @@ def _abort_managed_job(self, job_id: str) -> None: 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() diff --git a/tests/conftest.py b/tests/conftest.py index ee43c1d..99708ea 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,6 +1,12 @@ import importlib import importlib.util +import os +import sys from pathlib import Path +from typing import Union + +import dlt +import pytest def import_from_path(file_path: Path) -> None: @@ -16,3 +22,59 @@ def pytest_configure(config): py_files = [path for path in sources_root.rglob("**/models/*.py") if path.is_file()] for module in py_files: import_from_path(module) + + +@pytest.fixture(scope="session") +def aws_extension_auth(): + """Load the real extension models/auth, without its unrelated native engine. + + Use the sibling repo, OPENHOUND_AWS_SOURCE, or an installed extension. These + model integration tests are optional when the internal extension is absent. + """ + default = Path(__file__).resolve().parents[2] / "openhound-aws" + repo = Path(os.environ.get("OPENHOUND_AWS_SOURCE", default)) + path = repo / "src" / "openhound_aws" / "auth.py" + if not path.is_file(): + if "OPENHOUND_AWS_SOURCE" in os.environ: + pytest.fail("OPENHOUND_AWS_SOURCE must point to an openhound-aws checkout") + return pytest.importorskip( + "openhound_aws.auth", + reason="AWS model integration requires openhound-aws or OPENHOUND_AWS_SOURCE", + ) + name = "_openhound_aws_auth_integration" + spec = importlib.util.spec_from_file_location(name, path) + module = importlib.util.module_from_spec(spec) + sys.modules[name] = module + spec.loader.exec_module(module) + return module + + +@pytest.fixture +def resolve_aws_settings(aws_extension_auth): + auth = aws_extension_auth + observed = {} + + # Use the real model union and DLT source resolution used by the extension. + # Resolving the configuration alone does not need its Rust policy evaluator. + @dlt.source(name="aws", section="source") + def source( + credentials: Union[ # noqa: UP007 - match the extension source's union + auth.AWSOrganizationProfileCredentials, + auth.AWSMultiProfileCredentials, + auth.AWSProfileCredentials, + ] = dlt.secrets.value, + bucket_policy: auth.BucketPolicyConfiguration = dlt.config.value, + max_attempts: int = 3, + ): + observed.update( + credentials=credentials, + bucket_policy=bucket_policy, + max_attempts=max_attempts, + ) + return dlt.resource([], name="configuration_only") + + def resolve(): + source() + return observed.copy() + + return resolve diff --git a/tests/test_aws_profile_adapter.py b/tests/test_aws_profile_adapter.py new file mode 100644 index 0000000..f1f4745 --- /dev/null +++ b/tests/test_aws_profile_adapter.py @@ -0,0 +1,661 @@ +import logging +import multiprocessing +import os +import stat +import tempfile +import time +import traceback +import zipfile +from concurrent.futures import ProcessPoolExecutor +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.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 openhound.core import logging as openhound_logging +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" +MEMBER_ACCOUNT = "222222222222" + + +@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() + + +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 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_partial_profile_write_failure_restores_worker_state( + 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) + with ( + pytest.raises(ManagedConfigurationError) as error, + managed_aws_profile({}, 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_real_sdk_and_extension_profile_resolution( + workspace, aws_extension_auth, resolve_aws_settings +): + 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 = resolve_aws_settings() + model = settings["credentials"] + assert type(model) is aws_extension_auth.AWSProfileCredentials + assert model.profile_name == profile + assert model.region_name == "eu-west-1" + assert model.endpoint_url == "https://aws.example.test" + _assert_source_credentials(model.session) + client = model.session.client("sts", endpoint_url=model.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 + bucket = settings["bucket_policy"] + assert bucket.include is True and bucket.max_objects == 9 + assert bucket.bucket_whitelist[0].match == "public/.*" + 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, resolve_aws_settings, aws_extension_auth, 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: + model = resolve_aws_settings()["credentials"] + assert type(model) is aws_extension_auth.AWSProfileCredentials + assert model.profile_name == profile + _assert_source_credentials(model.session) + 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( + "dedicated_source", [False, True], ids=["management-profile", "source-account"] +) +def test_organization_models_and_assume_role_use_generated_source_identity( + dedicated_source, workspace, aws_extension_auth, resolve_aws_settings, monkeypatch +): + 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", + ) + params = {"sources": {"source": {"aws": {"credentials": options}}}} + clients = [] + sessions = [] + original_client = boto3.session.Session.client + + def stub_sts(session, name, **kwargs): + assert name == "sts" + _assert_source_credentials(session) + sessions.append(session) + client = original_client(session, name, **kwargs) + stubber = Stubber(client) + if not clients: + stubber.add_response("get_caller_identity", _identity(), {}) + else: + account, role = ( + (MANAGEMENT_ACCOUNT, "ManagementCollectorRole") + if dedicated_source and len(clients) == 1 + else (MEMBER_ACCOUNT, "MemberCollectorRole") + ) + stubber.add_response( + "assume_role", + { + "Credentials": { + "AccessKeyId": "ASIAIOSFODNN7EXAMPLE", + "SecretAccessKey": "T" * 40, + "SessionToken": "assumed-session-token", + "Expiration": datetime.now(UTC) + timedelta(hours=1), + } + }, + { + "RoleArn": f"arn:aws:iam::{account}:role/{role}", + "RoleSessionName": "job-session", + }, + ) + stubber.activate() + clients.append(stubber) + return client + + monkeypatch.setattr(boto3.session.Session, "client", stub_sts) + try: + with managed_aws_profile(params, PAIR, workspace.directory) as profile: + model = resolve_aws_settings()["credentials"] + assert type(model) is aws_extension_auth.AWSOrganizationProfileCredentials + assert model.profile_name == profile + assert model.org_role_name == "MemberCollectorRole" + assert model.collect_management_account is False + auth = aws_extension_auth.AWSAuth(model) + assert auth.source_ctx.account_id == SOURCE_ACCOUNT + if dedicated_source: + management = auth.management_ctx + assert management.account_id == MANAGEMENT_ACCOUNT + assert ( + management.session.get_credentials().token + == "assumed-session-token" + ) + member = auth.context_for_account(MEMBER_ACCOUNT) + assert member.account_id == MEMBER_ACCOUNT + assert member.session.get_credentials().token == "assumed-session-token" + # Both management and member assumptions originate from the raw pair, + # rather than chaining through management's temporary credentials. + assert all(session is auth.source_ctx.session for session in sessions) + for stubber in clients: + stubber.assert_no_pending_responses() + finally: + for stubber in clients: + stubber.deactivate() + + +@pytest.mark.parametrize("failed", [False, True], ids=["success", "failure"]) +def test_successive_jobs_restore_environment_providers_and_remove_files( + failed, resolve_aws_settings, 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 {} + ) + try: + with managed_aws_profile(options, PAIR, workspace.directory) as profile: + profiles.append(profile) + model = resolve_aws_settings()["credentials"] + assert model.region_name == ("eu-west-1" if index == 0 else "us-east-1") + _assert_source_credentials(model.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._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_bhe_job_scheduling.py b/tests/test_bhe_job_scheduling.py index 2ca23ab..02c06c1 100644 --- a/tests/test_bhe_job_scheduling.py +++ b/tests/test_bhe_job_scheduling.py @@ -3,6 +3,7 @@ import hashlib import json import logging +import os import subprocess import sys import time @@ -15,10 +16,12 @@ from urllib.parse import urlsplit from uuid import UUID +import boto3 import dlt import pytest import requests from botocore.exceptions import ClientError +from botocore.stub import Stubber from fastapi import FastAPI, Request, Response from fastapi.testclient import TestClient @@ -42,6 +45,7 @@ ) 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, @@ -893,6 +897,7 @@ def test_collector_job_queue_response_parses_contract_fields( } } 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" @@ -908,6 +913,7 @@ def test_collector_job_queue_response_parses_contract_fields( 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" @@ -1791,6 +1797,579 @@ def submit(function, *args): assert credential not in caplog.text +AWS_ACCESS_KEY = "AKIAIOSFODNN7EXAMPLE" +AWS_SECRET_KEY = "wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY" + + +@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, + resolve_aws_settings, + 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 + + def load_collectors(): + assert state.collector_job_start_requests == [ + (state.collector_job_claim_response["id"], b"") + ] + model = resolve_aws_settings()["credentials"] + assert model.profile_name.startswith("openhound-job-") + assert model.region_name == "us-west-2" + credentials = model.session.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) + + +@pytest.mark.parametrize( + "field,value", + [ + ("external_secret_key_id", None), + ("external_secret_key_id", ""), + ("external_secret_key_id", "ASIAIOSFODNN7EXAMPLE"), + ("external_secret_key_id", AWS_ACCESS_KEY + "\n"), + ("external_secret_key_id", "not-an-access-key"), + ("secret_key_id", None), + ("secret_key_id", ""), + ("secret_key_id", " "), + ("secret_key_id", " private-secret-reference"), + ("secret_key_id", "reference\nother"), + ], +) +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": ""}, + {"SecretString": " " + AWS_SECRET_KEY}, + {"SecretString": AWS_SECRET_KEY + "\n"}, + {"SecretString": AWS_SECRET_KEY[:-1]}, + {"SecretString": "x" * 41}, + {"SecretString": "%" * 40}, + {"SecretString": json.dumps(AWS_SECRET_KEY)}, + {"SecretString": '["key"]'}, + {"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}, + ], +) +def test_raw_aws_authentication_conflicts_fail_before_start( + 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 _: pytest.fail("invalid authentication must fail before retrieval"), + ) + aws_managed_service._poll() + assert state.collector_job_start_requests == [] + assert aws_managed_service.executor.submitted == [] + 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": {"profile_name": "local"}}}} + } + ) + 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"]}}}}, + {"profile_name": "flat-profile"}, + {"aws_collector": {"credentials": {"profile_names": ["other"]}}}, + { + "aws_collector": { + "sources": { + "source": {"aws": {"credentials": {"profile_name": "other"}}} + } + } + }, + ], +) +def test_raw_aws_rejects_ambiguous_authentication_layouts_before_start( + 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 state.collector_job_start_requests == [] + assert aws_managed_service.executor.submitted == [] + assert state.collector_job_end_requests[0]["body"]["outcome"] == "failed" + + @pytest.mark.parametrize( "association_index", [0, 1], ids=["null-associations", "uuid-associations"] ) 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 index 3e66f3b..0c3610c 100644 --- a/tests/test_data/api/jobs/collector_jobs_available_with_job.json +++ b/tests/test_data/api/jobs/collector_jobs_available_with_job.json @@ -23,6 +23,7 @@ } }, "scope_client_id": null, + "external_secret_key_id": null, "secret_key_id": null, "priority": 10, "status": "ready", 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 index c51be26..e1ad94f 100644 --- a/tests/test_data/api/jobs/collector_jobs_available_with_jobs.json +++ b/tests/test_data/api/jobs/collector_jobs_available_with_jobs.json @@ -23,6 +23,7 @@ } }, "scope_client_id": null, + "external_secret_key_id": null, "secret_key_id": null, "priority": 10, "status": "ready", @@ -46,6 +47,7 @@ "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", diff --git a/tests/test_log_handlers.py b/tests/test_log_handlers.py index f89dae9..454a6ba 100644 --- a/tests/test_log_handlers.py +++ b/tests/test_log_handlers.py @@ -1,16 +1,24 @@ import json import logging +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, + _ValueRedactor, + redact_log_values, ) @@ -293,6 +301,105 @@ 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} + assert secret not in repr(record.__dict__) + 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.""" diff --git a/uv.lock b/uv.lock index 80fc76b..9aac18a 100644 --- a/uv.lock +++ b/uv.lock @@ -1310,6 +1310,7 @@ dependencies = [ { name = "cookiecutter" }, { name = "dlt" }, { name = "duckdb" }, + { name = "filelock" }, { name = "griffe" }, { name = "griffe-fieldz" }, { name = "jinja2" }, @@ -1362,6 +1363,7 @@ requires-dist = [ { 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" }, From 5ee3388e982d1c7579ee6b84353ba94ac3b15d6b Mon Sep 17 00:00:00 2001 From: Stran Dutton <65475509+StranDutton@users.noreply.github.com> Date: Wed, 28 Oct 2026 16:17:57 -0500 Subject: [PATCH 33/34] feat(managed): support AWS role profiles - BED-10272 (#111) * add support for roles (for openhound-aws specific AWS profile adapter) --- src/openhound/core/logging.py | 207 +++++++++++- src/openhound/scheduler/aws_profile.py | 154 ++++----- src/openhound/scheduler/service.py | 111 ++++-- tests/conftest.py | 62 ---- tests/test_aws_profile_adapter.py | 445 ++++++++++++++++++++----- tests/test_bhe_job_scheduling.py | 192 +++++++++-- tests/test_log_handlers.py | 309 ++++++++++++++++- 7 files changed, 1166 insertions(+), 314 deletions(-) diff --git a/src/openhound/core/logging.py b/src/openhound/core/logging.py index 2a6ae8d..46630da 100644 --- a/src/openhound/core/logging.py +++ b/src/openhound/core/logging.py @@ -11,9 +11,13 @@ 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 from rich.console import Console @@ -27,7 +31,7 @@ class _ValueRedactor(logging.Filter): def __init__(self, values: tuple[str, ...]): super().__init__() - self.values = values + self.values = tuple(value for value in values if value) self.handlers: set[logging.Handler] = set() def _redact(self, value): @@ -66,6 +70,167 @@ def filter(self, record: logging.LogRecord) -> bool: _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: @@ -177,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 @@ -330,6 +499,7 @@ def __init__( # 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_handler_lock = RLock() self.handlers = { LogMode.CLI: self.cli_handlers, @@ -400,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) @@ -410,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) @@ -473,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 @@ -485,15 +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 - 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) + # 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""" diff --git a/src/openhound/scheduler/aws_profile.py b/src/openhound/scheduler/aws_profile.py index 2434517..82b2255 100644 --- a/src/openhound/scheduler/aws_profile.py +++ b/src/openhound/scheduler/aws_profile.py @@ -1,4 +1,4 @@ -"""Adapt BHE's raw IAM access-key pair to the AWS extension's named profile. +"""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. @@ -6,7 +6,6 @@ import logging import os -import re import shutil import stat import tempfile @@ -16,6 +15,7 @@ from contextlib import 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 @@ -41,16 +41,7 @@ logger = logging.getLogger(__name__) -_CREDENTIAL_SETTINGS = { - "region_name", - "endpoint_url", - "org_role_name", - "org_role_session_name", - "include_management_account", - "source_account", - "management_account_id", - "management_role_name", -} +_ROLE_PROFILE_SETTINGS = {"role_arn", "role_session_name", "external_id"} class AWSAccessKeyError(ManagedConfigurationError): @@ -59,7 +50,7 @@ class AWSAccessKeyError(ManagedConfigurationError): def __init__( self, message: str, - summary: str = "Managed AWS credentials or authentication settings are invalid", + summary: str = "Managed AWS credentials could not be prepared", **context: Any, ) -> None: super().__init__(message) @@ -67,20 +58,8 @@ def __init__( self.context = context -def validate_access_key_id(value: Any) -> None: - # ASIA keys are temporary and require a session token, absent in BHE's beta - # contract. Do not strip input: INI/SDK whitespace normalization would change it. - if not isinstance(value, str) or re.fullmatch(r"AKIA[A-Z0-9]{16}", value) is None: - raise AWSAccessKeyError( - "Managed AWS job requires a long-lived IAM Access Key ID in external_secret_key_id" - ) - - def validate_secret_reference(value: Any) -> str: - if ( - not isinstance(value, str) - or re.fullmatch(r"[A-Za-z0-9/_+=.@:-]{1,2048}", value) is None - ): + if not isinstance(value, str) or not value.strip(): raise AWSAccessKeyError( "Managed AWS job requires a valid secret_key_id reference" ) @@ -93,13 +72,12 @@ class AWSAccessKeyCredentials: secret_access_key: str = field(repr=False) def __post_init__(self) -> None: - validate_access_key_id(self.access_key_id) - if ( - not isinstance(self.secret_access_key, str) - or re.fullmatch(r"[A-Za-z0-9/+=]{40}", self.secret_access_key) is 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 credential secret must be the raw 40-character Secret Access Key" + "Managed AWS raw credentials require a key identifier and text secret payload" ) @@ -110,12 +88,11 @@ def prepare_aws_credentials( An explicit key ID selects the beta raw-key contract. Without it, retain compatibility with nested credential documents; a raw payload still needs - a valid key ID and must never fall back to the scheduler's AWS identity. + 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: - validate_access_key_id(job.external_secret_key_id) aws_profile_params(job.params) try: credentials = load_secret(reference) @@ -143,7 +120,7 @@ def prepare_aws_credentials( def aws_profile_params( params: dict[str, Any], *, collector_name: str = "aws" ) -> dict[str, Any]: - """Validate the raw-key mode and copy its full nested settings before /start.""" + """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( @@ -153,7 +130,6 @@ def aws_profile_params( "Managed AWS params must use nested configuration sections" ) document = deepcopy(params) - _validate_credential_sections(document) node = document for section in ("sources", "source", "aws"): node = node.setdefault(section, {}) @@ -162,47 +138,18 @@ def aws_profile_params( "Managed AWS source configuration must be a nested object" ) credentials = node.setdefault("credentials", {}) - if not isinstance(credentials, dict) or credentials.keys() - _CREDENTIAL_SETTINGS: + if not isinstance(credentials, dict): raise AWSAccessKeyError( - "Managed AWS access-key jobs cannot select another profile or authentication mode" + "Managed AWS credentials configuration must be a nested object" ) - if any( - key in node - for key in ("profile_name", "profile_names", "auth_mode", "authentication_mode") - ): - raise AWSAccessKeyError( - "Managed AWS access-key jobs contain conflicting authentication settings" - ) - # An incomplete organization configuration must not fall through to the single - # profile model and silently ignore account/role settings. - organization = _CREDENTIAL_SETTINGS - {"region_name", "endpoint_url"} - if credentials.keys() & organization and ( - not isinstance(credentials.get("org_role_name"), str) - or not credentials["org_role_name"].strip() - ): - raise AWSAccessKeyError( - "Managed AWS organization settings require org_role_name" - ) - return document - - -def _validate_credential_sections( - document: dict[str, Any], path: tuple[str, ...] = () -) -> None: - # DLT accepts compact, less-specific and pipeline-prefixed layouts too. - # Keep every source-credential setting in the canonical raw-key section so - # another union alternative cannot win during pipeline-scoped resolution. - for key, value in document.items(): - if ( - key == "credentials" - and path != ("sources", "source", "aws") - and "destination" not in path - ): + 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 credential settings must use sources.source.aws.credentials" + "Managed AWS role settings require a non-empty role_arn" ) - if isinstance(value, dict): - _validate_credential_sections(value, (*path, key)) + return document def _workspace_root() -> Path: @@ -274,11 +221,10 @@ def __init__(self, provider: ConfigProvider) -> None: def get_value( self, key: str, hint: type[Any], pipeline_name: str, *sections: str ) -> tuple[Any, str]: - if not self.provider.supports_sections and key in _CREDENTIAL_SETTINGS | { - "profile_name", - "profile_names", - "credentials", - }: + # 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 @@ -299,9 +245,32 @@ def supports_sections(self) -> bool: return self.provider.supports_sections -def _write_ini(path: Path, sections: dict[str, dict[str, str]]) -> None: +def _write_ini(path: Path, sections: dict[str, dict[str, Any]]) -> None: config = ConfigParser(interpolation=None) - config.read_dict(sections) + 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( @@ -309,7 +278,7 @@ def _write_ini(path: Path, sections: dict[str, dict[str, str]]) -> None: "w", encoding="utf-8", ) as stream: - config.write(stream) + stream.write(serialized) @contextmanager @@ -369,10 +338,25 @@ def _log_collection_failure(error: Exception) -> None: def managed_aws_profile( params: dict[str, Any], credentials: AWSAccessKeyCredentials, directory: Path ) -> Iterator[str]: - """Install one profile before loading collectors; restore all worker state.""" + """Install a direct-key or role profile; restore all worker state afterward.""" document = aws_profile_params(params) profile = f"openhound-job-{uuid4().hex}" - document["sources"]["source"]["aws"]["credentials"]["profile_name"] = profile + 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", @@ -396,13 +380,15 @@ def managed_aws_profile( _write_ini( files["AWS_SHARED_CREDENTIALS_FILE"], { - profile: { + 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}": {}}) + _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 diff --git a/src/openhound/scheduler/service.py b/src/openhound/scheduler/service.py index 2f61ce6..e40d08a 100644 --- a/src/openhound/scheduler/service.py +++ b/src/openhound/scheduler/service.py @@ -1,5 +1,5 @@ import logging -import math +import multiprocessing import shutil import signal import tempfile @@ -50,6 +50,7 @@ ) from openhound.scheduler.configuration import ( ManagedConfigurationError, + _validate_document, managed_job_configuration, ) @@ -107,26 +108,6 @@ class PreparedManagedJob: credentials: dict[str, Any] | AWSAccessKeyCredentials = field(repr=False) -def _usable_credential_value(value: Any) -> bool: - """Validate generic JSON values without inventing a collector field schema.""" - if isinstance(value, str): - return bool(value.strip()) - if isinstance(value, dict): - return bool(value) and all( - isinstance(key, str) - and bool(key.strip()) - and _usable_credential_value(item) - for key, item in value.items() - ) - if isinstance(value, list): - return bool(value) and all(_usable_credential_value(item) for item in value) - if isinstance(value, (bool, int)): - return True - if isinstance(value, float): - return math.isfinite(value) - return False - - @dataclass class Result: results: dict @@ -161,6 +142,57 @@ def _subprocess_collect( 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( @@ -300,7 +332,13 @@ def __init__( # 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 @@ -327,7 +365,10 @@ def _shutdown(self) -> None: "Unexpected error handling completed job during shutdown." ) self._clear_running_job() - self.executor.shutdown(wait=True, cancel_futures=True) + 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: @@ -362,10 +403,19 @@ def _cleanup_aws_profile_workspace(self) -> 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. @@ -439,7 +489,8 @@ def _load_managed_credentials( credentials = self.secrets_manager.get_secret(job.secret_key_id) if isinstance(credentials, AWSAccessKeyCredentials): return credentials - if isinstance(credentials, dict) and _usable_credential_value(credentials): + if isinstance(credentials, dict): + _validate_document(credentials, allow_empty=True) return credentials except AWSAccessKeyError as error: raise CredentialValidationError( @@ -447,6 +498,8 @@ def _load_managed_credentials( ) 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 @@ -455,9 +508,7 @@ def _load_managed_credentials( f"Credential secret reader failed ({type(error).__name__})", "Credential secret retrieval failed", ) from None - raise CredentialValidationError( - "Credential document must be a non-empty JSON object with usable values" - ) + 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.""" @@ -799,7 +850,7 @@ def _abort_managed_job(self, job_id: str) -> None: ): self._remove_job_output(future.result()) self._clear_running_job() - self.executor = ProcessPoolExecutor(max_workers=1, max_tasks_per_child=1) + 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.""" diff --git a/tests/conftest.py b/tests/conftest.py index 99708ea..ee43c1d 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,12 +1,6 @@ import importlib import importlib.util -import os -import sys from pathlib import Path -from typing import Union - -import dlt -import pytest def import_from_path(file_path: Path) -> None: @@ -22,59 +16,3 @@ def pytest_configure(config): py_files = [path for path in sources_root.rglob("**/models/*.py") if path.is_file()] for module in py_files: import_from_path(module) - - -@pytest.fixture(scope="session") -def aws_extension_auth(): - """Load the real extension models/auth, without its unrelated native engine. - - Use the sibling repo, OPENHOUND_AWS_SOURCE, or an installed extension. These - model integration tests are optional when the internal extension is absent. - """ - default = Path(__file__).resolve().parents[2] / "openhound-aws" - repo = Path(os.environ.get("OPENHOUND_AWS_SOURCE", default)) - path = repo / "src" / "openhound_aws" / "auth.py" - if not path.is_file(): - if "OPENHOUND_AWS_SOURCE" in os.environ: - pytest.fail("OPENHOUND_AWS_SOURCE must point to an openhound-aws checkout") - return pytest.importorskip( - "openhound_aws.auth", - reason="AWS model integration requires openhound-aws or OPENHOUND_AWS_SOURCE", - ) - name = "_openhound_aws_auth_integration" - spec = importlib.util.spec_from_file_location(name, path) - module = importlib.util.module_from_spec(spec) - sys.modules[name] = module - spec.loader.exec_module(module) - return module - - -@pytest.fixture -def resolve_aws_settings(aws_extension_auth): - auth = aws_extension_auth - observed = {} - - # Use the real model union and DLT source resolution used by the extension. - # Resolving the configuration alone does not need its Rust policy evaluator. - @dlt.source(name="aws", section="source") - def source( - credentials: Union[ # noqa: UP007 - match the extension source's union - auth.AWSOrganizationProfileCredentials, - auth.AWSMultiProfileCredentials, - auth.AWSProfileCredentials, - ] = dlt.secrets.value, - bucket_policy: auth.BucketPolicyConfiguration = dlt.config.value, - max_attempts: int = 3, - ): - observed.update( - credentials=credentials, - bucket_policy=bucket_policy, - max_attempts=max_attempts, - ) - return dlt.resource([], name="configuration_only") - - def resolve(): - source() - return observed.copy() - - return resolve diff --git a/tests/test_aws_profile_adapter.py b/tests/test_aws_profile_adapter.py index f1f4745..550f136 100644 --- a/tests/test_aws_profile_adapter.py +++ b/tests/test_aws_profile_adapter.py @@ -7,6 +7,7 @@ import traceback import zipfile from concurrent.futures import ProcessPoolExecutor +from configparser import ConfigParser from copy import deepcopy from datetime import UTC, datetime, timedelta from types import SimpleNamespace @@ -16,6 +17,7 @@ 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 @@ -39,7 +41,11 @@ PAIR = AWSAccessKeyCredentials(ACCESS_KEY, SECRET_KEY) SOURCE_ACCOUNT = "999999999999" MANAGEMENT_ACCOUNT = "111111111111" -MEMBER_ACCOUNT = "222222222222" +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) @@ -70,6 +76,62 @@ def workspace(): 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, @@ -89,6 +151,123 @@ def _assert_source_credentials(session): ) +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() + + def test_boto3_file_provider_uses_exact_pair_without_a_session_token( workspace, monkeypatch ): @@ -108,8 +287,90 @@ def test_boto3_file_provider_uses_exact_pair_without_a_session_token( 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( - workspace, monkeypatch, caplog + role_profile, workspace, monkeypatch, caplog ): original_environment = dict(os.environ) providers = Container()[PluggableRunContext].providers @@ -122,9 +383,14 @@ def fail_config_write(path, sections): 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({}, PAIR, workspace.directory), + 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)) @@ -234,9 +500,7 @@ def test_sdk_debug_responses_are_suppressed_and_logging_state_restored( sdk_logger.handlers, sdk_logger.disabled, sdk_logger.propagate = previous[1:] -def test_real_sdk_and_extension_profile_resolution( - workspace, aws_extension_auth, resolve_aws_settings -): +def test_profile_resolves_in_sdk_and_preserves_nested_configuration(workspace): params = { "sources": { "source": { @@ -257,22 +521,27 @@ def test_real_sdk_and_extension_profile_resolution( } original = deepcopy(params) with managed_aws_profile(params, PAIR, workspace.directory) as profile: - settings = resolve_aws_settings() - model = settings["credentials"] - assert type(model) is aws_extension_auth.AWSProfileCredentials - assert model.profile_name == profile - assert model.region_name == "eu-west-1" - assert model.endpoint_url == "https://aws.example.test" - _assert_source_credentials(model.session) - client = model.session.client("sts", endpoint_url=model.endpoint_url) + 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 - bucket = settings["bucket_policy"] - assert bucket.include is True and bucket.max_objects == 9 - assert bucket.bucket_whitelist[0].match == "public/.*" + 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 @@ -287,7 +556,7 @@ def test_real_sdk_and_extension_profile_resolution( def test_runtime_and_local_configuration_cannot_override_generated_identity( - workspace, resolve_aws_settings, aws_extension_auth, monkeypatch + workspace, monkeypatch ): for key, value in { "AWS_ACCESS_KEY_ID": "scheduler-access-key", @@ -339,10 +608,22 @@ def get_value(self, key, hint, pipeline_name, *sections): original_environment = dict(os.environ) try: with managed_aws_profile({}, PAIR, workspace.directory) as profile: - model = resolve_aws_settings()["credentials"] - assert type(model) is aws_extension_auth.AWSProfileCredentials - assert model.profile_name == profile - _assert_source_credentials(model.session) + 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" @@ -352,11 +633,17 @@ def get_value(self, key, hint, pipeline_name, *sections): 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_models_and_assume_role_use_generated_source_identity( - dedicated_source, workspace, aws_extension_auth, resolve_aws_settings, monkeypatch +def test_organization_settings_pass_through_with_generated_source_profile( + dedicated_source, + role_profile, + workspace, + profile_assumptions, ): options = { "org_role_name": "MemberCollectorRole", @@ -370,77 +657,35 @@ def test_organization_models_and_assume_role_use_generated_source_identity( 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}}}} - clients = [] - sessions = [] - original_client = boto3.session.Session.client - - def stub_sts(session, name, **kwargs): - assert name == "sts" - _assert_source_credentials(session) - sessions.append(session) - client = original_client(session, name, **kwargs) - stubber = Stubber(client) - if not clients: - stubber.add_response("get_caller_identity", _identity(), {}) + 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: - account, role = ( - (MANAGEMENT_ACCOUNT, "ManagementCollectorRole") - if dedicated_source and len(clients) == 1 - else (MEMBER_ACCOUNT, "MemberCollectorRole") - ) - stubber.add_response( - "assume_role", - { - "Credentials": { - "AccessKeyId": "ASIAIOSFODNN7EXAMPLE", - "SecretAccessKey": "T" * 40, - "SessionToken": "assumed-session-token", - "Expiration": datetime.now(UTC) + timedelta(hours=1), - } - }, - { - "RoleArn": f"arn:aws:iam::{account}:role/{role}", - "RoleSessionName": "job-session", - }, - ) - stubber.activate() - clients.append(stubber) - return client - - monkeypatch.setattr(boto3.session.Session, "client", stub_sts) - try: - with managed_aws_profile(params, PAIR, workspace.directory) as profile: - model = resolve_aws_settings()["credentials"] - assert type(model) is aws_extension_auth.AWSOrganizationProfileCredentials - assert model.profile_name == profile - assert model.org_role_name == "MemberCollectorRole" - assert model.collect_management_account is False - auth = aws_extension_auth.AWSAuth(model) - assert auth.source_ctx.account_id == SOURCE_ACCOUNT - if dedicated_source: - management = auth.management_ctx - assert management.account_id == MANAGEMENT_ACCOUNT - assert ( - management.session.get_credentials().token - == "assumed-session-token" - ) - member = auth.context_for_account(MEMBER_ACCOUNT) - assert member.account_id == MEMBER_ACCOUNT - assert member.session.get_credentials().token == "assumed-session-token" - # Both management and member assumptions originate from the raw pair, - # rather than chaining through management's temporary credentials. - assert all(session is auth.source_ctx.session for session in sessions) - for stubber in clients: - stubber.assert_no_pending_responses() - finally: - for stubber in clients: - stubber.deactivate() + _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, resolve_aws_settings, monkeypatch + failed, role_profile, profile_assumptions, monkeypatch ): monkeypatch.setenv("AWS_ACCESS_KEY_ID", "scheduler-access") monkeypatch.setenv("AWS_SECRET_ACCESS_KEY", "scheduler-secret") @@ -463,12 +708,27 @@ def test_successive_jobs_restore_environment_providers_and_remove_files( 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) - model = resolve_aws_settings()["credentials"] - assert model.region_name == ("eu-west-1" if index == 0 else "us-east-1") - _assert_source_credentials(model.session) + 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: @@ -607,6 +867,7 @@ def test_parent_removes_credentials_after_forced_worker_termination( 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() diff --git a/tests/test_bhe_job_scheduling.py b/tests/test_bhe_job_scheduling.py index 02c06c1..d8ba37f 100644 --- a/tests/test_bhe_job_scheduling.py +++ b/tests/test_bhe_job_scheduling.py @@ -10,6 +10,7 @@ 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 @@ -388,6 +389,11 @@ 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", @@ -1835,7 +1841,6 @@ def aws_managed_service(managed_service, mock_bloodhound_api): def test_active_aws_poll_uses_claimed_key_and_raw_secret_in_worker_adapter( aws_managed_service, mock_bloodhound_api, - resolve_aws_settings, failed, binary_secret, monkeypatch, @@ -1868,10 +1873,12 @@ def load_collectors(): assert state.collector_job_start_requests == [ (state.collector_job_claim_response["id"], b"") ] - model = resolve_aws_settings()["credentials"] - assert model.profile_name.startswith("openhound-job-") - assert model.region_name == "us-west-2" - credentials = model.session.get_credentials() + 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) == ( @@ -1958,15 +1965,9 @@ def submit(function, *args): "field,value", [ ("external_secret_key_id", None), - ("external_secret_key_id", ""), - ("external_secret_key_id", "ASIAIOSFODNN7EXAMPLE"), - ("external_secret_key_id", AWS_ACCESS_KEY + "\n"), - ("external_secret_key_id", "not-an-access-key"), ("secret_key_id", None), ("secret_key_id", ""), ("secret_key_id", " "), - ("secret_key_id", " private-secret-reference"), - ("secret_key_id", "reference\nother"), ], ) def test_invalid_aws_metadata_fails_before_start_without_runtime_fallback( @@ -2005,14 +2006,6 @@ def retrieve(reference): "response", [ {}, - {"SecretString": ""}, - {"SecretString": " " + AWS_SECRET_KEY}, - {"SecretString": AWS_SECRET_KEY + "\n"}, - {"SecretString": AWS_SECRET_KEY[:-1]}, - {"SecretString": "x" * 41}, - {"SecretString": "%" * 40}, - {"SecretString": json.dumps(AWS_SECRET_KEY)}, - {"SecretString": '["key"]'}, {"SecretString": json.dumps({"aws_secret_access_key": AWS_SECRET_KEY})}, {"SecretString": json.dumps(COLLECTOR_SECRETS)}, {"SecretBinary": b"\xff\xfe"}, @@ -2047,9 +2040,44 @@ def test_invalid_raw_aws_payloads_are_not_trimmed_wrapped_or_treated_as_document {"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_authentication_conflicts_fail_before_start( +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 @@ -2059,11 +2087,52 @@ def test_raw_aws_authentication_conflicts_fail_before_start( monkeypatch.setattr( aws_managed_service.secrets_manager, "get_secret", - lambda _: pytest.fail("invalid authentication must fail before retrieval"), + 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" @@ -2323,9 +2392,7 @@ def test_aws_launch_paths_validate_before_start( params = ( {} if wrong_collector - else { - "sources": {"source": {"aws": {"credentials": {"profile_name": "local"}}}} - } + else {"sources": {"source": {"aws": {"credentials": "not-an-object"}}}} ) kwargs = {"managed": True} if launch_path == "_launch_job" else {} with pytest.raises(AWSAccessKeyError): @@ -2348,7 +2415,6 @@ def test_aws_launch_paths_validate_before_start( {"sources": {"credentials": {"org_role_name": "OtherRole"}}}, {"sources": {"source": {"credentials": {"profile_name": "other"}}}}, {"sources": {"aws": {"credentials": {"profile_names": ["other"]}}}}, - {"profile_name": "flat-profile"}, {"aws_collector": {"credentials": {"profile_names": ["other"]}}}, { "aws_collector": { @@ -2359,15 +2425,75 @@ def test_aws_launch_paths_validate_before_start( }, ], ) -def test_raw_aws_rejects_ambiguous_authentication_layouts_before_start( +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 state.collector_job_start_requests == [] - assert aws_managed_service.executor.submitted == [] - assert state.collector_job_end_requests[0]["body"]["outcome"] == "failed" + 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( @@ -2612,12 +2738,9 @@ def unavailable(job_id): [ "plain-secret", '{"malformed":', - "{}", '{"value":" "}', '{"value":null}', - '{"value":{}}', '{"value":[]}', - '{"value":{"nested":""}}', '{"value":NaN}', '{" ":"private-credential-value"}', ], @@ -3487,6 +3610,9 @@ def test_managed_configuration_failure_is_logged_and_reported_with_the_same_reas 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}" @@ -4502,7 +4628,7 @@ def fail(extension, output_path): raise RuntimeError("conversion failed") monkeypatch.setattr(scheduler_service.dataflow, "pipeline", fail) - with pytest.raises(RuntimeError, match="conversion failed"): + with pytest.raises(ManagedConfigurationError, match="check worker logs"): _subprocess_collect( "example", "6bc0857e-9d9b-48c6-bbd8-77a31c5d351d", managed=True ) diff --git a/tests/test_log_handlers.py b/tests/test_log_handlers.py index 454a6ba..f848814 100644 --- a/tests/test_log_handlers.py +++ b/tests/test_log_handlers.py @@ -1,5 +1,8 @@ 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 @@ -17,6 +20,7 @@ OpenHoundRichFormatter, OpenHoundTextFormatter, RotatingFileHandler, + _ParentWorkerHandler, _ValueRedactor, redact_log_values, ) @@ -337,7 +341,13 @@ def test_redaction_formats_original_arguments(message, args, secret, expected): assert record.msg == record.getMessage() == expected assert record.args == () assert record.response == {"credential": "[REDACTED]", "count": 7} - assert secret not in repr(record.__dict__) + # 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) @@ -431,6 +441,32 @@ def test_get_file_handler_caches_handler_per_path(tmp_path): handler.close() +@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: + record = logging.LogRecord( + "dlt", logging.ERROR, __file__, 1, "one-delivery", (), None + ) + _ParentWorkerHandler(configured_logger).handle(record) + assert output.getvalue().count("one-delivery") == 1 + finally: + configured_logger.root_logger.removeHandler(handlers[0]) + configured_logger.dlt_logger.removeHandler(handlers[1]) + for handler in handlers: + handler.close() + + 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 @@ -490,3 +526,274 @@ 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 + 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 From 4865ab53d4e74d6ca58cd42ee9b4678b3351047e Mon Sep 17 00:00:00 2001 From: Stran Dutton <65475509+StranDutton@users.noreply.github.com> Date: Fri, 30 Oct 2026 09:48:06 -0500 Subject: [PATCH 34/34] fix(managed): fail AWS jobs on initial auth failure (BED-10368) (#114) --- src/openhound/scheduler/aws_profile.py | 109 +++++++++++- src/openhound/scheduler/service.py | 9 +- tests/test_aws_profile_adapter.py | 222 ++++++++++++++++++++++++- tests/test_bhe_job_scheduling.py | 189 +++++++++++++++++++++ tests/test_log_handlers.py | 2 + 5 files changed, 528 insertions(+), 3 deletions(-) diff --git a/src/openhound/scheduler/aws_profile.py b/src/openhound/scheduler/aws_profile.py index 82b2255..0056047 100644 --- a/src/openhound/scheduler/aws_profile.py +++ b/src/openhound/scheduler/aws_profile.py @@ -2,6 +2,10 @@ 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 @@ -12,7 +16,7 @@ import traceback from collections.abc import Callable, Iterator from configparser import ConfigParser -from contextlib import contextmanager +from contextlib import closing, contextmanager from copy import deepcopy from dataclasses import dataclass, field from io import StringIO @@ -20,7 +24,10 @@ 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 @@ -334,6 +341,106 @@ def _log_collection_failure(error: Exception) -> None: 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 diff --git a/src/openhound/scheduler/service.py b/src/openhound/scheduler/service.py index e40d08a..604f80d 100644 --- a/src/openhound/scheduler/service.py +++ b/src/openhound/scheduler/service.py @@ -47,6 +47,8 @@ cleanup_orphaned_aws_profiles, managed_aws_profile, prepare_aws_credentials, + validate_aws_profile, + validate_configured_aws_profiles, ) from openhound.scheduler.configuration import ( ManagedConfigurationError, @@ -198,7 +200,10 @@ def _configured_worker_collect( raise ManagedConfigurationError( "Managed AWS profile worker context is required" ) - with managed_aws_profile(params or {}, secrets, aws_profile_directory): + 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 = ( @@ -209,6 +214,8 @@ def _configured_worker_collect( # 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) diff --git a/tests/test_aws_profile_adapter.py b/tests/test_aws_profile_adapter.py index 550f136..aa14407 100644 --- a/tests/test_aws_profile_adapter.py +++ b/tests/test_aws_profile_adapter.py @@ -6,7 +6,7 @@ import time import traceback import zipfile -from concurrent.futures import ProcessPoolExecutor +from concurrent.futures import Future, ProcessPoolExecutor from configparser import ConfigParser from copy import deepcopy from datetime import UTC, datetime, timedelta @@ -24,8 +24,11 @@ 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 ( @@ -268,6 +271,223 @@ def test_denied_profile_role_fails_without_source_fallback( 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 ): diff --git a/tests/test_bhe_job_scheduling.py b/tests/test_bhe_job_scheduling.py index d8ba37f..97a77a5 100644 --- a/tests/test_bhe_job_scheduling.py +++ b/tests/test_bhe_job_scheduling.py @@ -22,6 +22,8 @@ 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 @@ -1780,6 +1782,10 @@ def collect(extension, output_path): 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() @@ -1807,6 +1813,162 @@ def submit(function, *args): 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 @@ -1869,6 +2031,29 @@ def handoff(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"") @@ -1959,6 +2144,10 @@ def submit(function, *args): 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( diff --git a/tests/test_log_handlers.py b/tests/test_log_handlers.py index f848814..a6792ad 100644 --- a/tests/test_log_handlers.py +++ b/tests/test_log_handlers.py @@ -585,6 +585,8 @@ def collect(*args): 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 )