diff --git a/README.md b/README.md index 0db2c437..0639e3e9 100644 --- a/README.md +++ b/README.md @@ -52,6 +52,137 @@ Follow the docs for setup, CLI usage, and collector development: - **Convert**: The raw resources are read from disk and converted to OpenGraph nodes and edges. +## Uploading to BloodHound + +BloodHound needs two things from a collector: the extension **schema**, which describes the node and +relationship kinds, and the **collection**, which is the data itself. OpenHound can send both. + +### Configure where BloodHound is + +Put the BloodHound address in `.dlt/config.toml` and the credentials in `.dlt/secrets.toml`, under the +`destination.bloodhound` section. This is the same section the saved-search and privilege-zone uploads use. +Use an API token pair (created under *My Profile → API Key Management* in BloodHound) or a bearer token from +a signed-in session. + +```toml +# .dlt/config.toml +[destination.bloodhound] +url = "http://localhost:8080" +``` + +```toml +# .dlt/secrets.toml, with an API token pair... +[destination.bloodhound] +token_id = "..." +token_key = "..." + +# ...or with a bearer token instead +[destination.bloodhound] +token = "..." +``` + +As environment variables these are `DESTINATION__BLOODHOUND__URL`, `DESTINATION__BLOODHOUND__TOKEN_ID`, +`DESTINATION__BLOODHOUND__TOKEN_KEY`, and `DESTINATION__BLOODHOUND__TOKEN`. The token pair is used when both +halves are present; otherwise the bearer token is used. + +When `destination.bloodhound` holds no credentials, the upload commands read the +`destination.bloodhoundenterprise` section, the one Enterprise ingest uses. `convert --method ingest` to +Enterprise reads `destination.bloodhoundenterprise`. + +### Install or update a schema + +```console +# One schema file, or a zip bundle containing schema.json +openhound schema upload ./schema.json +openhound schema upload ./my-extension-bundle.zip + +# Every .json and .zip file in a folder +openhound schema upload ./schemas + +# The schema.json that an installed collector ships next to its extension.yaml +openhound schema upload sccm + +# Check the files without contacting BloodHound +openhound schema upload ./schemas --dry-run +``` + +The command prints one line per file (`created`, `updated`, `skipped`, or `failed`) and exits with +status 1 when any file failed. BloodHound matches extensions by the `schema.name` field, so uploading a +changed file with the same name updates the existing extension. + +Collectors that want `openhound schema upload ` to work should ship their schema as `schema.json` +in the same package folder as `extension.yaml`. + +### Upload a collection + +After `openhound convert` has written OpenGraph JSON files, send them to BloodHound: + +```console +# An OpenHound convert output folder (every .json and .zip directly inside it) +openhound graph upload ./out/graph + +# A single OpenGraph JSON file, or a zip of them such as ConfigManBearPig's bundle +openhound graph upload ./out/graph/users-1.json +openhound graph upload ./out/configmanbearpig_collection.zip + +# Check the files without contacting BloodHound +openhound graph upload ./out/graph --dry-run + +# Return as soon as the files are sent instead of waiting for BloodHound to process them +openhound graph upload ./out/graph --no-wait + +# Cut large JSON files into parts of at most 2,000 nodes plus edges (default 5,000; 0 sends each file whole) +openhound graph upload ./out/graph --batch-size 2000 + +# Send JSON files uncompressed (they are gzip-compressed by default) +openhound graph upload ./out/graph --no-compress +``` + +The command opens one BloodHound file-upload job, attaches every file to it, and ends the job. It +prints one line per file (`uploaded`, `skipped`, or `failed`) and the job number. By default it then +waits, for up to ten minutes (`--timeout`), until BloodHound reports the job complete and prints the +result. It exits with status 1 when any file failed or the job did not complete. Progress is also +visible under *Administration → File Ingest* in BloodHound. + +Each JSON file is checked locally before anything is sent: it must parse and must have a top-level +`graph` section. JSON files are gzip-compressed before sending, which makes large collections much +smaller on the wire; BloodHound unpacks them on arrival. Zip files are already compressed, so they are +sent as they are and checked by BloodHound. + +A JSON file with more than `--batch-size` nodes plus edges (default 5,000) is cut into parts that are +all attached to the same job. Each part keeps the file's `metadata` section, and all node parts are sent +before any edge part. Parts are named `.part-0001.json`, `.part-0002.json`, and so on. Zip +files cannot be cut and are always sent whole. + +### Convert straight into BloodHound + +`openhound convert` can send its output to BloodHound instead of writing files, with `--method ingest`: + +```console +openhound convert ./collected/ ./unused --method ingest +``` + +The output folder argument is required but not used with `--method ingest`; nothing is written to it. +The default, `--method write`, saves OpenGraph JSON files to the output folder. + +OpenHound asks the configured BloodHound which edition it is and sends the data the way that edition +expects: + +- **Community Edition**: the settings come from `destination.bloodhound`, or from + `destination.bloodhoundenterprise` when `destination.bloodhound` holds no credentials (see + [Configure where BloodHound is](#configure-where-bloodhound-is)). The run opens one BloodHound + file-upload job, attaches every batch of 1,000 converted items to it as a gzip-compressed OpenGraph + file, and ends the job. Each file carries the collector's source kind in its `metadata` section, so + the data can later be removed by source kind under *Administration → Database Management*. The command + does not wait for BloodHound to process the job; watch it under *Administration → File Ingest*. + BloodHound runs its analysis once, after the job ends. If the conversion fails part-way, the job is + ended anyway, so BloodHound processes whatever arrived. +- **Enterprise**: the data goes to `/api/v2/ingest` using the `destination.bloodhoundenterprise` + settings. OpenHound also takes this route when the edition cannot be determined: neither section + holds credentials, or BloodHound cannot be reached or does not answer the question. + +The Enterprise scheduler service always takes the Enterprise route. + ## Available extensions Extend OpenHound with pre-built extensions for other services. Additional collectors can be installed diff --git a/src/openhound/cli/graph.py b/src/openhound/cli/graph.py new file mode 100644 index 00000000..e5b373fc --- /dev/null +++ b/src/openhound/cli/graph.py @@ -0,0 +1,179 @@ +"""The ``openhound graph`` command group: send collected OpenGraph data to BloodHound.""" + +import logging + +import requests +import typer + +from openhound.core.clients.bloodhound import BloodHoundHTTPError +from openhound.core.clients.from_secrets import ( + MissingBloodHoundCredentials, + bloodhound_client_from_secrets, +) +from openhound.core.clients.models import FileUploadJobStatus +from openhound.core.graph_upload import ( + DEFAULT_BATCH_SIZE, + CollectionNotFound, + find_collection_files, + upload_collection_files, + wait_for_job, +) + +logger = logging.getLogger(__name__) + +graph = typer.Typer(help="OpenGraph collection upload") + + +@graph.callback() +def graph_group(): + """Send collected OpenGraph data to BloodHound. + + This callback keeps ``graph`` a command group with named subcommands, such as + ``openhound graph upload``, even while it has only one. + """ + + +@graph.command(help="Upload collection files to BloodHound through a file-upload job") +def upload( + target: str = typer.Argument( + ..., + help=( + "An OpenHound convert output folder, a single OpenGraph JSON file, " + "or a zip of OpenGraph JSON files" + ), + ), + dry_run: bool = typer.Option( + False, + "--dry-run", + help="Check the files and show what would be uploaded without contacting BloodHound", + ), + compress: bool = typer.Option( + True, + "--compress/--no-compress", + help="Gzip JSON files before sending them (zip files are never recompressed)", + ), + batch_size: int = typer.Option( + DEFAULT_BATCH_SIZE, + "--batch-size", + help=( + "Most nodes plus edges per request; larger JSON files are cut into parts of " + "this size within the same job. Use 0 to send each file whole" + ), + ), + wait: bool = typer.Option( + True, + "--wait/--no-wait", + help="After uploading, wait until BloodHound has finished processing the job", + ), + timeout: float = typer.Option( + 600.0, + "--timeout", + help="How many seconds to wait for the job before giving up", + ), +): + """Find the files, attach them to one job, end the job, then optionally wait for it. + + The command exits with status 1 when any file failed, when the job could not be started + or ended, or when BloodHound reports the job failed. + """ + try: + files = find_collection_files(target) + except CollectionNotFound as error: + typer.secho(str(error), fg=typer.colors.RED, err=True) + raise typer.Exit(code=1) + + if dry_run: + # A dry run never needs credentials, so a missing configuration must not stop it. + client = None + else: + try: + client = bloodhound_client_from_secrets() + except MissingBloodHoundCredentials as error: + typer.secho(str(error), fg=typer.colors.RED, err=True) + raise typer.Exit(code=1) + + result = upload_collection_files( + client=client, # type: ignore[arg-type] + files=files, + dry_run=dry_run, + compress=compress, + batch_size=batch_size, + ) + + for file_result in result.files: + if file_result.succeeded: + note = f" ({file_result.detail})" if file_result.detail else "" + typer.secho( + f"{file_result.outcome:>8} {file_result.path}{note}", + fg=typer.colors.GREEN, + ) + else: + typer.secho( + f"{file_result.outcome:>8} {file_result.path}: {file_result.detail}", + fg=typer.colors.RED, + err=True, + ) + + failed = not result.all_files_succeeded + if result.detail: + failed = True + typer.secho(result.detail, fg=typer.colors.RED, err=True) + + if result.job_id is not None: + typer.echo(f"BloodHound file-upload job: {result.job_id}") + + # Wait only for a job that was ended: BloodHound never processes one left open, so + # waiting would only use up the whole --timeout before reporting the failure. + if wait and client is not None and result.job_ended: + try: + job = wait_for_job(client, result.job_id, timeout_seconds=timeout) + except (BloodHoundHTTPError, requests.RequestException) as error: + failed = True + reason = ( + error.reason + if isinstance(error, BloodHoundHTTPError) + else str(error) + ) + typer.secho( + f"Lost contact with BloodHound while waiting for job " + f"{result.job_id}: {reason}. Check Administration > File Ingest.", + fg=typer.colors.RED, + err=True, + ) + raise typer.Exit(code=1) + if job is None: + failed = True + typer.secho( + f"BloodHound no longer lists job {result.job_id}.", + fg=typer.colors.RED, + err=True, + ) + else: + summary = f"Job {job.id} is {job.status_text}" + if job.total_files: + summary += ( + f" ({job.failed_files} of {job.total_files} files failed)" + ) + if job.status_message: + summary += f": {job.status_message}" + job_failed = job.status_code in ( + FileUploadJobStatus.FAILED, + FileUploadJobStatus.CANCELED, + FileUploadJobStatus.TIMED_OUT, + FileUploadJobStatus.PARTIALLY_COMPLETE, + ) + still_running = ( + job.status_code is None or not job.status_code.is_finished + ) + if job_failed or still_running: + failed = True + typer.secho(summary, fg=typer.colors.RED, err=True) + else: + typer.secho(summary, fg=typer.colors.GREEN) + elif result.job_ended: + typer.echo( + "Watch its progress under Administration > File Ingest in BloodHound." + ) + + if failed: + raise typer.Exit(code=1) diff --git a/src/openhound/cli/schema.py b/src/openhound/cli/schema.py new file mode 100644 index 00000000..a35c7890 --- /dev/null +++ b/src/openhound/cli/schema.py @@ -0,0 +1,82 @@ +"""The ``openhound schema`` command group: install extension schemas into BloodHound.""" + +import logging + +import typer + +from openhound.core.clients.from_secrets import ( + MissingBloodHoundCredentials, + bloodhound_client_from_secrets, +) +from openhound.core.schema_upload import ( + SchemaNotFound, + find_schema_files, + upload_schema_files, +) + +logger = logging.getLogger(__name__) + +schema = typer.Typer(help="OpenGraph extension schema management") + + +@schema.callback() +def schema_group(): + """Manage OpenGraph extension schemas in BloodHound. + + This callback exists so that ``schema`` behaves as a command group with named + subcommands, such as ``openhound schema upload``, even while it has only one. Without + it, Typer would collapse the group and the subcommand name would be lost. + """ + + +@schema.command(help="Install or update extension schemas in BloodHound") +def upload( + target: str = typer.Argument( + ..., + help=( + "A schema file (.json or .zip bundle), a folder of schema files, " + "or the name of an installed collector that ships a schema.json" + ), + ), + dry_run: bool = typer.Option( + False, + "--dry-run", + help="Check the files and show what would be uploaded without contacting BloodHound", + ), +): + """Find the schema files, upload each one, and print a line per file. + + The command exits with status 1 when any file failed, so scripts can rely on it. + """ + try: + files = find_schema_files(target) + except SchemaNotFound as error: + typer.secho(str(error), fg=typer.colors.RED, err=True) + raise typer.Exit(code=1) + + if dry_run: + # A dry run never needs credentials, so a missing configuration must not stop it. + client = None + else: + try: + client = bloodhound_client_from_secrets() + except MissingBloodHoundCredentials as error: + typer.secho(str(error), fg=typer.colors.RED, err=True) + raise typer.Exit(code=1) + + results = upload_schema_files(client=client, files=files, dry_run=dry_run) # type: ignore[arg-type] + + any_failed = False + for result in results: + if result.succeeded: + typer.secho(f"{result.outcome:>8} {result.path}", fg=typer.colors.GREEN) + else: + any_failed = True + typer.secho( + f"{result.outcome:>8} {result.path}: {result.detail}", + fg=typer.colors.RED, + err=True, + ) + + if any_failed: + raise typer.Exit(code=1) diff --git a/src/openhound/core/app.py b/src/openhound/core/app.py index ba78ad04..b8dda37d 100644 --- a/src/openhound/core/app.py +++ b/src/openhound/core/app.py @@ -16,7 +16,9 @@ from openhound.cli.preproc import preprocess from openhound.core import validate from openhound.core.asset import BaseAsset, EdgeDef, NodeDef +from openhound.core.clients.bloodhound import BloodHoundClient from openhound.core.collect import CollectContext, Collector +from openhound.core.community_ingest import community_edition_client_for_ingest from openhound.core.convert import ConvertContext, Converter, Method from openhound.core.models.extension import Extension from openhound.core.preproc import PreProcContext, PreProcessor @@ -170,6 +172,7 @@ def run_convert( lookup_file: Path = DEFAULT_LOOKUP_FILE, progress: Progress = Progress.tqdm, method: Method = Method.write, + community_edition_client: BloodHoundClient | None = None, ) -> LoadInfo: lookup_session = None if lookup: @@ -187,6 +190,7 @@ def run_convert( lookup=lookup_session, progress=progress, method=method, + community_edition_client=community_edition_client, ) source_method, extra_context = func( @@ -229,13 +233,32 @@ def wrapper( help="DuckDB lookup file path", ), ] = DEFAULT_LOOKUP_FILE, + method: Method = typer.Option( + Method.write, + help=( + "write: save OpenGraph JSON files to the output path. " + "ingest: send the data straight to BloodHound. Community Edition " + "uses the destination.bloodhound settings, or " + "destination.bloodhoundenterprise when those hold no credentials; " + "Enterprise uses destination.bloodhoundenterprise. The output path " + "is then not used" + ), + ), ) -> LoadInfo: + # For ingest, find out whether the configured BloodHound is Community Edition. + # Only this command line does that; the Enterprise scheduler calls run_convert + # directly without a client, so it always keeps the Enterprise route. + community_edition_client = None + if method == Method.ingest: + community_edition_client = community_edition_client_for_ingest() + return run_convert( input_path=input_path, output_path=output_path, lookup_file=lookup_file, progress=progress, - method=Method.write, + method=method, + community_edition_client=community_edition_client, ) logger.debug(f"Registering convert command for {self.name}") diff --git a/src/openhound/core/clients/bloodhound.py b/src/openhound/core/clients/bloodhound.py index c33427bc..cbd557d0 100644 --- a/src/openhound/core/clients/bloodhound.py +++ b/src/openhound/core/clients/bloodhound.py @@ -16,6 +16,9 @@ from .models import ( AssetGroupsTags, CustomNodes, + FileUploadJob, + FileUploadJobResponse, + FileUploadJobsResponse, RunCypher, SavedQueries, SavedQuery, @@ -24,6 +27,19 @@ logger = logging.getLogger(__name__) +# Seconds to wait for BloodHound when installing an extension schema: the first number is for +# opening the connection, the second for each wait on the reply. Without them a stalled server +# would block `schema upload` forever. +EXTENSION_UPLOAD_CONNECT_TIMEOUT_SECONDS = 10 +EXTENSION_UPLOAD_READ_TIMEOUT_SECONDS = 120 + +# The same two kinds of limit, in seconds, for the edition question and the file-upload job calls +# used by `graph upload` and `convert --method ingest`. Attaching a file gets a longer reply limit, +# because BloodHound answers only after it has received and stored the whole file. Without these, +# a server that accepts the connection and then stops replying would block either command forever. +QUICK_REQUEST_TIMEOUT_SECONDS = (10, 60) +FILE_ATTACH_TIMEOUT_SECONDS = (10, 300) + class BloodHoundHTTPError(DltException): def __init__(self, reason: str, code: int): @@ -134,6 +150,109 @@ def custom_node(self, body: str) -> CustomNodes: response = self.request(method="POST", path=path, body=body.encode()) return CustomNodes.model_validate(response.json()) + def upsert_extension(self, body: bytes, content_type: str) -> bool: + """Install a new OpenGraph extension schema, or update the one with the same name. + + BloodHound calls this operation an "upsert": it creates the extension when the + schema name is new and updates it when the name already exists. The request body + is either the schema as JSON, or a zip bundle that contains a ``schema.json`` file + and optional ``saved_queries.json`` and ``pz_rules.json`` files. The content type + tells BloodHound which of the two it is receiving. + + Args: + body (bytes): The raw file contents to send. + content_type (str): Either ``application/json`` or ``application/zip``. + + Returns: + bool: True when an existing extension was updated, False when a new one was created. + """ + path = "/api/v2/extensions" + response = self.request( + method="PUT", + path=path, + body=body, + extra_headers={"Content-Type": content_type}, + timeout=( + EXTENSION_UPLOAD_CONNECT_TIMEOUT_SECONDS, + EXTENSION_UPLOAD_READ_TIMEOUT_SECONDS, + ), + ) + # BloodHound answers 200 (OK) for an update and 201 (Created) for a brand new extension. + return response.status_code == 200 + + def product_edition(self) -> str: + """Ask BloodHound which edition it is, for example ``community`` or ``enterprise``. + + The answer comes from the ``product_edition`` field of ``GET /api/version``, which + both editions provide to any signed-in caller. + """ + response = self.request( + method="GET", path="/api/version", timeout=QUICK_REQUEST_TIMEOUT_SECONDS + ) + return response.json()["data"]["product_edition"] + + def start_file_upload_job(self) -> FileUploadJob: + """Open a new file-upload job and return it. + + Files are attached to the job with :meth:`upload_file_to_job`, and the job is closed + with :meth:`end_file_upload_job`. BloodHound only starts reading the files after the + job is ended. + """ + path = "/api/v2/file-upload/start" + response = self.request( + method="POST", path=path, timeout=QUICK_REQUEST_TIMEOUT_SECONDS + ) + return FileUploadJobResponse.model_validate(response.json()).data + + def upload_file_to_job( + self, + job_id: int, + body: bytes, + content_type: str, + file_name: str, + content_encoding: str | None = None, + ) -> None: + """Attach one collection file to an open job. + + Args: + job_id (int): The job returned by :meth:`start_file_upload_job`. + body (bytes): The raw file contents, either an OpenGraph JSON document or a zip + of such documents. + content_type (str): ``application/json`` or ``application/zip``. + file_name (str): The original file name, shown in BloodHound's file ingest page. + content_encoding (str | None): ``gzip`` when ``body`` is gzip-compressed, so + BloodHound unpacks it before reading. ``content_type`` still names the + uncompressed file. + """ + path = f"/api/v2/file-upload/{job_id}" + extra_headers = { + "Content-Type": content_type, + "X-File-Upload-Name": file_name, + } + if content_encoding: + extra_headers["Content-Encoding"] = content_encoding + self.request( + method="POST", + path=path, + body=body, + extra_headers=extra_headers, + timeout=FILE_ATTACH_TIMEOUT_SECONDS, + ) + + def end_file_upload_job(self, job_id: int) -> None: + """Close a job so BloodHound starts reading the attached files.""" + path = f"/api/v2/file-upload/{job_id}/end" + self.request(method="POST", path=path, timeout=QUICK_REQUEST_TIMEOUT_SECONDS) + + def file_upload_job(self, job_id: int) -> FileUploadJob | None: + """Fetch the current state of one job, or None when BloodHound no longer lists it.""" + path = f"/api/v2/file-upload?id=eq:{job_id}" + response = self.request( + method="GET", path=path, timeout=QUICK_REQUEST_TIMEOUT_SECONDS + ) + jobs = FileUploadJobsResponse.model_validate(response.json()).data + return jobs[0] if jobs else None + @property def asset_group_tags(self) -> AssetGroupsTags: path = "/api/v2/asset-group-tags" diff --git a/src/openhound/core/clients/from_secrets.py b/src/openhound/core/clients/from_secrets.py new file mode 100644 index 00000000..292c5cef --- /dev/null +++ b/src/openhound/core/clients/from_secrets.py @@ -0,0 +1,159 @@ +"""Build a BloodHound API client from the settings stored in OpenHound's configuration. + +OpenHound keeps settings in dlt's two stores. Plain settings such as the BloodHound address +live in ``.dlt/config.toml``; credentials live in ``.dlt/secrets.toml``. Both can also be +given as environment variables. This module looks in two sections, in this order, and uses +the first one that holds credentials: + +``[destination.bloodhound]`` + The section that OpenHound's saved-search and privilege-zone uploads already use. It is + the home for BloodHound Community Edition settings. Holds ``url`` plus either an API token + pair (``token_id`` and ``token_key``, created under "My Profile > API Key Management") or + a bearer ``token`` from a signed-in session. + +``[destination.bloodhoundenterprise]`` + The section that OpenHound's Enterprise ingest and scheduler already use. Holds ``url`` + plus ``token_id`` and ``token_key``. Nothing about it changes; it is read here so that an + Enterprise user needs no extra configuration. + +Example, with the address in ``config.toml``:: + + [destination.bloodhound] + url = "http://localhost:8080" + +and the credentials in ``secrets.toml``:: + + [destination.bloodhound] + token_id = "..." + token_key = "..." + +As environment variables the same settings are ``DESTINATION__BLOODHOUND__URL``, +``DESTINATION__BLOODHOUND__TOKEN_ID``, and so on. +""" + +from dataclasses import dataclass + +import dlt + +from openhound.core.clients.bloodhound import ( + BloodHound, + BloodHoundClient, + BloodHoundJWT, +) + +# The configuration sections that may hold BloodHound settings, in the order they are tried. +SECTION_BLOODHOUND = "destination.bloodhound" +SECTION_ENTERPRISE = "destination.bloodhoundenterprise" +SECTIONS_IN_ORDER = (SECTION_BLOODHOUND, SECTION_ENTERPRISE) + +# The setting names inside each section. +SETTING_URL = "url" +SETTING_TOKEN_ID = "token_id" +SETTING_TOKEN_KEY = "token_key" +SETTING_BEARER_TOKEN = "token" + + +class MissingBloodHoundCredentials(Exception): + """Raised when the configuration does not contain enough information to reach BloodHound.""" + + +@dataclass +class BloodHoundSettings: + """The settings found in one configuration section. + + Attributes: + section: The section they came from, for example ``destination.bloodhound``. + url: The BloodHound address, or None when absent. + token_id: First half of an API token pair, or None. + token_key: Second half of an API token pair, or None. + bearer_token: A session token, or None. + """ + + section: str + url: str | None + token_id: str | None + token_key: str | None + bearer_token: str | None + + @property + def has_token_pair(self) -> bool: + return bool(self.token_id and self.token_key) + + @property + def has_credentials(self) -> bool: + return self.has_token_pair or bool(self.bearer_token) + + +def read_settings(section: str) -> BloodHoundSettings: + """Read one section's settings from the two stores. + + The address is read through the configuration store because that store sees + ``config.toml``, ``secrets.toml``, and environment variables. Credentials are read through + the secrets store, which sees ``secrets.toml`` and environment variables only, so a token + accidentally placed in ``config.toml`` is not picked up. + """ + url: str | None = dlt.config.get(f"{section}.{SETTING_URL}") + token_id: str | None = dlt.secrets.get(f"{section}.{SETTING_TOKEN_ID}") + token_key: str | None = dlt.secrets.get(f"{section}.{SETTING_TOKEN_KEY}") + bearer_token: str | None = dlt.secrets.get(f"{section}.{SETTING_BEARER_TOKEN}") + return BloodHoundSettings( + section=section, + url=url, + token_id=token_id, + token_key=token_key, + bearer_token=bearer_token, + ) + + +def bloodhound_client_from_secrets() -> BloodHoundClient: + """Create the right kind of BloodHound client from whichever section holds credentials. + + The sections are tried in the order of ``SECTIONS_IN_ORDER``. Within a section the API + token pair is preferred, because API tokens do not expire the way session tokens do. When + no section holds credentials, or the chosen section has no address, a plain-English error + names exactly which settings are needed. + + Returns: + BloodHoundClient: A client that signs every request with the configured credentials. + """ + for section in SECTIONS_IN_ORDER: + settings = read_settings(section) + if not settings.has_credentials: + continue + + if not settings.url: + raise MissingBloodHoundCredentials( + f"Credentials were found in [{section}] but it has no 'url'. Add " + f"'{section}.url' to .dlt/config.toml, or set the " + f"{environment_variable_name(section, SETTING_URL)} environment variable." + ) + + if settings.has_token_pair: + # The checks above guarantee both halves are present; the assertions tell the + # type checker the same thing. + assert settings.token_id is not None and settings.token_key is not None + return BloodHound( + token_key=settings.token_key, + token_id=settings.token_id, + bhe_uri=settings.url, + ) + + assert settings.bearer_token is not None + return BloodHoundJWT(token=settings.bearer_token, base_uri=settings.url) + + raise MissingBloodHoundCredentials( + "No BloodHound credentials are configured. Add a " + f"[{SECTION_BLOODHOUND}] section with '{SETTING_URL}' in .dlt/config.toml and either " + f"'{SETTING_TOKEN_ID}' and '{SETTING_TOKEN_KEY}' or '{SETTING_BEARER_TOKEN}' in " + f".dlt/secrets.toml. An existing [{SECTION_ENTERPRISE}] section is also accepted. " + "The same settings can be environment variables such as " + f"{environment_variable_name(SECTION_BLOODHOUND, SETTING_URL)}." + ) + + +def environment_variable_name(section: str, setting: str) -> str: + """Turn ``destination.bloodhound`` and ``url`` into ``DESTINATION__BLOODHOUND__URL``. + + This is dlt's naming rule: dots become double underscores and everything is upper case. + """ + return f"{section}.{setting}".replace(".", "__").upper() diff --git a/src/openhound/core/clients/models/__init__.py b/src/openhound/core/clients/models/__init__.py index b200f2e9..bc031ac2 100644 --- a/src/openhound/core/clients/models/__init__.py +++ b/src/openhound/core/clients/models/__init__.py @@ -1,13 +1,23 @@ from .asset_groups import AssetGroupsTags, Selectors from .custom_nodes import CustomNodes +from .file_upload import ( + FileUploadJob, + FileUploadJobResponse, + FileUploadJobsResponse, + FileUploadJobStatus, +) from .graph_cypher import RunCypher from .saved_query import SavedQueries, SavedQuery __all__ = [ + "AssetGroupsTags", "CustomNodes", + "FileUploadJob", + "FileUploadJobResponse", + "FileUploadJobStatus", + "FileUploadJobsResponse", "RunCypher", + "SavedQueries", "SavedQuery", - "AssetGroupsTags", "Selectors", - "SavedQueries", ] diff --git a/src/openhound/core/clients/models/file_upload.py b/src/openhound/core/clients/models/file_upload.py new file mode 100644 index 00000000..33c0568e --- /dev/null +++ b/src/openhound/core/clients/models/file_upload.py @@ -0,0 +1,83 @@ +"""Response shapes for BloodHound's file-upload job endpoints. + +A file-upload job is how BloodHound Community Edition and Enterprise both accept collected +data through the API. The client starts a job, attaches one or more files to it, and ends +it. BloodHound then reads the files ("ingesting"), runs its analysis, and marks the job +complete or failed. These models describe the job record that BloodHound returns. +""" + +from enum import IntEnum + +from pydantic import BaseModel + + +class FileUploadJobStatus(IntEnum): + """The numeric status codes BloodHound uses for a file-upload job, with readable names.""" + + INVALID = -1 + READY = 0 + RUNNING = 1 + COMPLETE = 2 + CANCELED = 3 + TIMED_OUT = 4 + FAILED = 5 + INGESTING = 6 + ANALYZING = 7 + PARTIALLY_COMPLETE = 8 + + @property + def is_finished(self) -> bool: + """True once BloodHound will not change this job's status any further.""" + return self in ( + FileUploadJobStatus.COMPLETE, + FileUploadJobStatus.CANCELED, + FileUploadJobStatus.TIMED_OUT, + FileUploadJobStatus.FAILED, + FileUploadJobStatus.PARTIALLY_COMPLETE, + ) + + @property + def readable(self) -> str: + """The status as lower-case words, for messages to people.""" + return self.name.lower().replace("_", " ") + + +class FileUploadJob(BaseModel): + """One file-upload job as BloodHound reports it. + + Only the fields this project reads are listed; BloodHound returns more and they are + ignored. ``status`` is kept as the raw number so an unknown future value cannot break + parsing; use ``status_code`` for the readable form. + """ + + id: int + status: int + status_message: str = "" + total_files: int = 0 + failed_files: int = 0 + + @property + def status_code(self) -> FileUploadJobStatus | None: + """The status as a named value, or None when BloodHound sent a number we do not know.""" + try: + return FileUploadJobStatus(self.status) + except ValueError: + return None + + @property + def status_text(self) -> str: + """A readable status for messages, falling back to the raw number.""" + code = self.status_code + return code.readable if code is not None else f"status {self.status}" + + +class FileUploadJobResponse(BaseModel): + """The ``{"data": {...}}`` wrapper BloodHound puts around one job.""" + + data: FileUploadJob + + +class FileUploadJobsResponse(BaseModel): + """The ``{"data": [...]}`` wrapper BloodHound puts around a list of jobs.""" + + data: list[FileUploadJob] diff --git a/src/openhound/core/community_ingest.py b/src/openhound/core/community_ingest.py new file mode 100644 index 00000000..1e9aaa99 --- /dev/null +++ b/src/openhound/core/community_ingest.py @@ -0,0 +1,155 @@ +"""Send ``openhound convert --method ingest`` output to BloodHound Community Edition. + +BloodHound Enterprise receives converted data through its ``/api/v2/ingest`` endpoint, which +OpenHound's Enterprise destination already uses. Community Edition does not have that endpoint. +It accepts data only through file-upload jobs, the same mechanism its File Ingest page uses: + +1. start a job, +2. attach files to it, +3. end the job, after which BloodHound reads the files and runs its analysis. + +This module decides which of the two routes a ``convert --method ingest`` run takes, and provides +the destination that attaches each batch of converted data to the job. The job itself is started +and ended by :class:`openhound.core.convert.Converter`, because a dlt destination is called once +per batch and has no "before the first batch" or "after the last batch" step. + +Only the ``convert`` command line uses this module. The Enterprise scheduler service calls the +converter directly and always takes the unchanged Enterprise route. +""" + +import gzip +import itertools +import logging +import threading + +import dlt +import requests +from dlt.common import json +from dlt.common.destination import Destination +from dlt.common.schema import TTableSchema +from dlt.common.typing import TDataItems + +from openhound.core.clients.bloodhound import BloodHoundClient, BloodHoundHTTPError +from openhound.core.clients.from_secrets import ( + MissingBloodHoundCredentials, + bloodhound_client_from_secrets, +) +from openhound.core.graph_upload import CONTENT_TYPE_JSON + +logger = logging.getLogger(__name__) + +# The value of ``product_edition`` that BloodHound Community Edition reports from /api/version. +EDITION_COMMUNITY = "community" + +# How many converted items go into one attached file. This matches the Enterprise destination, +# so both editions receive the same amount of data per request. An item is one node or a list +# of up to 150 relationships, as produced by the OpenGraph source. +ITEMS_PER_FILE = 1000 + + +def community_edition_client_for_ingest() -> BloodHoundClient | None: + """Return a client when the configured BloodHound is Community Edition, otherwise None. + + The credentials come from ``[destination.bloodhound]``, falling back to + ``[destination.bloodhoundenterprise]`` (see :mod:`openhound.core.clients.from_secrets`). + BloodHound is then asked which edition it is. + + None means "use the Enterprise route exactly as before". That happens when no credentials + are configured, when the server says it is Enterprise, and when the question cannot be + answered at all. Falling back in the last case keeps every setup that worked before this + change working the same way. + """ + try: + client = bloodhound_client_from_secrets() + except MissingBloodHoundCredentials: + logger.debug("No BloodHound settings found; using the Enterprise ingest route") + return None + + try: + edition = client.product_edition() + except ( + BloodHoundHTTPError, + requests.RequestException, + KeyError, + ValueError, + ) as error: + logger.warning( + f"Could not ask BloodHound which edition it is ({error}); " + "using the Enterprise ingest route" + ) + return None + + if edition != EDITION_COMMUNITY: + logger.debug( + f"BloodHound reports edition '{edition}'; using the Enterprise route" + ) + return None + + logger.info( + "BloodHound Community Edition detected; sending data through a file-upload job" + ) + return client + + +def build_document(items: TDataItems, source_kind: str | None) -> dict: + """Turn one batch of converted items into an OpenGraph document. + + The ``metadata`` section names the source kind. BloodHound uses it to register the source + kind, which makes the data removable by source kind under Database Management, and to label + the imported objects. It is left out when the collector declares no source kind. + """ + nodes = [] + edges = [] + for item in items: + if item["graph"]["entity_type"] == "node": + nodes.append(item["graph"]["content"]) + elif item["graph"]["entity_type"] == "edge": + edges.extend(item["graph"]["content"]) + + document: dict = {"graph": {"nodes": nodes, "edges": edges}} + if source_kind: + document = {"metadata": {"source_kind": source_kind}, **document} + return document + + +def community_edition_destination( + client: BloodHoundClient, job_id: int, source_kind: str | None +) -> Destination: + """Create a dlt destination that attaches every batch to one Community Edition job. + + Each batch becomes one gzip-compressed OpenGraph JSON file named after its table and a + running number, for example ``users-0001.json``. That name is what BloodHound's File Ingest + page shows. dlt may deliver batches from several threads at once, so the running number is + protected by a lock. + + Args: + client (BloodHoundClient): A signed-in Community Edition client. + job_id (int): The file-upload job the converter started. + source_kind (str | None): The collector's source kind, written into each file's metadata. + """ + file_numbers = itertools.count(1) + file_number_lock = threading.Lock() + + @dlt.destination( + name="bloodhound_community", + skip_dlt_columns_and_tables=True, + batch_size=ITEMS_PER_FILE, + ) + def attach_to_job(items: TDataItems, table: TTableSchema) -> None: + with file_number_lock: + file_number = next(file_numbers) + file_name = f"{table['name']}-{file_number:04d}.json" + + body = gzip.compress(json.dumps(build_document(items, source_kind)).encode()) + logger.info( + f"Attaching {len(items)} items to BloodHound job {job_id} as {file_name}" + ) + client.upload_file_to_job( + job_id=job_id, + body=body, + content_type=CONTENT_TYPE_JSON, + file_name=file_name, + content_encoding="gzip", + ) + + return attach_to_job() diff --git a/src/openhound/core/convert.py b/src/openhound/core/convert.py index 5df942d5..747755de 100644 --- a/src/openhound/core/convert.py +++ b/src/openhound/core/convert.py @@ -5,11 +5,14 @@ from typing import Callable import dlt +import requests from dlt.common.pipeline import LoadInfo from dlt.extract.source import DltSource from dlt.pipeline.pipeline import Pipeline from openhound.core.asset import BaseAsset +from openhound.core.clients.bloodhound import BloodHoundClient, BloodHoundHTTPError +from openhound.core.community_ingest import community_edition_destination from openhound.core.logging import logger_override from openhound.core.lookup import LookupManager from openhound.core.pipeline import BasePipeline @@ -22,6 +25,13 @@ class Method(str, Enum): + """Where the converter sends its output. + + ``write`` saves OpenGraph JSON files to the output folder. ``ingest`` sends the data + straight to BloodHound: through ``/api/v2/ingest`` for Enterprise, or through one + file-upload job for Community Edition (see :mod:`openhound.core.community_ingest`). + """ + write = "write" ingest = "ingest" @@ -36,7 +46,14 @@ def __init__( source_kind: str | None = None, progress: Progress = Progress.tqdm, method: Method = Method.write, + community_edition_client: BloodHoundClient | None = None, ): + """ + Args: + community_edition_client: Set only by the ``convert`` command line when ``ingest`` + is chosen and the configured BloodHound is Community Edition. Without it, + ``ingest`` takes the Enterprise route, as the Enterprise scheduler always does. + """ self.name = name self.input_path = input_path self.output_path = output_path @@ -44,10 +61,31 @@ def __init__( self.lookup = lookup self.progress = progress self.method = method + self.community_edition_client = community_edition_client + # The Community Edition file-upload job for the current run, set by run(). + self.community_edition_job_id: int | None = None + + @property + def sends_to_community_edition(self) -> bool: + return ( + self.method == Method.ingest and self.community_edition_client is not None + ) @property def pipeline(self) -> Pipeline: - if self.method == Method.ingest: + if self.sends_to_community_edition: + assert self.community_edition_client is not None + assert self.community_edition_job_id is not None + logger.debug( + "Using BloodHound Community Edition file-upload job for ingest" + ) + dest = community_edition_destination( + client=self.community_edition_client, + job_id=self.community_edition_job_id, + source_kind=self.source_kind, + ) + + elif self.method == Method.ingest: logger.debug( "Initializing BloodHound Enterprise client for converter ingest method" ) @@ -91,19 +129,58 @@ def run( ) logger_override.set_handler(self.name) - result = self._run( - opengraph( - valid_resources, - lookup=self.lookup, - bucket_url=str(self.input_path), - extras=extra_context, - source_kind=self.source_kind, - ) + source = opengraph( + valid_resources, + lookup=self.lookup, + bucket_url=str(self.input_path), + extras=extra_context, + source_kind=self.source_kind, ) + if self.sends_to_community_edition: + result = self._run_into_community_edition_job(source) + else: + result = self._run(source) + logger.info(f"Converter run completed with method '{self.method}'") return result + def _run_into_community_edition_job(self, source: DltSource) -> LoadInfo: + """Run the conversion inside one Community Edition file-upload job. + + The job is opened before the first batch and always ended afterwards, even when the + conversion fails part-way. Ending it makes BloodHound process whatever arrived; an + open job would otherwise sit unfinished on the File Ingest page, because BloodHound + offers no way to cancel one. BloodHound runs its analysis once, after the job ends. + """ + assert self.community_edition_client is not None + job = self.community_edition_client.start_file_upload_job() + self.community_edition_job_id = job.id + logger.info(f"Started BloodHound file-upload job {job.id}") + + try: + result = self._run(source) + except BaseException: + # The conversion failed. Still end the job, but never let a second failure here + # hide the original error, which is the one the person needs to see. + try: + self.community_edition_client.end_file_upload_job(job.id) + logger.info( + f"Ended BloodHound file-upload job {job.id} after a failure" + ) + except (BloodHoundHTTPError, requests.RequestException) as end_error: + logger.error( + f"Could not end BloodHound file-upload job {job.id}: {end_error}" + ) + raise + + self.community_edition_client.end_file_upload_job(job.id) + logger.info( + f"Ended BloodHound file-upload job {job.id}. BloodHound now processes " + "the data; watch its progress under Administration > File Ingest." + ) + return result + @dataclass class ConvertContext: diff --git a/src/openhound/core/graph_upload.py b/src/openhound/core/graph_upload.py new file mode 100644 index 00000000..e8f8fada --- /dev/null +++ b/src/openhound/core/graph_upload.py @@ -0,0 +1,383 @@ +"""Find OpenGraph collection files and send them to BloodHound through a file-upload job. + +A collection is the data a collector produced: OpenGraph JSON documents, each holding nodes +and edges plus a ``metadata`` section naming the source kind. OpenHound's ``convert`` stage +writes a folder of such files; other collectors, such as ConfigManBearPig, write a zip of +them. BloodHound accepts both through its file-upload job endpoints, which work the same way +in Community Edition and Enterprise: + +1. start a job, +2. attach every file to it, +3. end the job, after which BloodHound reads the files and runs its analysis. + +This module does the finding and the sending; the ``openhound graph upload`` command in +:mod:`openhound.cli.graph` is the thin layer on top. +""" + +import gzip +import logging +from collections.abc import Iterator +from dataclasses import dataclass, field +from pathlib import Path + +# Imported by name so tests can replace only this module's clock and pause, without +# touching the time module that other libraries rely on. +from time import monotonic, sleep + +import requests +from dlt.common import json + +from openhound.core.clients.bloodhound import BloodHoundClient, BloodHoundHTTPError +from openhound.core.clients.models import FileUploadJob + +logger = logging.getLogger(__name__) + +# The content types BloodHound accepts for a collection file. +CONTENT_TYPE_JSON = "application/json" +CONTENT_TYPE_ZIP = "application/zip" + +# The most nodes plus edges sent in one request by default. Collections can hold millions of +# entities; sending them in bounded parts keeps each request small enough to succeed, lets a +# failure cost one part instead of the whole file, and lets BloodHound start working earlier. +DEFAULT_BATCH_SIZE = 5000 + +# How often to ask BloodHound about a job while waiting for it to finish, in seconds. +DEFAULT_POLL_SECONDS = 2.0 + + +class CollectionNotFound(Exception): + """Raised when the path does not lead to any collection file.""" + + +class InvalidCollectionFile(Exception): + """Raised when a collection file cannot be read or is not an OpenGraph JSON document.""" + + +@dataclass +class FileResult: + """What happened to one collection file. + + Attributes: + path: The file. + outcome: "uploaded", "skipped" (dry run), or "failed". + detail: Extra explanation, used for failures and for files sent in several parts. + """ + + path: Path + outcome: str + detail: str = "" + + @property + def succeeded(self) -> bool: + return self.outcome in ("uploaded", "skipped") + + +@dataclass +class GraphUploadResult: + """The outcome of one ``graph upload`` run. + + Attributes: + job_id: The BloodHound job the files were attached to, or None for a dry run or when + the job could not be started. + files: One result per file, in upload order. + job: The job's final state when the caller waited for it; otherwise None. + detail: Explanation when the job itself could not be started or ended. + job_ended: True once BloodHound accepted the request to end the job. BloodHound only + starts processing a job after it ends, so there is nothing to wait for otherwise. + """ + + job_id: int | None + files: list[FileResult] = field(default_factory=list) + job: FileUploadJob | None = None + detail: str = "" + job_ended: bool = False + + @property + def all_files_succeeded(self) -> bool: + return all(result.succeeded for result in self.files) + + +def find_collection_files(target: str) -> list[Path]: + """Work out which files to upload from what the person typed. + + Args: + target (str): A collection file (``.json`` or ``.zip``) or a folder holding them. + + Returns: + list[Path]: The files to upload, in name order. + """ + candidate = Path(target) + + if candidate.is_file(): + return [candidate] + + if candidate.is_dir(): + files = sorted( + path + for path in candidate.iterdir() + if path.is_file() and path.suffix.lower() in (".json", ".zip") + ) + if not files: + raise CollectionNotFound( + f"The folder '{candidate}' contains no .json or .zip collection files." + ) + return files + + raise CollectionNotFound( + f"'{target}' is not a file and not a folder. Point the command at an OpenHound " + "convert output folder, a single OpenGraph JSON file, or a zip of them." + ) + + +def _load_collection_file(path: Path) -> tuple[bytes, str, dict | None]: + """Read a collection file; for JSON also return the parsed document, otherwise None.""" + try: + contents = path.read_bytes() + except OSError as error: + raise InvalidCollectionFile(f"Could not read '{path}': {error}") from error + + if path.suffix.lower() == ".zip": + return contents, CONTENT_TYPE_ZIP, None + + try: + document = json.loads(contents.decode("utf-8")) + except (ValueError, UnicodeDecodeError) as error: + raise InvalidCollectionFile(f"'{path}' is not valid JSON: {error}") from error + + if not isinstance(document, dict) or "graph" not in document: + raise InvalidCollectionFile( + f"'{path}' is not an OpenGraph document: it has no top-level 'graph' section." + ) + + return contents, CONTENT_TYPE_JSON, document + + +def read_collection_file(path: Path) -> tuple[bytes, str]: + """Read a collection file and decide which content type to send it as. + + JSON files are parsed once locally and checked for the ``graph`` section every OpenGraph + document has, so a wrong file produces a clear local error naming it instead of a job + that fails later inside BloodHound. + + Returns: + tuple[bytes, str]: The raw bytes and the content type for the upload. + """ + contents, content_type, _ = _load_collection_file(path) + return contents, content_type + + +def split_document(document: dict, batch_size: int) -> Iterator[dict]: + """Cut one OpenGraph document into documents of at most ``batch_size`` nodes plus edges. + + Every part keeps the original ``metadata`` section. All node parts come before any edge + part, so the nodes an edge points at have been sent before the edge itself. + + Args: + document (dict): A parsed OpenGraph document. + batch_size (int): The most nodes plus edges allowed in one part. + + Yields: + dict: Each part, a complete OpenGraph document. + """ + graph = document["graph"] + metadata = document.get("metadata") + nodes = graph.get("nodes") or [] + edges = graph.get("edges") or [] + + def part(kind: str, items: list) -> dict: + piece: dict = {"graph": {kind: items}} + if metadata is not None: + piece = {"metadata": metadata, **piece} + return piece + + for start in range(0, len(nodes), batch_size): + yield part("nodes", nodes[start : start + batch_size]) + for start in range(0, len(edges), batch_size): + yield part("edges", edges[start : start + batch_size]) + + +def _send_part( + client: BloodHoundClient, + job_id: int, + body: bytes, + content_type: str, + file_name: str, + compress: bool, +) -> None: + """Attach one request body to the job, gzip-compressing JSON when asked.""" + content_encoding = None + if compress and content_type == CONTENT_TYPE_JSON: + original_size = len(body) + body = gzip.compress(body) + content_encoding = "gzip" + logger.info( + f"Compressed {file_name} from {original_size:,} to {len(body):,} bytes" + ) + client.upload_file_to_job( + job_id=job_id, + body=body, + content_type=content_type, + file_name=file_name, + content_encoding=content_encoding, + ) + + +def _upload_one_file( + client: BloodHoundClient, + job_id: int, + path: Path, + compress: bool, + batch_size: int, +) -> FileResult: + """Send one file, cutting a large JSON document into parts of at most ``batch_size``.""" + try: + body, content_type, document = _load_collection_file(path) + except InvalidCollectionFile as error: + return FileResult(path=path, outcome="failed", detail=str(error)) + + graph = document["graph"] if document is not None else {} + entity_count = len(graph.get("nodes") or []) + len(graph.get("edges") or []) + + if document is None or batch_size <= 0 or entity_count <= batch_size: + parts = [(path.name, body)] + else: + parts = [ + (f"{path.stem}.part-{number:04d}{path.suffix}", json.dumps(piece).encode()) + for number, piece in enumerate(split_document(document, batch_size), 1) + ] + logger.info( + f"Split {path.name} ({entity_count:,} items) into {len(parts)} parts" + ) + # The parsed document and the original bytes are no longer needed once parts exist. + del document, body + + failures = [] + for file_name, part_body in parts: + try: + _send_part(client, job_id, part_body, content_type, file_name, compress) + except BloodHoundHTTPError as error: + failures.append(f"{file_name}: {error.reason}") + except requests.RequestException as error: + failures.append(f"{file_name}: {error}") + + if failures: + return FileResult(path=path, outcome="failed", detail="; ".join(failures)) + detail = f"sent in {len(parts)} parts" if len(parts) > 1 else "" + return FileResult(path=path, outcome="uploaded", detail=detail) + + +def upload_collection_files( + client: BloodHoundClient, + files: list[Path], + dry_run: bool = False, + compress: bool = True, + batch_size: int = DEFAULT_BATCH_SIZE, +) -> GraphUploadResult: + """Attach every file to one new job and end the job. + + Files that fail the local check are reported and skipped. Files that BloodHound rejects + are reported, and the remaining files are still attached. The job is always ended once + it was started, so BloodHound processes whatever did arrive. + + Args: + client (BloodHoundClient): A signed-in BloodHound client. + files (list[Path]): The files to upload. + dry_run (bool): When True, read and check every file but contact BloodHound for nothing. + compress (bool): When True, gzip each JSON request before sending it; collection files + repeat the same keys millions of times, so they shrink a great deal. Zip files + are already compressed and are sent as they are. + batch_size (int): The most nodes plus edges in one request. A larger JSON file is + cut into parts of this size, all attached to the same job. Zip files cannot be + cut and are always sent whole. Zero or less turns batching off. + + Returns: + GraphUploadResult: The job id and one result per file. + """ + result = GraphUploadResult(job_id=None) + + # Check every file before touching BloodHound, so a typo never leaves a half-filled job. + # Each file is read and then released, so only one file is in memory at a time. + checked: list[Path] = [] + for path in files: + try: + read_collection_file(path) + except InvalidCollectionFile as error: + result.files.append( + FileResult(path=path, outcome="failed", detail=str(error)) + ) + continue + checked.append(path) + + if dry_run: + for path in checked: + result.files.append( + FileResult(path=path, outcome="skipped", detail="dry run") + ) + return result + + if not checked: + result.detail = "No file passed the local check, so no job was started." + return result + + try: + job = client.start_file_upload_job() + except (BloodHoundHTTPError, requests.RequestException) as error: + reason = error.reason if isinstance(error, BloodHoundHTTPError) else str(error) + result.detail = f"BloodHound refused to start a file-upload job: {reason}" + for path in checked: + result.files.append( + FileResult(path=path, outcome="failed", detail="job not started") + ) + return result + + result.job_id = job.id + logger.info(f"Started BloodHound file-upload job {job.id}") + + for path in checked: + result.files.append( + _upload_one_file(client, job.id, path, compress, batch_size) + ) + + try: + client.end_file_upload_job(job.id) + result.job_ended = True + except (BloodHoundHTTPError, requests.RequestException) as error: + reason = error.reason if isinstance(error, BloodHoundHTTPError) else str(error) + result.detail = f"Files were attached but the job could not be ended: {reason}" + + return result + + +def wait_for_job( + client: BloodHoundClient, + job_id: int, + timeout_seconds: float, + poll_seconds: float = DEFAULT_POLL_SECONDS, +) -> FileUploadJob | None: + """Ask BloodHound about the job until it finishes or the time runs out. + + Args: + client (BloodHoundClient): A signed-in BloodHound client. + job_id (int): The job to watch. + timeout_seconds (float): Give up after this long. + poll_seconds (float): How long to pause between questions. + + Returns: + FileUploadJob | None: The job's last known state, or None if BloodHound never listed it. + """ + deadline = monotonic() + timeout_seconds + job = client.file_upload_job(job_id) + + while job is not None: + status = job.status_code + if status is not None and status.is_finished: + return job + if monotonic() >= deadline: + logger.warning( + f"Job {job_id} is still '{job.status_text}' after {timeout_seconds:.0f} seconds" + ) + return job + sleep(poll_seconds) + job = client.file_upload_job(job_id) + + return None diff --git a/src/openhound/core/schema_upload.py b/src/openhound/core/schema_upload.py new file mode 100644 index 00000000..b0f5180a --- /dev/null +++ b/src/openhound/core/schema_upload.py @@ -0,0 +1,225 @@ +"""Find OpenGraph extension schema files and send them to BloodHound. + +An extension schema tells BloodHound which node kinds and relationship kinds a collector +produces, with their display names, icons, and colors. BloodHound must know the schema before +it can show collected data properly. This module does the finding and the sending; the +``openhound schema upload`` command in :mod:`openhound.cli.schema` is the thin layer on top. + +The thing to upload can be named three ways: + +* A single file: ``schema.json`` or a ``bundle.zip`` that contains a ``schema.json``. +* A folder: every ``.json`` and ``.zip`` file directly inside it is uploaded. +* The name of an installed collector, such as ``sccm``: the ``schema.json`` that the + collector ships next to its ``extension.yaml`` is uploaded. +""" + +import logging +from dataclasses import dataclass +from importlib import resources +from importlib.metadata import entry_points +from pathlib import Path + +import requests +from dlt.common import json + +from openhound.core.clients.bloodhound import BloodHoundClient, BloodHoundHTTPError + +logger = logging.getLogger(__name__) + +# The entry point group that every OpenHound collector registers itself under. The same +# value is used by the collector manager when it discovers installed collectors. +COLLECTOR_ENTRY_POINT_GROUP = "openhound.sources" + +# The file name a collector uses for the schema it ships with its code. +PACKAGED_SCHEMA_FILE_NAME = "schema.json" + +# The content types BloodHound accepts for an extension upload. +CONTENT_TYPE_JSON = "application/json" +CONTENT_TYPE_ZIP = "application/zip" + + +class SchemaNotFound(Exception): + """Raised when the path or collector name does not lead to any schema file.""" + + +class InvalidSchemaFile(Exception): + """Raised when a schema file cannot be read or is not valid JSON.""" + + +@dataclass +class SchemaUploadResult: + """What happened to one schema file. + + Attributes: + path: The file that was uploaded. + outcome: "created", "updated", "skipped" (dry run), or "failed". + detail: Extra explanation, used for failures. + """ + + path: Path + outcome: str + detail: str = "" + + @property + def succeeded(self) -> bool: + return self.outcome in ("created", "updated", "skipped") + + +def find_schema_files(target: str) -> list[Path]: + """Work out which files to upload from what the person typed. + + Args: + target (str): A file path, a folder path, or the name of an installed collector. + + Returns: + list[Path]: The schema files to upload, in a stable order. + """ + candidate = Path(target) + + if candidate.is_file(): + return [candidate] + + if candidate.is_dir(): + files = sorted( + path + for path in candidate.iterdir() + if path.is_file() and path.suffix.lower() in (".json", ".zip") + ) + if not files: + raise SchemaNotFound( + f"The folder '{candidate}' contains no .json or .zip schema files." + ) + return files + + packaged_schema = find_packaged_schema(target) + if packaged_schema is not None: + return [packaged_schema] + + raise SchemaNotFound( + f"'{target}' is not a file, not a folder, and not the name of an installed " + f"collector that ships a '{PACKAGED_SCHEMA_FILE_NAME}'. Installed collectors " + f"are: {', '.join(installed_collector_names()) or 'none'}." + ) + + +def installed_collector_names() -> list[str]: + """List the names of every collector installed in this Python environment.""" + return sorted( + entry_point.name + for entry_point in entry_points(group=COLLECTOR_ENTRY_POINT_GROUP) + ) + + +def find_packaged_schema(collector_name: str) -> Path | None: + """Locate the schema file that an installed collector ships with its code. + + A collector registers itself with a line like ``sccm = "openhound_sccm.main:app"``. + The part before the first dot, ``openhound_sccm`` here, is the Python package that holds + the collector's files. The schema is expected at ``schema.json`` inside that package, + beside the ``extension.yaml`` that OpenHound already requires there. + + Args: + collector_name (str): The registered collector name, for example ``sccm``. + + Returns: + Path | None: The schema file, or None when the collector or its schema is not installed. + """ + for entry_point in entry_points(group=COLLECTOR_ENTRY_POINT_GROUP): + if entry_point.name != collector_name: + continue + + package_name = entry_point.module.split(".")[0] + schema_location = resources.files(package_name) / PACKAGED_SCHEMA_FILE_NAME + if schema_location.is_file(): + # resources.files() can describe files inside a zipped package; converting + # through str() gives us a normal filesystem path for the ordinary case. + return Path(str(schema_location)) + + logger.warning( + f"Collector '{collector_name}' is installed but ships no " + f"'{PACKAGED_SCHEMA_FILE_NAME}' in its '{package_name}' package." + ) + return None + + return None + + +def read_schema_file(path: Path) -> tuple[bytes, str]: + """Read a schema file and decide which content type to send it as. + + JSON files are parsed once before sending so that a typo in the file produces a clear + local error instead of a confusing answer from the server. + + Returns: + tuple[bytes, str]: The raw bytes and the content type for the upload. + """ + try: + contents = path.read_bytes() + except OSError as error: + raise InvalidSchemaFile(f"Could not read '{path}': {error}") from error + + if path.suffix.lower() == ".zip": + return contents, CONTENT_TYPE_ZIP + + try: + json.loads(contents.decode("utf-8")) + except (ValueError, UnicodeDecodeError) as error: + raise InvalidSchemaFile(f"'{path}' is not valid JSON: {error}") from error + + return contents, CONTENT_TYPE_JSON + + +def upload_schema_files( + client: BloodHoundClient, + files: list[Path], + dry_run: bool = False, +) -> list[SchemaUploadResult]: + """Upload each schema file in turn and report what happened to each one. + + One bad file does not stop the others: every file gets its own result, and the caller + decides what to do when some failed. + + Args: + client (BloodHoundClient): A signed-in BloodHound client. + files (list[Path]): The files to upload. + dry_run (bool): When True, read and check every file but send nothing. + + Returns: + list[SchemaUploadResult]: One result per file, in the same order as ``files``. + """ + results: list[SchemaUploadResult] = [] + + for path in files: + try: + body, content_type = read_schema_file(path) + except InvalidSchemaFile as error: + results.append( + SchemaUploadResult(path=path, outcome="failed", detail=str(error)) + ) + continue + + if dry_run: + results.append( + SchemaUploadResult(path=path, outcome="skipped", detail="dry run") + ) + continue + + try: + updated = client.upsert_extension(body=body, content_type=content_type) + except BloodHoundHTTPError as error: + results.append( + SchemaUploadResult(path=path, outcome="failed", detail=error.reason) + ) + continue + except requests.RequestException as error: + # Connection refused, DNS failure, or a timeout: record it and carry on. + results.append( + SchemaUploadResult(path=path, outcome="failed", detail=str(error)) + ) + continue + + results.append( + SchemaUploadResult(path=path, outcome="updated" if updated else "created") + ) + + return results diff --git a/src/openhound/main.py b/src/openhound/main.py index fd21556d..3f79ea8b 100644 --- a/src/openhound/main.py +++ b/src/openhound/main.py @@ -4,10 +4,12 @@ from openhound.cli.collect import collect from openhound.cli.convert import convert from openhound.cli.create import create_app +from openhound.cli.graph import graph from openhound.cli.override import TyperOverride from openhound.cli.preproc import preprocess from openhound.cli.privilege_zone import privilege_zone from openhound.cli.saved_search import saved_searches +from openhound.cli.schema import schema BASE_SOUCE_PATH = Path(__file__).parent / "sources" @@ -19,6 +21,8 @@ app.add_typer(create_app, name="create") app.add_typer(saved_searches, name="searches") app.add_typer(privilege_zone, name="rules") +app.add_typer(schema, name="schema") +app.add_typer(graph, name="graph") if __name__ == "__main__": app() diff --git a/tests/test_community_ingest.py b/tests/test_community_ingest.py new file mode 100644 index 00000000..3cd491ea --- /dev/null +++ b/tests/test_community_ingest.py @@ -0,0 +1,425 @@ +"""Tests for ``openhound convert --method ingest`` against BloodHound Community Edition. + +A fake BloodHound, built with FastAPI, answers the calls the converter makes. It can pretend to be +either edition: Community Edition receives data through a file-upload job, and Enterprise through +``/api/v2/ingest``. Every ``requests.request`` call is sent to it instead of the network, so the +tests check the real paths, headers and bodies. The full tests run the installed ``faker`` +collector, the same one the Enterprise scheduler test uses. +""" + +import gzip +import json +import os +from pathlib import Path +from types import SimpleNamespace +from urllib.parse import urlsplit + +os.environ.setdefault("RUNTIME__LOG_PATH", "/tmp/openhound-test-logs") + +import pytest +import requests +from fastapi import FastAPI, Request, Response +from fastapi.testclient import TestClient +from typer.testing import CliRunner + +from openhound.core.clients.bloodhound import BloodHoundJWT +from openhound.core.community_ingest import ( + build_document, + community_edition_client_for_ingest, +) +from openhound.core.convert import Converter, Method +from openhound.core.manager import CollectorManager +from openhound.core.progress import Progress +from openhound.main import app as openhound_cli + +# The environment variables that could leak settings from the developer's machine into a test. +ALL_SETTING_VARIABLES = [ + f"DESTINATION__{section}__{setting}" + for section in ("BLOODHOUND", "BLOODHOUNDENTERPRISE") + for setting in ("URL", "TOKEN", "TOKEN_ID", "TOKEN_KEY") +] + + +@pytest.fixture +def fake_bloodhound(): + """A fake BloodHound that records every job, attached file, and Enterprise ingest call. + + Set ``app.state.edition`` to "community" or "enterprise" to choose what /api/version reports. + """ + fake = FastAPI() + fake.state.edition = "community" + fake.state.jobs = {} + fake.state.enterprise_ingests = [] + + @fake.get("/api/version") + async def version(): + return { + "data": { + "API": {"current_version": "v2"}, + "server_version": "v9.8.0", + "product_edition": fake.state.edition, + } + } + + @fake.post("/api/v2/file-upload/start", status_code=201) + async def start_job(): + job_id = 100 + len(fake.state.jobs) + fake.state.jobs[job_id] = {"files": [], "ended": False} + return {"data": {"id": job_id, "status": 1, "status_message": ""}} + + @fake.post("/api/v2/file-upload/{job_id}", status_code=202) + async def attach_file(job_id: int, request: Request): + fake.state.jobs[job_id]["files"].append( + { + "name": request.headers.get("x-file-upload-name", ""), + "content_type": request.headers.get("content-type", ""), + "content_encoding": request.headers.get("content-encoding", ""), + "body": await request.body(), + } + ) + return Response(status_code=202) + + @fake.post("/api/v2/file-upload/{job_id}/end") + async def end_job(job_id: int): + fake.state.jobs[job_id]["ended"] = True + return Response(status_code=200) + + @fake.post("/api/v2/ingest") + async def enterprise_ingest(request: Request): + fake.state.enterprise_ingests.append(gzip.decompress(await request.body())) + return Response(status_code=202) + + return TestClient(fake) + + +@pytest.fixture +def send_requests_to_fake(monkeypatch, fake_bloodhound): + """Send every ``requests.request`` call to the fake BloodHound instead of the network.""" + + def fake_request(method, url, **kwargs): + parsed = urlsplit(url) + path = f"{parsed.path}?{parsed.query}" if parsed.query else parsed.path + # The OpenHound client passes the body as ``data``; the test client calls it ``content``. + content = kwargs.pop("data", None) + kwargs.pop("timeout", None) + return fake_bloodhound.request(method.upper(), path, content=content, **kwargs) + + monkeypatch.setattr("requests.request", fake_request) + + +@pytest.fixture +def no_settings(monkeypatch, tmp_path): + """Start from a clean slate, with dlt's working files kept inside the test folder.""" + for name in ALL_SETTING_VARIABLES: + monkeypatch.delenv(name, raising=False) + monkeypatch.setenv("DLT_DATA_DIR", str(tmp_path / ".dlt")) + + +@pytest.fixture +def community_settings(no_settings, monkeypatch): + monkeypatch.setenv("DESTINATION__BLOODHOUND__URL", "http://localhost:8080") + monkeypatch.setenv("DESTINATION__BLOODHOUND__TOKEN", "test-bearer-token") + + +@pytest.fixture +def enterprise_settings(no_settings, monkeypatch): + monkeypatch.setenv( + "DESTINATION__BLOODHOUNDENTERPRISE__URL", "http://localhost:8000" + ) + monkeypatch.setenv("DESTINATION__BLOODHOUNDENTERPRISE__TOKEN_ID", "test-id") + monkeypatch.setenv("DESTINATION__BLOODHOUNDENTERPRISE__TOKEN_KEY", "test-key") + + +@pytest.fixture +def collected_faker_data(tmp_path): + """Run the installed faker collector and its preprocessing step into the test folder. + + Returns: + tuple[Path, Path, str | None]: The collected data folder, the lookup file, and the + collector's source kind. + """ + extension = next( + collector + for collector in CollectorManager.from_entrypoint().collectors + if collector.name == "faker" # pyright: ignore[reportAttributeAccessIssue] + ) + collected = tmp_path / "collected" + lookup_file = tmp_path / "lookup.duckdb" + extension.collector(output_path=collected, resources=[], progress=Progress.log) + if extension.preprocessor is not None: + extension.preprocessor( + input_path=collected / extension.name, + output_file=lookup_file, + progress=Progress.log, + ) + return collected / extension.name, lookup_file, extension.source_kind + + +def run_convert_command(collected: Path, lookup_file: Path, tmp_path: Path, *extra): + return CliRunner().invoke( + openhound_cli, + [ + "convert", + "faker", + str(collected), + str(tmp_path / "converted"), + "--lookup-file", + str(lookup_file), + "--progress", + "log", + *extra, + ], + ) + + +def count_entities(documents: list[dict]) -> tuple[int, int]: + nodes = sum(len(document["graph"]["nodes"]) for document in documents) + edges = sum(len(document["graph"]["edges"]) for document in documents) + return nodes, edges + + +# --- Choosing the route ------------------------------------------------------------------ + + +def test_community_edition_is_detected(send_requests_to_fake, community_settings): + assert community_edition_client_for_ingest() is not None + + +def test_enterprise_keeps_the_enterprise_route( + send_requests_to_fake, fake_bloodhound, community_settings +): + fake_bloodhound.app.state.edition = "enterprise" + + assert community_edition_client_for_ingest() is None + + +def test_no_settings_keeps_the_enterprise_route(no_settings): + assert community_edition_client_for_ingest() is None + + +def test_an_unanswered_edition_question_keeps_the_enterprise_route( + monkeypatch, community_settings +): + def refuse(method, url, **kwargs): + raise requests.ConnectionError("connection refused") + + monkeypatch.setattr("requests.request", refuse) + + assert community_edition_client_for_ingest() is None + + +# --- Building each file ------------------------------------------------------------------ + + +def test_documents_carry_the_source_kind_in_their_metadata(): + items = [ + {"graph": {"entity_type": "node", "content": {"id": "a"}}}, + {"graph": {"entity_type": "edge", "content": [{"kind": "K"}, {"kind": "K"}]}}, + ] + + document = build_document(items, source_kind="GHBase") + + assert document["metadata"] == {"source_kind": "GHBase"} + assert document["graph"] == {"nodes": [{"id": "a"}], "edges": [{"kind": "K"}] * 2} + + +def test_documents_without_a_source_kind_have_no_metadata(): + assert "metadata" not in build_document([], source_kind=None) + + +# --- Running a conversion ---------------------------------------------------------------- + + +def test_convert_ingest_sends_everything_through_one_community_edition_job( + send_requests_to_fake, + fake_bloodhound, + community_settings, + collected_faker_data, + tmp_path, +): + collected, lookup_file, source_kind = collected_faker_data + + result = run_convert_command(collected, lookup_file, tmp_path, "--method", "ingest") + + assert result.exit_code == 0, result.output + jobs = fake_bloodhound.app.state.jobs + assert list(jobs) == [100], "exactly one job per run" + job = jobs[100] + assert job["ended"] is True + assert job["files"], "at least one file attached" + assert all(file["content_type"] == "application/json" for file in job["files"]) + assert all(file["content_encoding"] == "gzip" for file in job["files"]) + assert all(file["name"].endswith(".json") for file in job["files"]) + + documents = [json.loads(gzip.decompress(file["body"])) for file in job["files"]] + assert source_kind, "the faker collector declares a source kind" + assert all(doc["metadata"] == {"source_kind": source_kind} for doc in documents) + # The faker collector produces the same counts the Enterprise scheduler test checks. + assert count_entities(documents) == (1000, 10000) + assert fake_bloodhound.app.state.enterprise_ingests == [] + assert not (tmp_path / "converted").exists(), "ingest writes no output files" + + +def test_convert_ingest_to_enterprise_is_unchanged( + send_requests_to_fake, + fake_bloodhound, + enterprise_settings, + collected_faker_data, + tmp_path, +): + fake_bloodhound.app.state.edition = "enterprise" + collected, lookup_file, _ = collected_faker_data + + result = run_convert_command(collected, lookup_file, tmp_path, "--method", "ingest") + + assert result.exit_code == 0, result.output + assert fake_bloodhound.app.state.jobs == {}, "Enterprise opens no file-upload job" + documents = [ + json.loads(body) for body in fake_bloodhound.app.state.enterprise_ingests + ] + assert all("metadata" not in document for document in documents) + assert count_entities(documents) == (1000, 10000) + + +def test_convert_without_method_still_writes_files( + send_requests_to_fake, + fake_bloodhound, + community_settings, + collected_faker_data, + tmp_path, +): + collected, lookup_file, _ = collected_faker_data + + result = run_convert_command(collected, lookup_file, tmp_path) + + assert result.exit_code == 0, result.output + assert list((tmp_path / "converted").glob("*.json")) + assert fake_bloodhound.app.state.jobs == {} + assert fake_bloodhound.app.state.enterprise_ingests == [] + + +def test_the_job_is_ended_even_when_the_conversion_fails( + send_requests_to_fake, fake_bloodhound, monkeypatch, tmp_path +): + def fail(self, source, **kwargs): + raise RuntimeError("conversion broke") + + monkeypatch.setattr( + "openhound.core.convert.logger_override.set_handler", lambda n: None + ) + monkeypatch.setattr(Converter, "_run", fail) + converter = Converter( + name="test", + input_path=tmp_path, + lookup=None, + output_path=tmp_path / "unused", + source_kind="TestBase", + progress=Progress.log, + method=Method.ingest, + community_edition_client=BloodHoundJWT( + token="t", base_uri="http://localhost:8080" + ), + ) + + with pytest.raises(RuntimeError, match="conversion broke"): + converter.run( + SimpleNamespace(resources={}), graph_resources=[], extra_context={} + ) + + assert fake_bloodhound.app.state.jobs[100]["ended"] is True + + +def test_a_failure_to_end_the_job_does_not_hide_the_original_error( + monkeypatch, tmp_path +): + client = BloodHoundJWT(token="t", base_uri="http://localhost:8080") + monkeypatch.setattr(client, "start_file_upload_job", lambda: SimpleNamespace(id=7)) + + def end_fails(job_id): + raise requests.ConnectionError("BloodHound went away") + + def conversion_fails(self, source, **kwargs): + raise RuntimeError("conversion broke") + + monkeypatch.setattr(client, "end_file_upload_job", end_fails) + monkeypatch.setattr( + "openhound.core.convert.logger_override.set_handler", lambda n: None + ) + monkeypatch.setattr(Converter, "_run", conversion_fails) + converter = Converter( + name="test", + input_path=tmp_path, + lookup=None, + output_path=tmp_path / "unused", + source_kind=None, + progress=Progress.log, + method=Method.ingest, + community_edition_client=client, + ) + + with pytest.raises(RuntimeError, match="conversion broke"): + converter.run( + SimpleNamespace(resources={}), graph_resources=[], extra_context={} + ) + + +@pytest.fixture +def record_edition_checks_and_skip_conversion(monkeypatch): + """Record every edition check and every converter, and skip the conversion itself. + + The edition check is replaced where ``app.py`` looks it up, because ``app.py`` imports it + by name; replacing it in ``community_ingest`` would not be seen there. + + Returns: + dict: ``edition_checks`` counts the checks; ``converters`` holds each converter that ran. + """ + record = {"edition_checks": 0, "converters": []} + + def check_edition(): + # Returning nothing means "take the Enterprise route", like a real Enterprise answer. + record["edition_checks"] += 1 + + def skip_conversion(self, source_object, graph_resources, extra_context, **kwargs): + record["converters"].append(self) + return "skipped" + + monkeypatch.setattr( + "openhound.core.app.community_edition_client_for_ingest", check_edition + ) + monkeypatch.setattr(Converter, "run", skip_conversion) + return record + + +def test_the_scheduler_path_never_asks_for_the_edition( + record_edition_checks_and_skip_conversion, collected_faker_data +): + """The Enterprise scheduler calls the collector's converter directly, like this.""" + collected, lookup_file, _ = collected_faker_data + extension = next( + collector + for collector in CollectorManager.from_entrypoint().collectors + if collector.name == "faker" # pyright: ignore[reportAttributeAccessIssue] + ) + + extension.converter( + input_path=collected, + lookup_file=lookup_file, + method=Method.ingest, + progress=Progress.log, + ) + + record = record_edition_checks_and_skip_conversion + assert record["edition_checks"] == 0 + assert record["converters"][0].sends_to_community_edition is False + + +def test_the_command_line_asks_for_the_edition( + record_edition_checks_and_skip_conversion, collected_faker_data, tmp_path +): + """The control for the test above: the command line does ask, so the check is watched.""" + collected, lookup_file, _ = collected_faker_data + + result = run_convert_command(collected, lookup_file, tmp_path, "--method", "ingest") + + assert result.exit_code == 0, result.output + assert record_edition_checks_and_skip_conversion["edition_checks"] == 1 diff --git a/tests/test_graph_upload.py b/tests/test_graph_upload.py new file mode 100644 index 00000000..ad7abd6a --- /dev/null +++ b/tests/test_graph_upload.py @@ -0,0 +1,585 @@ +"""Tests for ``openhound graph upload`` and the pieces underneath it. + +A small fake BloodHound API, built with FastAPI, stands in for the real server. It implements +the four file-upload job calls the client makes: start a job, attach a file, end the job, and +list jobs. Every request is redirected to it, so the tests check the real request shape +without needing a running BloodHound. +""" + +import gzip +import json +import zipfile +from pathlib import Path +from types import SimpleNamespace +from urllib.parse import parse_qs, urlsplit + +import pytest +import requests +from fastapi import FastAPI, Request, Response +from fastapi.testclient import TestClient +from typer.testing import CliRunner + +from openhound.cli.graph import graph as graph_cli +from openhound.core.clients.bloodhound import ( + FILE_ATTACH_TIMEOUT_SECONDS, + QUICK_REQUEST_TIMEOUT_SECONDS, + BloodHoundJWT, +) +from openhound.core.clients.models import FileUploadJobStatus +from openhound.core.graph_upload import ( + CONTENT_TYPE_JSON, + CONTENT_TYPE_ZIP, + CollectionNotFound, + find_collection_files, + upload_collection_files, + wait_for_job, +) + +EXAMPLE_DOCUMENT = { + "metadata": {"source_kind": "GH"}, + "graph": { + "nodes": [ + {"id": "GH:alice", "kinds": ["GH_User"], "properties": {"name": "alice"}} + ], + "edges": [], + }, +} + +# The environment variables that could leak settings from the developer's machine into a test. +ALL_SETTING_VARIABLES = [ + f"DESTINATION__{section}__{setting}" + for section in ("BLOODHOUND", "BLOODHOUNDENTERPRISE") + for setting in ("URL", "TOKEN", "TOKEN_ID", "TOKEN_KEY") +] + + +@pytest.fixture +def mock_bloodhound_api(): + """A fake BloodHound that runs file-upload jobs in memory. + + Each job records the files attached to it. Ending a job moves it to "ingesting"; every + later status question advances it one step, through "analyzing" to "complete", which + lets the waiting logic be tested without real delays. A file whose name contains + "reject" is refused with 400, which lets per-file failures be tested. + """ + app = FastAPI() + app.state.jobs = {} + app.state.next_job_id = 100 + # Set to True to make the fake refuse every request to end a job. + app.state.refuse_to_end_jobs = False + + def job_record(job_id: int) -> dict: + job = app.state.jobs[job_id] + return { + "id": job_id, + "status": job["status"], + "status_message": job["status_message"], + "total_files": len(job["files"]), + "failed_files": 0, + } + + @app.post("/api/v2/file-upload/start", status_code=201) + async def start_job(): + job_id = app.state.next_job_id + app.state.next_job_id += 1 + app.state.jobs[job_id] = { + "status": int(FileUploadJobStatus.RUNNING), + "status_message": "", + "files": [], + "ended": False, + } + return {"data": job_record(job_id)} + + @app.post("/api/v2/file-upload/{job_id}", status_code=202) + async def attach_file(job_id: int, request: Request): + job = app.state.jobs.get(job_id) + if job is None: + return Response(status_code=404, content="no such job") + if job["ended"]: + return Response(status_code=400, content="job must be in running status") + name = request.headers.get("x-file-upload-name", "") + if "reject" in name: + return Response(status_code=400, content=f"{name} failed validation") + job["files"].append( + { + "name": name, + "content_type": request.headers.get("content-type", ""), + "content_encoding": request.headers.get("content-encoding", ""), + "body": await request.body(), + } + ) + return Response(status_code=202) + + @app.post("/api/v2/file-upload/{job_id}/end") + async def end_job(job_id: int): + job = app.state.jobs.get(job_id) + if job is None: + return Response(status_code=404, content="no such job") + if app.state.refuse_to_end_jobs: + return Response(status_code=500, content="could not end the job") + job["ended"] = True + job["status"] = int(FileUploadJobStatus.INGESTING) + return Response(status_code=200) + + @app.get("/api/v2/file-upload") + async def list_jobs(request: Request): + wanted = parse_qs(request.url.query).get("id", [""])[0] + job_id = int(wanted.removeprefix("eq:")) if wanted else None + records = [] + for candidate_id, job in app.state.jobs.items(): + if job_id is not None and candidate_id != job_id: + continue + records.append(job_record(candidate_id)) + # Each look advances an ended job one step towards completion. + if job["ended"] and job["status"] == int(FileUploadJobStatus.INGESTING): + job["status"] = int(FileUploadJobStatus.ANALYZING) + elif job["ended"] and job["status"] == int(FileUploadJobStatus.ANALYZING): + job["status"] = int(FileUploadJobStatus.COMPLETE) + return {"data": records, "count": len(records)} + + return TestClient(app) + + +@pytest.fixture +def mock_bh_requests(monkeypatch, mock_bloodhound_api): + """Send every ``requests.request`` call to the fake BloodHound instead of the network.""" + + def mock_request(method, url, **kwargs): + parsed = urlsplit(url) + path = parsed.path + if parsed.query: + path = f"{path}?{parsed.query}" + content = kwargs.pop("data", None) + kwargs.pop("timeout", None) + return mock_bloodhound_api.request( + method.upper(), path, content=content, **kwargs + ) + + monkeypatch.setattr("requests.request", mock_request) + # Waiting must not actually sleep in tests. + monkeypatch.setattr("openhound.core.graph_upload.sleep", lambda seconds: None) + + +@pytest.fixture +def no_credentials(monkeypatch): + for name in ALL_SETTING_VARIABLES: + monkeypatch.delenv(name, raising=False) + + +@pytest.fixture +def bearer_credentials(no_credentials, monkeypatch): + monkeypatch.setenv("DESTINATION__BLOODHOUND__URL", "http://localhost:8080") + monkeypatch.setenv("DESTINATION__BLOODHOUND__TOKEN", "test-bearer-token") + + +def write_document(path: Path, payload: dict | None = None) -> Path: + path.write_text(json.dumps(payload or EXAMPLE_DOCUMENT)) + return path + + +def make_client() -> BloodHoundJWT: + return BloodHoundJWT(token="t", base_uri="http://localhost:8080") + + +# --- Finding files ----------------------------------------------------------------------- + + +def test_find_collection_files_accepts_a_single_file(tmp_path): + document = write_document(tmp_path / "users-1.json") + assert find_collection_files(str(document)) == [document] + + +def test_find_collection_files_lists_json_and_zip_in_a_folder(tmp_path): + second = write_document(tmp_path / "groups-1.json") + first = write_document(tmp_path / "computers-1.json") + bundle = tmp_path / "sccm_collection.zip" + with zipfile.ZipFile(bundle, "w") as archive: + archive.writestr("nodes-1.json", json.dumps(EXAMPLE_DOCUMENT)) + (tmp_path / "lookup.duckdb").write_bytes(b"ignored") + + assert find_collection_files(str(tmp_path)) == [first, second, bundle] + + +def test_find_collection_files_explains_an_empty_folder(tmp_path): + with pytest.raises(CollectionNotFound, match="contains no .json or .zip"): + find_collection_files(str(tmp_path)) + + +def test_find_collection_files_explains_a_missing_path(): + with pytest.raises(CollectionNotFound, match="not a file and not a folder"): + find_collection_files("/no/such/place") + + +# --- Uploading --------------------------------------------------------------------------- + + +def test_upload_attaches_every_file_to_one_job_and_ends_it( + mock_bh_requests, mock_bloodhound_api, tmp_path +): + first = write_document(tmp_path / "a.json") + second = write_document(tmp_path / "b.json") + + result = upload_collection_files(make_client(), [first, second]) + + assert [file.outcome for file in result.files] == ["uploaded", "uploaded"] + assert result.job_id == 100 + job = mock_bloodhound_api.app.state.jobs[100] + assert [file["name"] for file in job["files"]] == ["a.json", "b.json"] + assert all(file["content_type"] == CONTENT_TYPE_JSON for file in job["files"]) + assert json.loads(gzip.decompress(job["files"][0]["body"])) == EXAMPLE_DOCUMENT + assert job["ended"] is True + + +def test_upload_compresses_json_files_with_gzip( + mock_bh_requests, mock_bloodhound_api, tmp_path +): + document = write_document(tmp_path / "a.json") + + result = upload_collection_files(make_client(), [document]) + + assert result.files[0].outcome == "uploaded" + attached = mock_bloodhound_api.app.state.jobs[100]["files"][0] + assert attached["content_encoding"] == "gzip" + # The content type still names the file inside, not the compression. + assert attached["content_type"] == CONTENT_TYPE_JSON + assert json.loads(gzip.decompress(attached["body"])) == EXAMPLE_DOCUMENT + + +def make_large_document(nodes: int, edges: int) -> dict: + return { + "metadata": {"source_kind": "GH"}, + "graph": { + "nodes": [ + {"id": f"GH:{n}", "kinds": ["GH_User"], "properties": {"name": f"u{n}"}} + for n in range(nodes) + ], + "edges": [ + { + "start": {"match_by": "id", "value": f"GH:{e}"}, + "end": {"match_by": "id", "value": f"GH:{e + 1}"}, + "kind": "GH_Knows", + } + for e in range(edges) + ], + }, + } + + +def test_large_files_are_cut_into_parts_nodes_before_edges( + mock_bh_requests, mock_bloodhound_api, tmp_path +): + big = tmp_path / "big.json" + big.write_text(json.dumps(make_large_document(nodes=25, edges=12))) + + result = upload_collection_files(make_client(), [big], batch_size=10) + + assert result.files[0].outcome == "uploaded" + assert result.files[0].detail == "sent in 5 parts" + attached = mock_bloodhound_api.app.state.jobs[100]["files"] + assert [f["name"] for f in attached] == [ + f"big.part-000{n}.json" for n in range(1, 6) + ] + parts = [json.loads(gzip.decompress(f["body"])) for f in attached] + assert all(part["metadata"] == {"source_kind": "GH"} for part in parts) + sizes = [ + len(p["graph"].get("nodes", [])) + len(p["graph"].get("edges", [])) + for p in parts + ] + assert sizes == [10, 10, 5, 10, 2] + # Every node part comes before any edge part. + kinds = ["nodes" if "nodes" in p["graph"] else "edges" for p in parts] + assert kinds == ["nodes", "nodes", "nodes", "edges", "edges"] + + +def test_files_within_the_batch_size_are_sent_whole( + mock_bh_requests, mock_bloodhound_api, tmp_path +): + small = write_document(tmp_path / "small.json") + + upload_collection_files(make_client(), [small], batch_size=10) + + attached = mock_bloodhound_api.app.state.jobs[100]["files"] + assert [f["name"] for f in attached] == ["small.json"] + + +def test_batch_size_zero_turns_batching_off( + mock_bh_requests, mock_bloodhound_api, tmp_path +): + big = tmp_path / "big.json" + big.write_text(json.dumps(make_large_document(nodes=25, edges=12))) + + upload_collection_files(make_client(), [big], batch_size=0) + + assert len(mock_bloodhound_api.app.state.jobs[100]["files"]) == 1 + + +def test_a_rejected_part_fails_the_file_but_later_parts_are_still_sent( + mock_bh_requests, mock_bloodhound_api, tmp_path +): + # The fake server rejects any file name containing "reject". + big = tmp_path / "reject.json" + big.write_text(json.dumps(make_large_document(nodes=25, edges=0))) + + result = upload_collection_files(make_client(), [big], batch_size=10) + + assert result.files[0].outcome == "failed" + assert "reject.part-0001.json" in result.files[0].detail + assert "reject.part-0003.json" in result.files[0].detail + + +def test_network_errors_are_reported_per_file_not_raised(monkeypatch, tmp_path): + def refuse(method, url, **kwargs): + raise requests.ConnectionError("connection refused") + + monkeypatch.setattr("requests.request", refuse) + + result = upload_collection_files( + make_client(), [write_document(tmp_path / "a.json")] + ) + + assert result.job_id is None + assert "connection refused" in result.detail + + +def test_cli_reports_losing_contact_while_waiting( + mock_bh_requests, mock_bloodhound_api, bearer_credentials, tmp_path, monkeypatch +): + def lose_contact(*args, **kwargs): + raise requests.ConnectionError("Connection reset by peer") + + monkeypatch.setattr("openhound.cli.graph.wait_for_job", lose_contact) + + result = CliRunner().invoke( + graph_cli, ["upload", str(write_document(tmp_path / "a.json"))] + ) + + assert result.exit_code == 1 + assert "Lost contact with BloodHound" in result.output + assert "Traceback" not in result.output + + +def test_upload_can_skip_compression(mock_bh_requests, mock_bloodhound_api, tmp_path): + document = write_document(tmp_path / "a.json") + + upload_collection_files(make_client(), [document], compress=False) + + attached = mock_bloodhound_api.app.state.jobs[100]["files"][0] + assert attached["content_encoding"] == "" + assert json.loads(attached["body"]) == EXAMPLE_DOCUMENT + + +def test_upload_sends_zip_files_with_the_zip_content_type( + mock_bh_requests, mock_bloodhound_api, tmp_path +): + bundle = tmp_path / "collection.zip" + with zipfile.ZipFile(bundle, "w") as archive: + archive.writestr("nodes-1.json", json.dumps(EXAMPLE_DOCUMENT)) + + result = upload_collection_files(make_client(), [bundle]) + + assert result.files[0].outcome == "uploaded" + attached = mock_bloodhound_api.app.state.jobs[100]["files"][0] + assert attached["content_type"] == CONTENT_TYPE_ZIP + assert attached["content_encoding"] == "" + assert attached["name"] == "collection.zip" + + +def test_upload_rejects_bad_files_locally_before_starting_a_job( + mock_bh_requests, mock_bloodhound_api, tmp_path +): + broken = tmp_path / "broken.json" + broken.write_text("{ not json") + wrong_shape = write_document(tmp_path / "wrong.json", {"nodes": []}) + + result = upload_collection_files(make_client(), [broken, wrong_shape]) + + assert [file.outcome for file in result.files] == ["failed", "failed"] + assert "not valid JSON" in result.files[0].detail + assert "no top-level 'graph' section" in result.files[1].detail + assert result.job_id is None + assert "no job was started" in result.detail + assert mock_bloodhound_api.app.state.jobs == {} + + +def test_upload_reports_a_server_rejection_and_still_ends_the_job( + mock_bh_requests, mock_bloodhound_api, tmp_path +): + rejected = write_document(tmp_path / "reject-me.json") + accepted = write_document(tmp_path / "fine.json") + + result = upload_collection_files(make_client(), [rejected, accepted]) + + assert result.files[0].outcome == "failed" + assert "failed validation" in result.files[0].detail + assert result.files[1].outcome == "uploaded" + job = mock_bloodhound_api.app.state.jobs[100] + assert [file["name"] for file in job["files"]] == ["fine.json"] + assert job["ended"] is True + + +def test_dry_run_checks_files_but_contacts_nothing( + mock_bh_requests, mock_bloodhound_api, tmp_path +): + good = write_document(tmp_path / "good.json") + broken = tmp_path / "broken.json" + broken.write_text("{ not json") + + result = upload_collection_files(client=None, files=[good, broken], dry_run=True) # type: ignore[arg-type] + + outcomes = {file.path.name: file.outcome for file in result.files} + assert outcomes == {"good.json": "skipped", "broken.json": "failed"} + assert result.job_id is None + assert mock_bloodhound_api.app.state.jobs == {} + + +# --- Waiting ----------------------------------------------------------------------------- + + +def test_wait_for_job_returns_once_the_job_is_finished( + mock_bh_requests, mock_bloodhound_api, tmp_path +): + document = write_document(tmp_path / "a.json") + result = upload_collection_files(make_client(), [document]) + + job = wait_for_job(make_client(), result.job_id, timeout_seconds=30) + + assert job is not None + assert job.status_code is FileUploadJobStatus.COMPLETE + assert job.status_text == "complete" + + +def test_wait_for_job_gives_up_after_the_timeout( + mock_bh_requests, mock_bloodhound_api, tmp_path, monkeypatch +): + document = write_document(tmp_path / "a.json") + result = upload_collection_files(make_client(), [document]) + # Pretend the clock jumps past the deadline immediately after the first look. + clock = iter([0.0, 1000.0, 2000.0]) + monkeypatch.setattr("openhound.core.graph_upload.monotonic", lambda: next(clock)) + + job = wait_for_job(make_client(), result.job_id, timeout_seconds=10) + + assert job is not None + assert job.status_code is FileUploadJobStatus.INGESTING + + +def test_wait_for_job_returns_none_for_an_unknown_job( + mock_bh_requests, mock_bloodhound_api +): + assert wait_for_job(make_client(), 999, timeout_seconds=1) is None + + +# --- The command line -------------------------------------------------------------------- + + +def test_cli_uploads_waits_and_reports_the_job( + mock_bh_requests, mock_bloodhound_api, bearer_credentials, tmp_path +): + document = write_document(tmp_path / "a.json") + + result = CliRunner().invoke(graph_cli, ["upload", str(document)]) + + assert result.exit_code == 0, result.output + assert "uploaded" in result.output + assert "file-upload job: 100" in result.output + assert "Job 100 is complete" in result.output + + +def test_cli_no_wait_points_at_the_ingest_page( + mock_bh_requests, mock_bloodhound_api, bearer_credentials, tmp_path +): + document = write_document(tmp_path / "a.json") + + result = CliRunner().invoke(graph_cli, ["upload", str(document), "--no-wait"]) + + assert result.exit_code == 0, result.output + assert "File Ingest" in result.output + assert "Job 100 is" not in result.output + + +def test_cli_does_not_wait_for_a_job_that_could_not_be_ended( + mock_bh_requests, mock_bloodhound_api, bearer_credentials, tmp_path, monkeypatch +): + mock_bloodhound_api.app.state.refuse_to_end_jobs = True + waited = [] + monkeypatch.setattr( + "openhound.cli.graph.wait_for_job", lambda *args, **kwargs: waited.append(1) + ) + document = write_document(tmp_path / "a.json") + + result = CliRunner().invoke(graph_cli, ["upload", str(document)]) + + assert result.exit_code == 1 + assert "could not be ended" in result.output + assert waited == [], ( + "an open job is never processed, so there is nothing to wait for" + ) + assert "Watch its progress" not in result.output + + +def test_every_job_request_has_a_time_limit(monkeypatch): + """A server that accepts the connection and then stays silent must not block forever.""" + time_limits = {} + replies = { + "/api/version": {"data": {"product_edition": "community"}}, + "/api/v2/file-upload/start": {"data": {"id": 1, "status": 1}}, + "/api/v2/file-upload": {"data": []}, + } + + def record(method, url, **kwargs): + path = urlsplit(url).path + time_limits[(method, path)] = kwargs.get("timeout") + return SimpleNamespace( + status_code=200, text="", json=lambda: replies.get(path, {}) + ) + + monkeypatch.setattr("requests.request", record) + client = make_client() + + client.product_edition() + client.start_file_upload_job() + client.upload_file_to_job( + job_id=1, body=b"{}", content_type="application/json", file_name="a.json" + ) + client.end_file_upload_job(1) + client.file_upload_job(1) + + assert time_limits == { + ("GET", "/api/version"): QUICK_REQUEST_TIMEOUT_SECONDS, + ("POST", "/api/v2/file-upload/start"): QUICK_REQUEST_TIMEOUT_SECONDS, + ("POST", "/api/v2/file-upload/1"): FILE_ATTACH_TIMEOUT_SECONDS, + ("POST", "/api/v2/file-upload/1/end"): QUICK_REQUEST_TIMEOUT_SECONDS, + ("GET", "/api/v2/file-upload"): QUICK_REQUEST_TIMEOUT_SECONDS, + } + + +def test_cli_dry_run_needs_no_credentials( + mock_bh_requests, mock_bloodhound_api, no_credentials, tmp_path +): + document = write_document(tmp_path / "a.json") + + result = CliRunner().invoke(graph_cli, ["upload", str(document), "--dry-run"]) + + assert result.exit_code == 0, result.output + assert "skipped" in result.output + assert mock_bloodhound_api.app.state.jobs == {} + + +def test_cli_exits_with_status_one_when_a_file_fails( + mock_bh_requests, mock_bloodhound_api, bearer_credentials, tmp_path +): + write_document(tmp_path / "reject-me.json") + write_document(tmp_path / "fine.json") + + result = CliRunner().invoke(graph_cli, ["upload", str(tmp_path)]) + + assert result.exit_code == 1 + assert "failed validation" in result.output + assert "uploaded" in result.output + + +def test_cli_explains_a_missing_path(bearer_credentials): + result = CliRunner().invoke(graph_cli, ["upload", "/no/such/place"]) + + assert result.exit_code == 1 + assert "not a file and not a folder" in result.output diff --git a/tests/test_schema_upload.py b/tests/test_schema_upload.py new file mode 100644 index 00000000..0c05b1e3 --- /dev/null +++ b/tests/test_schema_upload.py @@ -0,0 +1,423 @@ +"""Tests for ``openhound schema upload`` and the pieces underneath it. + +A small fake BloodHound API, built with FastAPI, stands in for the real server. Every request +the OpenHound client makes is redirected to it, so the tests check the real request shape +(path, method, headers, body) without needing a running BloodHound. +""" + +import json +import subprocess +import sys +import zipfile +from importlib.metadata import EntryPoint +from pathlib import Path +from urllib.parse import urlsplit + +import pytest +import requests +from fastapi import FastAPI, Request, Response +from fastapi.testclient import TestClient +from typer.testing import CliRunner + +from openhound.cli.schema import schema as schema_cli +from openhound.core.clients.bloodhound import BloodHound, BloodHoundJWT +from openhound.core.clients.from_secrets import ( + MissingBloodHoundCredentials, + bloodhound_client_from_secrets, +) +from openhound.core.schema_upload import ( + CONTENT_TYPE_JSON, + CONTENT_TYPE_ZIP, + SchemaNotFound, + find_packaged_schema, + find_schema_files, + upload_schema_files, +) + +EXAMPLE_SCHEMA = { + "schema": { + "name": "GH", + "display_name": "GitHub", + "version": "v0.0.1", + "namespace": "GH", + }, + "node_kinds": [ + {"name": "GH_User", "display_name": "User", "is_display_kind": True} + ], + "relationship_kinds": [{"name": "GH_ContributedTo", "is_traversable": True}], +} + + +@pytest.fixture +def mock_bloodhound_api(): + """A fake BloodHound that remembers every extension it was sent. + + It answers 201 the first time a schema name is seen and 200 afterwards, exactly like + the real "upsert" endpoint. A body without a top-level "schema" key is rejected with 400. + """ + app = FastAPI() + app.state.extensions = {} + app.state.received = [] + + @app.put("/api/v2/extensions") + async def upsert_extension(request: Request): + content_type = request.headers.get("content-type", "") + body = await request.body() + app.state.received.append({"content_type": content_type, "body": body}) + + if content_type.startswith(CONTENT_TYPE_ZIP): + # The real server unpacks the bundle; here it is enough to note that it arrived. + name = "zip-bundle" + elif content_type.startswith(CONTENT_TYPE_JSON): + payload = json.loads(body) + if "schema" not in payload: + return Response(status_code=400, content="missing schema section") + name = payload["schema"]["name"] + else: + return Response(status_code=415, content="unsupported content type") + + existed = name in app.state.extensions + app.state.extensions[name] = body + return Response(status_code=200 if existed else 201) + + return TestClient(app) + + +@pytest.fixture +def mock_bh_requests(monkeypatch, mock_bloodhound_api): + """Send every ``requests.request`` call to the fake BloodHound instead of the network.""" + + def mock_request(method, url, **kwargs): + parsed = urlsplit(url) + path = parsed.path + if parsed.query: + path = f"{path}?{parsed.query}" + # The OpenHound client passes the body as ``data``; the test client calls it ``content``. + content = kwargs.pop("data", None) + kwargs.pop("timeout", None) + return mock_bloodhound_api.request( + method.upper(), path, content=content, **kwargs + ) + + monkeypatch.setattr("requests.request", mock_request) + + +# Every environment variable that any of the three configuration sections could contribute. +ALL_SETTING_VARIABLES = [ + f"DESTINATION__{section}__{setting}" + for section in ("BLOODHOUND", "BLOODHOUNDENTERPRISE") + for setting in ("URL", "TOKEN", "TOKEN_ID", "TOKEN_KEY") +] + + +@pytest.fixture +def no_credentials(monkeypatch): + """Start from a clean slate so settings on the developer's machine cannot leak into a test.""" + for name in ALL_SETTING_VARIABLES: + monkeypatch.delenv(name, raising=False) + + +@pytest.fixture +def bearer_credentials(no_credentials, monkeypatch): + """Configure a bearer token so the client can be built from secrets.""" + monkeypatch.setenv("DESTINATION__BLOODHOUND__URL", "http://localhost:8080") + monkeypatch.setenv("DESTINATION__BLOODHOUND__TOKEN", "test-bearer-token") + + +def write_schema(path: Path, payload: dict | None = None) -> Path: + path.write_text(json.dumps(payload or EXAMPLE_SCHEMA)) + return path + + +# --- Finding files ----------------------------------------------------------------------- + + +def test_find_schema_files_accepts_a_single_file(tmp_path): + schema_file = write_schema(tmp_path / "schema.json") + assert find_schema_files(str(schema_file)) == [schema_file] + + +def test_find_schema_files_lists_json_and_zip_in_a_folder(tmp_path): + second = write_schema(tmp_path / "b.json") + first = write_schema(tmp_path / "a.json") + bundle = tmp_path / "c.zip" + with zipfile.ZipFile(bundle, "w") as archive: + archive.writestr("schema.json", json.dumps(EXAMPLE_SCHEMA)) + (tmp_path / "notes.txt").write_text("ignored") + + assert find_schema_files(str(tmp_path)) == [first, second, bundle] + + +def test_find_schema_files_explains_an_empty_folder(tmp_path): + with pytest.raises(SchemaNotFound, match="contains no .json or .zip"): + find_schema_files(str(tmp_path)) + + +def test_find_schema_files_explains_an_unknown_target(monkeypatch): + monkeypatch.setattr("openhound.core.schema_upload.entry_points", lambda group: []) + with pytest.raises(SchemaNotFound, match="not the name of an installed collector"): + find_schema_files("does-not-exist") + + +def test_find_packaged_schema_reads_the_collectors_own_schema(tmp_path, monkeypatch): + """A collector registered as ``fake = openhound_fake.main:app`` ships schema.json in its package.""" + package_dir = tmp_path / "openhound_fake" + package_dir.mkdir() + (package_dir / "__init__.py").write_text("") + expected = write_schema(package_dir / "schema.json") + monkeypatch.syspath_prepend(str(tmp_path)) + + fake_entry_point = EntryPoint( + name="fake", value="openhound_fake.main:app", group="openhound.sources" + ) + monkeypatch.setattr( + "openhound.core.schema_upload.entry_points", lambda group: [fake_entry_point] + ) + + assert find_packaged_schema("fake") == expected + assert find_schema_files("fake") == [expected] + assert find_packaged_schema("other") is None + + +# --- Uploading --------------------------------------------------------------------------- + + +def test_upload_creates_then_updates(mock_bh_requests, mock_bloodhound_api, tmp_path): + schema_file = write_schema(tmp_path / "schema.json") + client = BloodHoundJWT(token="t", base_uri="http://localhost:8080") + + first = upload_schema_files(client, [schema_file]) + second = upload_schema_files(client, [schema_file]) + + assert [result.outcome for result in first] == ["created"] + assert [result.outcome for result in second] == ["updated"] + assert list(mock_bloodhound_api.app.state.extensions) == ["GH"] + sent = mock_bloodhound_api.app.state.received[0] + assert sent["content_type"] == CONTENT_TYPE_JSON + assert json.loads(sent["body"]) == EXAMPLE_SCHEMA + + +def test_upload_sends_zip_bundles_with_the_zip_content_type( + mock_bh_requests, mock_bloodhound_api, tmp_path +): + bundle = tmp_path / "bundle.zip" + with zipfile.ZipFile(bundle, "w") as archive: + archive.writestr("schema.json", json.dumps(EXAMPLE_SCHEMA)) + client = BloodHoundJWT(token="t", base_uri="http://localhost:8080") + + results = upload_schema_files(client, [bundle]) + + assert results[0].outcome == "created" + assert mock_bloodhound_api.app.state.received[0]["content_type"] == CONTENT_TYPE_ZIP + + +def test_upload_reports_invalid_json_without_sending_it( + mock_bh_requests, mock_bloodhound_api, tmp_path +): + broken = tmp_path / "broken.json" + broken.write_text("{ not json") + good = write_schema(tmp_path / "good.json") + client = BloodHoundJWT(token="t", base_uri="http://localhost:8080") + + results = upload_schema_files(client, [broken, good]) + + assert results[0].outcome == "failed" + assert "not valid JSON" in results[0].detail + assert results[1].outcome == "created" + # Only the good file reached the server. + assert len(mock_bloodhound_api.app.state.received) == 1 + + +def test_upload_reports_server_rejections_per_file( + mock_bh_requests, mock_bloodhound_api, tmp_path +): + rejected = write_schema(tmp_path / "rejected.json", {"node_kinds": []}) + client = BloodHoundJWT(token="t", base_uri="http://localhost:8080") + + results = upload_schema_files(client, [rejected]) + + assert results[0].outcome == "failed" + assert "missing schema section" in results[0].detail + + +def test_upload_sends_a_bounded_timeout(monkeypatch, tmp_path): + seen = {} + + def fake_request(method, url, **kwargs): + seen.update(kwargs) + raise requests.ConnectTimeout("timed out") + + monkeypatch.setattr("requests.request", fake_request) + client = BloodHoundJWT(token="t", base_uri="http://localhost:8080") + + upload_schema_files(client, [write_schema(tmp_path / "schema.json")]) + + assert seen["timeout"] == (10, 120) + + +def test_upload_reports_network_errors_and_continues(monkeypatch, tmp_path): + calls = [] + + def fake_request(method, url, **kwargs): + calls.append(url) + raise requests.ConnectionError("connection refused") + + monkeypatch.setattr("requests.request", fake_request) + client = BloodHoundJWT(token="t", base_uri="http://localhost:8080") + first = write_schema(tmp_path / "a.json") + second = write_schema(tmp_path / "b.json") + + results = upload_schema_files(client, [first, second]) + + assert [r.outcome for r in results] == ["failed", "failed"] + assert "connection refused" in results[0].detail + assert len(calls) == 2 # the second file was still attempted + + +def test_dry_run_reads_files_but_sends_nothing( + mock_bh_requests, mock_bloodhound_api, tmp_path +): + schema_file = write_schema(tmp_path / "schema.json") + + results = upload_schema_files(client=None, files=[schema_file], dry_run=True) # type: ignore[arg-type] + + assert results[0].outcome == "skipped" + assert mock_bloodhound_api.app.state.received == [] + + +# --- Choosing credentials ---------------------------------------------------------------- + + +def test_client_from_secrets_prefers_the_api_token_pair(no_credentials, monkeypatch): + monkeypatch.setenv("DESTINATION__BLOODHOUND__URL", "http://localhost:8080") + monkeypatch.setenv("DESTINATION__BLOODHOUND__TOKEN_ID", "id") + monkeypatch.setenv("DESTINATION__BLOODHOUND__TOKEN_KEY", "key") + monkeypatch.setenv("DESTINATION__BLOODHOUND__TOKEN", "bearer") + + client = bloodhound_client_from_secrets() + + assert isinstance(client, BloodHound) + assert client.base_uri == "http://localhost:8080" + + +def test_client_from_secrets_falls_back_to_the_bearer_token(bearer_credentials): + client = bloodhound_client_from_secrets() + assert isinstance(client, BloodHoundJWT) + + +def test_client_from_secrets_reads_the_existing_enterprise_section( + no_credentials, monkeypatch +): + """An Enterprise user whose ingest already works needs no new configuration.""" + monkeypatch.setenv( + "DESTINATION__BLOODHOUNDENTERPRISE__URL", "https://tenant.example" + ) + monkeypatch.setenv("DESTINATION__BLOODHOUNDENTERPRISE__TOKEN_ID", "id") + monkeypatch.setenv("DESTINATION__BLOODHOUNDENTERPRISE__TOKEN_KEY", "key") + + client = bloodhound_client_from_secrets() + + assert isinstance(client, BloodHound) + assert client.base_uri == "https://tenant.example" + + +def test_client_from_secrets_prefers_the_bloodhound_section( + no_credentials, monkeypatch +): + monkeypatch.setenv("DESTINATION__BLOODHOUND__URL", "http://community") + monkeypatch.setenv("DESTINATION__BLOODHOUND__TOKEN", "bearer") + monkeypatch.setenv("DESTINATION__BLOODHOUNDENTERPRISE__URL", "https://enterprise") + monkeypatch.setenv("DESTINATION__BLOODHOUNDENTERPRISE__TOKEN_ID", "id") + monkeypatch.setenv("DESTINATION__BLOODHOUNDENTERPRISE__TOKEN_KEY", "key") + + client = bloodhound_client_from_secrets() + + assert client.base_uri == "http://community" + + +def test_client_from_secrets_reads_the_url_from_config_toml(no_credentials, tmp_path): + """The address may live in config.toml, as OpenHound's example configurations show. + + dlt reads its settings files once per process, so this check runs a fresh interpreter in a + folder that holds only a config.toml with the address and a secrets.toml with the token. + """ + settings_dir = tmp_path / ".dlt" + settings_dir.mkdir() + (settings_dir / "config.toml").write_text( + '[destination.bloodhound]\nurl = "http://from-config-toml"\n' + ) + (settings_dir / "secrets.toml").write_text( + '[destination.bloodhound]\ntoken = "from-secrets-toml"\n' + ) + program = ( + "from openhound.core.clients.from_secrets import bloodhound_client_from_secrets\n" + "client = bloodhound_client_from_secrets()\n" + "print(type(client).__name__, client.base_uri)\n" + ) + + completed = subprocess.run( + [sys.executable, "-c", program], + cwd=tmp_path, + capture_output=True, + text=True, + check=True, + ) + + assert completed.stdout.strip() == "BloodHoundJWT http://from-config-toml" + + +def test_client_from_secrets_explains_missing_settings(no_credentials, monkeypatch): + with pytest.raises(MissingBloodHoundCredentials, match="No BloodHound credentials"): + bloodhound_client_from_secrets() + + monkeypatch.setenv("DESTINATION__BLOODHOUND__TOKEN", "bearer") + with pytest.raises(MissingBloodHoundCredentials, match="has no 'url'"): + bloodhound_client_from_secrets() + + +# --- The command line -------------------------------------------------------------------- + + +def test_cli_uploads_and_prints_one_line_per_file( + mock_bh_requests, mock_bloodhound_api, bearer_credentials, tmp_path +): + schema_file = write_schema(tmp_path / "schema.json") + + result = CliRunner().invoke(schema_cli, ["upload", str(schema_file)]) + + assert result.exit_code == 0, result.output + assert "created" in result.output + assert list(mock_bloodhound_api.app.state.extensions) == ["GH"] + + +def test_cli_dry_run_needs_no_credentials( + mock_bh_requests, mock_bloodhound_api, no_credentials, tmp_path +): + schema_file = write_schema(tmp_path / "schema.json") + + result = CliRunner().invoke(schema_cli, ["upload", str(schema_file), "--dry-run"]) + + assert result.exit_code == 0, result.output + assert "skipped" in result.output + assert mock_bloodhound_api.app.state.received == [] + + +def test_cli_exits_with_status_one_when_a_file_fails( + mock_bh_requests, mock_bloodhound_api, bearer_credentials, tmp_path +): + broken = tmp_path / "broken.json" + broken.write_text("{ not json") + + result = CliRunner().invoke(schema_cli, ["upload", str(broken)]) + + assert result.exit_code == 1 + assert "not valid JSON" in result.output + + +def test_cli_explains_an_unknown_target(bearer_credentials, monkeypatch): + monkeypatch.setattr("openhound.core.schema_upload.entry_points", lambda group: []) + + result = CliRunner().invoke(schema_cli, ["upload", "nope"]) + + assert result.exit_code == 1 + assert "not the name of an installed collector" in result.output