BED-7975 - Add Support Bundle upload functionality

.gitignore

@@ -4,6 +4,7 @@ ui_catalog.db
 site/
 addons/
 collectors/
+.worktrees/
 
 notebooks
 

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

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}}"

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

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:

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