diff --git a/.gitignore b/.gitignore index 0819916b..1fbfd384 100644 --- a/.gitignore +++ b/.gitignore @@ -4,6 +4,7 @@ ui_catalog.db site/ addons/ collectors/ +.worktrees/ notebooks diff --git a/example-configurations/bloodhound-enterprise/docker-compose.yml b/example-configurations/bloodhound-enterprise/docker-compose.yml index 657ccef8..3322ebda 100644 --- a/example-configurations/bloodhound-enterprise/docker-compose.yml +++ b/example-configurations/bloodhound-enterprise/docker-compose.yml @@ -1,5 +1,8 @@ x-scheduler: &scheduler image: specterops/openhound:${IMAGE_VERSION:-latest-enterprise} + build: + context: ../.. + target: enterprise restart: unless-stopped init: true volumes: @@ -30,6 +33,9 @@ services: environment: <<: *env DESTINATION__BLOODHOUNDENTERPRISE__COLLECTOR_NAME: github + # Docker mounts the GitHub App key at this absolute path. Do not use + # ~/.dlt/github.pem here: the collector does not expand `~`. + SOURCES__SOURCE__GITHUB__CREDENTIALS__KEY_PATH: /app/.dlt/github.pem secrets: - source: secrets_github target: /app/.dlt/secrets.toml diff --git a/justfile b/justfile index 2d5b92be..e70a9d12 100644 --- a/justfile +++ b/justfile @@ -31,3 +31,41 @@ typecheck: dashboard: @echo "Starting marimo openhound dashboard" marimo edit notebooks/explore.py --watch + +# Docker commands for the Enterprise example configuration. +# Examples: `just oh`, `just oh up github`, `just oh down github`, `just oh down` +oh action='up' collector='': + #!/usr/bin/env bash + set -euo pipefail + + compose=( + docker compose + -f ./example-configurations/bloodhound-enterprise/docker-compose.yml + ) + + case "{{action}}" in + up) + if [[ -n "{{collector}}" ]]; then + "${compose[@]}" up -d --build --force-recreate "scheduler-{{collector}}" + else + "${compose[@]}" up -d --build --force-recreate + fi + ;; + down) + if [[ -n "{{collector}}" ]]; then + "${compose[@]}" down "scheduler-{{collector}}" + else + "${compose[@]}" down + fi + ;; + *) + >&2 echo "Usage: just oh [up|down] [github|jamf|okta]" + exit 1 + ;; + esac + +# Follow the most recent logs for one collector, e.g. `just oh-logs github`. +oh-logs collector: + docker compose \ + -f ./example-configurations/bloodhound-enterprise/docker-compose.yml \ + logs --follow --tail=100 "scheduler-{{collector}}" diff --git a/src/openhound/core/clients/bloodhound.py b/src/openhound/core/clients/bloodhound.py index 9195df49..d9819495 100644 --- a/src/openhound/core/clients/bloodhound.py +++ b/src/openhound/core/clients/bloodhound.py @@ -72,7 +72,7 @@ def _request( data=body, ) - if response.status_code not in [200, 201, 202]: + if not 200 <= response.status_code < 300: raise BloodHoundHTTPError(code=response.status_code, reason=response.text) return response diff --git a/src/openhound/core/clients/bloodhound_enterprise.py b/src/openhound/core/clients/bloodhound_enterprise.py index 4517d1b0..b3268a74 100644 --- a/src/openhound/core/clients/bloodhound_enterprise.py +++ b/src/openhound/core/clients/bloodhound_enterprise.py @@ -1,22 +1,44 @@ +import base64 import gzip +import hashlib import json +import logging +import math import socket +import time from enum import Enum +from pathlib import Path +from typing import Callable, TypeVar -from openhound.core.clients.bloodhound import BloodHound +import openhound +import requests +from openhound.core.clients.bloodhound import BloodHound, BloodHoundHTTPError from openhound.core.clients.models.jobs import ( JobsAvailable, JobsCurrent, JobsEnd, JobStart, + ArtifactUploadSession, + ManagementAvailable, + ManagementOperationResult, + ManagementOperationStatus, ) +logger = logging.getLogger(__name__) + class JobStatus(str, Enum): COMPLETE = "complete" FAILED = "failed" +SUPPORT_BUNDLE_PART_SIZE = 8 * 1024 * 1024 # 8 MiB +SUPPORT_BUNDLE_MAX_RETRIES = 3 +SUPPORT_BUNDLE_RETRY_DELAY_SECONDS = 2 + +T = TypeVar("T") + + class BloodHoundEnterprise(BloodHound): @property def jobs_available(self) -> JobsAvailable: @@ -54,6 +76,156 @@ def ingest(self, data: str) -> None: method="POST", path=path, body=compressed_data, extra_headers=headers ) + @property + def management_available(self) -> ManagementAvailable: + response = self.request( + method="GET", path="/api/v2/clients/management/available" + ) + return ManagementAvailable.model_validate(response.json()) + + def start_operation(self, operation_id: str) -> ManagementOperationResult: + response = self._retry_support_bundle_request( + "start management operation", + lambda: self.request( + method="POST", + path="/api/v2/clients/management/start", + body=json.dumps({"operation_id": operation_id}).encode(), + ), + ) + return ManagementOperationResult.model_validate(response.json()) + + def end_operation( + self, operation_id: str, status: ManagementOperationStatus + ) -> ManagementOperationResult: + response = self._retry_support_bundle_request( + "end management operation", + lambda: self.request( + method="POST", + path="/api/v2/clients/management/end", + body=json.dumps( + {"operation_id": operation_id, "status": status} + ).encode(), + ), + ) + return ManagementOperationResult.model_validate(response.json()) + + def create_artifact_upload( + self, operation_id: str, bundle_path: Path + ) -> ArtifactUploadSession: + total_size = bundle_path.stat().st_size + + logger.info("Total size of the support bundle: %s", total_size) + if total_size <= 0: + raise ValueError("Support bundle must not be empty.") + + part_size = SUPPORT_BUNDLE_PART_SIZE + checksum = self._file_checksum(bundle_path) + response = self._retry_support_bundle_request( + "create support bundle upload", + lambda: self.request( + method="POST", + path="/api/v2/clients/management/artifacts", + body=json.dumps( + { + "operation_id": operation_id, + "artifact_type": "support_bundle", + "total_size": total_size, + "part_size": part_size, + "part_count": math.ceil(total_size / part_size), + "content_type": "application/zip", + "checksum_algorithm": "sha256", + "checksum": checksum, + } + ).encode(), + ), + ) + return ArtifactUploadSession.model_validate(response.json()["data"]) + + def upload_artifact_part( + self, artifact_id: str, part_number: int, content: bytes + ) -> None: + checksum = base64.b64encode(hashlib.sha256(content).digest()).decode("ascii") + self._retry_support_bundle_request( + f"upload support bundle part {part_number}", + lambda: self.request( + method="POST", + path=f"/api/v2/clients/management/artifacts/{artifact_id}/parts/{part_number}", + body=content, + extra_headers={ + "Content-Length": str(len(content)), + "Content-Type": "application/zip", + "Content-Digest": f"sha-256=:{checksum}:", + }, + ), + ) + + def complete_artifact_upload(self, artifact_id: str, operation_id: str) -> None: + self._retry_support_bundle_request( + "complete support bundle upload", + lambda: self.request( + method="POST", + path=f"/api/v2/clients/management/artifacts/{artifact_id}/complete", + body=json.dumps({"operation_id": operation_id}).encode(), + ), + ) + + def upload_support_bundle(self, operation_id: str, bundle_path: Path) -> None: + """Create an upload session, transfer every ZIP part, then complete it.""" + session = self.create_artifact_upload(operation_id, bundle_path) + with bundle_path.open("rb") as bundle: + for part_number in range(1, session.part_count + 1): + part = bundle.read(session.part_size) + if not part: + raise ValueError(f"Support bundle ended before part {part_number}.") + self.upload_artifact_part(session.artifact_id, part_number, part) + if bundle.read(1): + raise ValueError("Support bundle grew while it was being uploaded.") + self.complete_artifact_upload(session.artifact_id, operation_id) + + @staticmethod + def _file_checksum(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as bundle: + for chunk in iter(lambda: bundle.read(1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + @staticmethod + def _is_transient_support_bundle_error(error: Exception) -> bool: + if isinstance(error, requests.RequestException): + return True + return isinstance(error, BloodHoundHTTPError) and error.code in { + 408, + 429, + 500, + 502, + 503, + 504, + } + + def _retry_support_bundle_request( + self, description: str, request: Callable[[], T] + ) -> T: + for retry in range(SUPPORT_BUNDLE_MAX_RETRIES + 1): + try: + return request() + except Exception as error: + if not self._is_transient_support_bundle_error(error): + raise + if retry == SUPPORT_BUNDLE_MAX_RETRIES: + raise + logger.warning( + "%s failed transiently; retrying in %s seconds (%s/%s).", + description, + SUPPORT_BUNDLE_RETRY_DELAY_SECONDS, + retry + 1, + SUPPORT_BUNDLE_MAX_RETRIES, + exc_info=True, + ) + time.sleep(SUPPORT_BUNDLE_RETRY_DELAY_SECONDS) + + raise AssertionError("Support bundle retry loop exited unexpectedly.") + def update_client_metadata(self) -> None: path = "/api/v2/clients/update" try: diff --git a/src/openhound/core/clients/models/jobs.py b/src/openhound/core/clients/models/jobs.py index b02e6546..de24a378 100644 --- a/src/openhound/core/clients/models/jobs.py +++ b/src/openhound/core/clients/models/jobs.py @@ -1,6 +1,6 @@ from pydantic import BaseModel from datetime import datetime -from typing import Union +from enum import StrEnum class Job(BaseModel): @@ -29,9 +29,9 @@ class DateAt(BaseModel): class JobStartData(Job): start_time: datetime end_time: datetime - created_at: Union[datetime, DateAt] - updated_at: Union[datetime, DateAt] - deleted_at: Union[datetime, DateAt] + created_at: datetime | DateAt + updated_at: datetime | DateAt + deleted_at: datetime | DateAt log_path: str | None event_title: str last_ingest: datetime @@ -51,3 +51,75 @@ class JobsCurrent(BaseModel): class JobsEnd(BaseModel): data: Job + + +class ManagementOperationType(StrEnum): + SUPPORT_BUNDLE = "support_bundle" + + +class ManagementOperationStatus(StrEnum): + QUEUED = "queued" + RUNNING = "running" + SUCCEEDED = "succeeded" + FAILED = "failed" + CANCELED = "canceled" + + +class ManagementOperation(BaseModel): + id: str + client_id: str + artifact_id: str | None = None + type: ManagementOperationType + status: ManagementOperationStatus + created_at: datetime + updated_at: datetime | None = None + requested_by_user_id: str | None = None + started_at: datetime | None = None + completed_at: datetime | None = None + execution_time: datetime | None = None + + +class ManagementAvailable(BaseModel): + data: list[ManagementOperation] + + +class ManagementOperationResult(BaseModel): + data: ManagementOperation + + +class ArtifactPart(BaseModel): + part_number: int + size: int + checksum: str + offset_start: int + storage_key: str + created_at: datetime + completed_at: datetime | None = None + + +class ArtifactStatus(StrEnum): + PENDING = "pending" + UPLOADING = "uploading" + COMPLETE = "complete" + FAILED = "failed" + CANCELED = "canceled" + + +class ArtifactUploadSession(BaseModel): + """Response data from creating a management artifact upload session. + + This models every field returned by + ``POST /api/v2/clients/management/artifacts``. The upload flow currently + needs only ``artifact_id``, ``part_size``, and ``part_count``, but callers + can use the remaining session and operation metadata without consulting + the API implementation. + """ + + artifact_id: str + client_id: str + storage_key: str + status: ArtifactStatus + part_size: int + part_count: int + missing_parts: list[int] + management_operation: ManagementOperation diff --git a/src/openhound/core/support_bundle.py b/src/openhound/core/support_bundle.py new file mode 100644 index 00000000..238ad566 --- /dev/null +++ b/src/openhound/core/support_bundle.py @@ -0,0 +1,47 @@ +"""Create uploadable support bundles from OpenHound logs.""" + +import logging +import tempfile +import zipfile +from datetime import UTC, datetime +from pathlib import Path + +logger = logging.getLogger(__name__) + +_LOG_PATTERNS = ( + "openhound.log", + "openhound.log.*", + "ext_*.log", + "ext_*.log.*", +) + + +def collect_log_files(log_base_path: Path) -> list[Path]: + """Return current and rotated platform and extension logs, sorted by path.""" + if not log_base_path.is_dir(): + logger.warning("Log directory does not exist: %s", log_base_path) + return [] + + files = { + path + for pattern in _LOG_PATTERNS + for path in log_base_path.glob(pattern) + if path.is_file() + } + return sorted(files) + + +def create_support_bundle(collector_name: str, log_base_path: Path) -> Path: + """Archive logs in a temporary ZIP; the caller must remove the returned file.""" + timestamp = datetime.now(UTC).strftime("%Y-%m-%d-%H-%M-%S") + bundle_path = ( + Path(tempfile.mkdtemp()) / f"{collector_name}_support_bundle_{timestamp}.zip" + ) + + with zipfile.ZipFile(bundle_path, "w", compression=zipfile.ZIP_DEFLATED) as archive: + for log_file in collect_log_files(log_base_path): + archive.write(log_file, arcname=log_file.name) + + logger.info("Support Bundle size: %s", bundle_path.stat().st_size) + logger.info("Created support bundle at %s", bundle_path) + return bundle_path diff --git a/src/openhound/scheduler/service.py b/src/openhound/scheduler/service.py index 2b24643c..ccaf676e 100644 --- a/src/openhound/scheduler/service.py +++ b/src/openhound/scheduler/service.py @@ -4,11 +4,18 @@ from concurrent.futures import Future, ProcessPoolExecutor from concurrent.futures.process import BrokenProcessPool from dataclasses import dataclass +from pathlib import Path -import openhound.core.logging # noqa: F401 +import openhound.core.logging as openhound_logging from openhound.core.clients.bloodhound_enterprise import BloodHoundEnterprise, JobStatus -from openhound.core.clients.models.jobs import Job +from openhound.core.clients.models.jobs import ( + Job, + ManagementOperation, + ManagementOperationStatus, + ManagementOperationType, +) from openhound.core.manager import CollectorManager +from openhound.core.support_bundle import create_support_bundle from openhound.scheduler import dataflow logger = logging.getLogger(__name__) @@ -70,6 +77,7 @@ def __init__( token_key: str, token_id: str, collector_name: str, + log_base_path: Path | None = None, ): # BHE client settings self.bhe_uri = bhe_uri @@ -79,6 +87,7 @@ def __init__( ) # Interval how often to check for a job self.interval = POLL_INTERVAL + self.log_base_path = log_base_path or openhound_logging.logger_override.base_path # Stores the ID of currently running BHE job self.job_running: int | None = None @@ -136,6 +145,41 @@ def check_jobs(self) -> Job | None: 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.") + for operation in self.client.management_available.data: + if ( + operation.type is ManagementOperationType.SUPPORT_BUNDLE + and operation.status is ManagementOperationStatus.QUEUED + ): + return operation + return None + + def _send_support_bundle(self, operation: ManagementOperation) -> None: + """Claim, upload, and complete a support-bundle operation.""" + bundle_path: Path | None = None + try: + self.client.start_operation(operation.id) + bundle_path = create_support_bundle(self.collector_name, self.log_base_path) + self.client.upload_support_bundle(operation.id, bundle_path) + except Exception: + logger.exception("Support bundle operation %s failed.", operation.id) + try: + self.client.end_operation(operation.id, ManagementOperationStatus.FAILED) + except Exception: + logger.exception("Unable to mark management operation %s as failed.", operation.id) + raise + finally: + if bundle_path is not None: + try: + bundle_path.unlink(missing_ok=True) + bundle_path.parent.rmdir() + except OSError: + logger.exception( + "Unable to remove support bundle for operation %s.", operation.id + ) + 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 @@ -218,16 +262,32 @@ def _poll(self) -> None: self.future = None self.job_running = None - # If no job is currently running, check for new jobs available and start the job - try: - if self.job_running is None: + # Management operations have priority over new collections while idle. + if self.job_running is None: + try: + operation = self.check_management() + except Exception: + logger.exception("Error checking management operations.") + operation = None + + if operation is not None: + try: + self._send_support_bundle(operation) + except Exception: + logger.exception("Error executing management operation.") + return + + try: available_job = self.check_jobs() if available_job: self._start_job(available_job) - else: + except Exception: + logger.exception("Error checking for or starting jobs.") + else: + try: self.client.jobs_current - except Exception: - logger.exception("Error checking for or starting jobs.") + except Exception: + logger.exception("Error checking in-progress 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 6e483635..4f9804a3 100644 --- a/tests/test_bhe_job_scheduling.py +++ b/tests/test_bhe_job_scheduling.py @@ -1,4 +1,6 @@ +import base64 import gzip +import hashlib import json from concurrent.futures import Future from concurrent.futures.process import BrokenProcessPool @@ -10,7 +12,13 @@ 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.models.jobs import ( + ManagementOperation, + ManagementOperationStatus, + ManagementOperationType, +) from openhound.core.models.graph import Graph from openhound.scheduler import service as scheduler_service from openhound.scheduler.service import ( @@ -21,6 +29,7 @@ ) TEST_DATA_DIR = Path(__file__).parent / "test_data" / "api" / "jobs" +MANAGEMENT_DATA_DIR = Path(__file__).parent / "test_data" / "api" / "management" def load_json(filename: str) -> dict: @@ -43,6 +52,16 @@ def mock_bloodhound_api(): app.state.start_payload = None app.state.client_update_payload = None app.state.ingested_edges = 0 + app.state.management_operations = [] + app.state.operation_started = False + app.state.operation_ended = False + app.state.operation_start_payload = None + app.state.operation_end_payload = None + app.state.operation_completed_by_artifact_upload = False + app.state.bundle_content = None + app.state.artifact_create_payload = None + app.state.uploaded_parts = [] + app.state.artifact_completed = False app.state.ingested_nodes = 0 @app.get("/api/v2/jobs/available") @@ -81,6 +100,85 @@ async def update_client(body: dict): app.state.client_update_payload = body return {"status": "success"} + @app.get("/api/v2/clients/management/available") + async def management_available(): + return {"data": app.state.management_operations} + + @app.post("/api/v2/clients/management/start") + async def start_operation(body: dict): + app.state.operation_started = True + app.state.operation_start_payload = body + return { + "data": { + "id": body["operation_id"], + "client_id": "client-123", + "artifact_id": None, + "type": "support_bundle", + "status": "running", + "created_at": "2026-01-01T00:00:00Z", + } + } + + @app.post("/api/v2/clients/management/artifacts") + async def create_artifact_upload(body: dict): + app.state.artifact_create_payload = body + return { + "data": { + "artifact_id": "artifact-123", + "client_id": "client-123", + "storage_key": "client-123--openhound-faker_support_bundle_2026-01-01_00-00-00.zip", + "status": "pending", + "part_size": body["part_size"], + "part_count": body["part_count"], + "missing_parts": list(range(1, body["part_count"] + 1)), + "management_operation": { + "id": body["operation_id"], + "client_id": "client-123", + "artifact_id": "artifact-123", + "type": "support_bundle", + "status": "running", + "requested_by_user_id": None, + "created_at": "2026-01-01T00:00:00Z", + "updated_at": "2026-01-01T00:00:00Z", + "started_at": "2026-01-01T00:00:00Z", + "completed_at": None, + "execution_time": "2026-01-01T00:00:00Z", + }, + } + } + + @app.post("/api/v2/clients/management/artifacts/{artifact_id}/parts/{part_number}") + async def upload_artifact_part( + artifact_id: str, part_number: int, request: Request + ): + content = await request.body() + checksum = base64.b64encode(hashlib.sha256(content).digest()).decode("ascii") + assert request.headers["content-digest"] == f"sha-256=:{checksum}:" + app.state.uploaded_parts.append((artifact_id, part_number, content)) + return Response(status_code=200) + + @app.post("/api/v2/clients/management/artifacts/{artifact_id}/complete") + async def complete_artifact_upload(artifact_id: str, body: dict): + app.state.artifact_completed = body["operation_id"] is not None + # BHE completes the associated management operation as part of this endpoint. + app.state.operation_completed_by_artifact_upload = True + return Response(status_code=204) + + @app.post("/api/v2/clients/management/end") + async def end_operation(body: dict): + app.state.operation_ended = True + app.state.operation_end_payload = body + return { + "data": { + "id": body["operation_id"], + "client_id": "client-123", + "artifact_id": "artifact-123", + "type": "support_bundle", + "status": body["status"], + "created_at": "2026-01-01T00:00:00Z", + } + } + return TestClient(app) @@ -385,3 +483,228 @@ def test_scheduler_ingest_opengraph(mock_service, mock_bloodhound_api, monkeypat assert result.job_id == 123 assert mock_bloodhound_api.app.state.ingested_nodes == 1000 assert mock_bloodhound_api.app.state.ingested_edges == 10000 + + +def _support_bundle_operation() -> dict: + return json.loads( + (MANAGEMENT_DATA_DIR / "management_available_with_operation.json").read_text() + )["data"][0] + + +def test_check_management_returns_support_bundle_operation( + mock_service, mock_bloodhound_api +): + mock_bloodhound_api.app.state.management_operations = [_support_bundle_operation()] + + operation = mock_service.check_management() + + assert operation is not None + assert operation.type is ManagementOperationType.SUPPORT_BUNDLE + + +def test_check_management_ignores_non_queued_operations( + mock_service, mock_bloodhound_api +): + operation = _support_bundle_operation() + operation["status"] = ManagementOperationStatus.RUNNING.value + mock_bloodhound_api.app.state.management_operations = [operation] + + assert mock_service.check_management() is None + + +def test_poll_prioritizes_management_over_a_new_job( + mock_service, mock_bloodhound_api, monkeypatch +): + mock_bloodhound_api.app.state.management_operations = [_support_bundle_operation()] + sent = [] + monkeypatch.setattr(mock_service, "_send_support_bundle", sent.append) + + mock_service._poll() + + assert len(sent) == 1 + assert mock_bloodhound_api.app.state.job_started is False + + +def test_poll_starts_job_when_no_management_work( + mock_service, mock_bloodhound_api, monkeypatch +): + submitted = Future() + monkeypatch.setattr(mock_service.executor, "submit", lambda *args: submitted) + + mock_service._poll() + + assert mock_bloodhound_api.app.state.job_started is True + + +def test_poll_still_checks_jobs_when_management_endpoint_fails( + mock_service, mock_bloodhound_api, monkeypatch +): + submitted = Future() + monkeypatch.setattr(mock_service, "check_management", lambda: 1 / 0) + monkeypatch.setattr(mock_service.executor, "submit", lambda *args: submitted) + + mock_service._poll() + + assert mock_bloodhound_api.app.state.job_started is True + + +def test_poll_does_not_start_a_job_when_management_work_fails( + mock_service, mock_bloodhound_api, monkeypatch +): + mock_bloodhound_api.app.state.management_operations = [_support_bundle_operation()] + + def fail(operation): + raise RuntimeError("upload failed") + + monkeypatch.setattr(mock_service, "_send_support_bundle", fail) + + mock_service._poll() + + assert mock_bloodhound_api.app.state.job_started is False + + +def test_send_support_bundle_claims_uploads_completes_and_cleans_up( + mock_service, mock_bloodhound_api, tmp_path, monkeypatch +): + log = tmp_path / "openhound.log" + log.write_text("support log") + mock_service.log_base_path = tmp_path + created = [] + + from openhound.scheduler import service as scheduler_service + + original_create = scheduler_service.create_support_bundle + + def capture_bundle(*args): + bundle = original_create(*args) + created.append(bundle) + return bundle + + monkeypatch.setattr(scheduler_service, "create_support_bundle", capture_bundle) + operation = ManagementOperation.model_validate(_support_bundle_operation()) + + mock_service._send_support_bundle(operation) + + assert mock_bloodhound_api.app.state.operation_start_payload == { + "operation_id": operation.id + } + assert ( + mock_bloodhound_api.app.state.artifact_create_payload["operation_id"] + == operation.id + ) + assert mock_bloodhound_api.app.state.uploaded_parts + assert mock_bloodhound_api.app.state.artifact_completed is True + assert mock_bloodhound_api.app.state.operation_completed_by_artifact_upload is True + assert mock_bloodhound_api.app.state.operation_end_payload is None + assert created and not created[0].exists() + assert not created[0].parent.exists() + + +def test_create_artifact_upload_preserves_entire_create_response( + mock_service, tmp_path +): + bundle = tmp_path / "support-bundle.zip" + bundle.write_bytes(b"support bundle") + + session = mock_service.client.create_artifact_upload("operation-123", bundle) + + assert session.artifact_id == "artifact-123" + assert session.client_id == "client-123" + assert session.storage_key.endswith("support_bundle_2026-01-01_00-00-00.zip") + assert session.status == "pending" + assert session.missing_parts == [1] + assert session.management_operation.id == "operation-123" + assert session.management_operation.artifact_id == session.artifact_id + + +@pytest.mark.parametrize( + "failure_point", + [ + "start_operation", + "create_support_bundle", + "create_artifact_upload", + "upload_artifact_part", + "complete_artifact_upload", + ], +) +def test_send_support_bundle_marks_operation_failed_for_each_lifecycle_failure( + mock_service, mock_bloodhound_api, monkeypatch, caplog, failure_point +): + operation = ManagementOperation.model_validate(_support_bundle_operation()) + + def fail(*args, **kwargs): + raise RuntimeError(f"{failure_point} failed") + + if failure_point == "create_support_bundle": + monkeypatch.setattr(scheduler_service, "create_support_bundle", fail) + else: + monkeypatch.setattr(mock_service.client, failure_point, fail) + + with pytest.raises(RuntimeError, match=f"{failure_point} failed"): + mock_service._send_support_bundle(operation) + + assert mock_bloodhound_api.app.state.operation_end_payload == { + "operation_id": operation.id, + "status": ManagementOperationStatus.FAILED.value, + } + assert "Support bundle operation" in caplog.text + + +def test_send_support_bundle_retries_transient_part_upload_failure( + mock_service, monkeypatch, tmp_path +): + log = tmp_path / "openhound.log" + log.write_text("support log") + mock_service.log_base_path = tmp_path + operation = ManagementOperation.model_validate(_support_bundle_operation()) + original_request = mock_service.client.request + attempts = 0 + delays = [] + + def flaky_request(method, path, **kwargs): + nonlocal attempts + if "/parts/" in path: + attempts += 1 + if attempts < 3: + raise BloodHoundHTTPError("temporary failure", 503) + return original_request(method, path, **kwargs) + + monkeypatch.setattr(mock_service.client, "request", flaky_request) + monkeypatch.setattr(bloodhound_enterprise.time, "sleep", delays.append) + + mock_service._send_support_bundle(operation) + + assert attempts == 3 + assert delays == [2, 2] + + +def test_send_support_bundle_fails_after_transient_retries_are_exhausted( + mock_service, mock_bloodhound_api, monkeypatch, tmp_path +): + log = tmp_path / "openhound.log" + log.write_text("support log") + mock_service.log_base_path = tmp_path + operation = ManagementOperation.model_validate(_support_bundle_operation()) + original_request = mock_service.client.request + attempts = 0 + delays = [] + + def unavailable_part_upload(method, path, **kwargs): + nonlocal attempts + if "/parts/" in path: + attempts += 1 + raise BloodHoundHTTPError("temporarily unavailable", 503) + return original_request(method, path, **kwargs) + + monkeypatch.setattr(mock_service.client, "request", unavailable_part_upload) + monkeypatch.setattr(bloodhound_enterprise.time, "sleep", delays.append) + + with pytest.raises(BloodHoundHTTPError): + mock_service._send_support_bundle(operation) + + assert attempts == 4 + assert delays == [2, 2, 2] + assert mock_bloodhound_api.app.state.operation_end_payload == { + "operation_id": operation.id, + "status": ManagementOperationStatus.FAILED.value, + } diff --git a/tests/test_data/api/management/management_available_empty.json b/tests/test_data/api/management/management_available_empty.json new file mode 100644 index 00000000..268c73f0 --- /dev/null +++ b/tests/test_data/api/management/management_available_empty.json @@ -0,0 +1,3 @@ +{ + "data": [] +} diff --git a/tests/test_data/api/management/management_available_with_operation.json b/tests/test_data/api/management/management_available_with_operation.json new file mode 100644 index 00000000..ab090a16 --- /dev/null +++ b/tests/test_data/api/management/management_available_with_operation.json @@ -0,0 +1,12 @@ +{ + "data": [ + { + "id": "a1b2c3d4-e5f6-7890-abcd-ef1234567890", + "client_id": "client-123", + "artifact_id": null, + "type": "support_bundle", + "status": "queued", + "created_at": "2026-01-01T00:00:00Z" + } + ] +} diff --git a/tests/test_support_bundle.py b/tests/test_support_bundle.py new file mode 100644 index 00000000..b1861ed9 --- /dev/null +++ b/tests/test_support_bundle.py @@ -0,0 +1,32 @@ +import zipfile +from pathlib import Path + +from openhound.core.support_bundle import collect_log_files, create_support_bundle + + +def test_collect_log_files_includes_current_and_rotated_platform_and_extension_logs( + tmp_path: Path, +): + expected = [ + tmp_path / "openhound.log", + tmp_path / "openhound.log.2026-05-28_10", + tmp_path / "ext_faker.log", + ] + for path in expected: + path.write_text("log") + (tmp_path / "unrelated.txt").write_text("ignore") + + assert set(collect_log_files(tmp_path)) == set(expected) + + +def test_create_support_bundle_contains_collected_logs(tmp_path: Path): + log = tmp_path / "openhound.log" + log.write_text("log") + + bundle = create_support_bundle("openhound-faker", tmp_path) + try: + assert bundle.name.startswith("openhound-faker_support_bundle_") + with zipfile.ZipFile(bundle) as archive: + assert archive.namelist() == [log.name] + finally: + bundle.unlink(missing_ok=True)