diff --git a/src/openhound/__init__.py b/src/openhound/__init__.py index e69de29..26bb1fd 100644 --- a/src/openhound/__init__.py +++ b/src/openhound/__init__.py @@ -0,0 +1,6 @@ +from importlib.metadata import PackageNotFoundError, version + +try: + __version__ = version("openhound") +except PackageNotFoundError: + __version__ = "unknown" diff --git a/src/openhound/core/clients/bloodhound.py b/src/openhound/core/clients/bloodhound.py index a70dff1..1906270 100644 --- a/src/openhound/core/clients/bloodhound.py +++ b/src/openhound/core/clients/bloodhound.py @@ -5,12 +5,13 @@ import logging from abc import ABC, abstractmethod from datetime import timedelta -from importlib.metadata import version import requests from dlt.common import json from dlt.common.exceptions import DltException +import openhound + from .models import ( AssetGroupsTags, CustomNodes, @@ -23,9 +24,6 @@ logger = logging.getLogger(__name__) -__version__ = version("openhound") - - class BloodHoundHTTPError(DltException): def __init__(self, reason: str, code: int): self.reason = reason @@ -179,7 +177,7 @@ def request( sig = base64.b64encode(digester.digest()).decode() headers = { - "User-Agent": f"openhound/{__version__}", + "User-Agent": f"openhound/{openhound.__version__}", "Authorization": f"bhesignature {self.token_id}", "RequestDate": datetime_formatted, "Signature": sig, @@ -207,7 +205,7 @@ def request( extra_headers: dict[str, str] | None = None, ): headers = { - "User-Agent": f"openhound/{__version__}", + "User-Agent": f"openhound/{openhound.__version__}", "Content-Type": "application/json", "Authorization": f"Bearer {self.token}", } diff --git a/src/openhound/core/clients/bloodhound_enterprise.py b/src/openhound/core/clients/bloodhound_enterprise.py index a76d3cb..a5d31c4 100644 --- a/src/openhound/core/clients/bloodhound_enterprise.py +++ b/src/openhound/core/clients/bloodhound_enterprise.py @@ -1,7 +1,9 @@ import gzip import json +import socket from enum import Enum +import openhound from openhound.core.clients.bloodhound import BloodHound from openhound.core.clients.models.jobs import ( JobsAvailable, @@ -52,3 +54,27 @@ def ingest(self, data: str) -> None: self.request( method="POST", path=path, body=compressed_data, extra_headers=headers ) + + def update_client_metadata(self) -> None: + path = "/api/v2/clients/update" + try: + hostname = socket.gethostname() + except OSError: + hostname = "unknown" + + if hostname == "unknown": + ip_address = "unknown" + else: + try: + ip_address = socket.gethostbyname(hostname) + except OSError: + ip_address = "unknown" + + payload = { + "Address": ip_address, + "Hostname": hostname, + "Version": openhound.__version__, + } + body = json.dumps(payload) + + self.request(method="PUT", path=path, body=body.encode()) diff --git a/src/openhound/core/logging.py b/src/openhound/core/logging.py index a53233b..2582719 100644 --- a/src/openhound/core/logging.py +++ b/src/openhound/core/logging.py @@ -5,16 +5,14 @@ import sys import time from enum import Enum -from importlib.metadata import version from logging.handlers import TimedRotatingFileHandler from pathlib import Path import dlt +import openhound from rich.console import Console from rich.logging import RichHandler -__version__ = version("openhound") - VALID_LEVELS = ["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"] VALID_FORMATS = ["json", "text"] @@ -140,7 +138,7 @@ def format(self, record: logging.LogRecord) -> str: "function": record.funcName, "line": record.lineno, "message": record.getMessage(), - "openhound_version": __version__, + "openhound_version": openhound.__version__, } if record.exc_info: @@ -209,7 +207,7 @@ def format(self, record: logging.LogRecord) -> str: Returns: str: A formatted string for Rich logging """ - log_fmt = f"time={self.formatTime(record, '%Y-%m-%d %H:%M:%S')}, msg={record.getMessage()} (openhound_version={__version__})" + log_fmt = f"time={self.formatTime(record, '%Y-%m-%d %H:%M:%S')}, msg={record.getMessage()} (openhound_version={openhound.__version__})" return log_fmt diff --git a/src/openhound/scheduler/service.py b/src/openhound/scheduler/service.py index cb25ac5..2b24643 100644 --- a/src/openhound/scheduler/service.py +++ b/src/openhound/scheduler/service.py @@ -236,6 +236,12 @@ def start(self) -> None: logger.info( f"Service started, monitoring {self.bhe_uri} every {self.interval} seconds." ) + try: + self.client.update_client_metadata() + + except Exception as err: + logger.exception(f"Unable to update client metadata: {err}") + try: while not self.exit: self._poll() diff --git a/tests/test_bhe_job_scheduling.py b/tests/test_bhe_job_scheduling.py index 13ea67c..6e712a0 100644 --- a/tests/test_bhe_job_scheduling.py +++ b/tests/test_bhe_job_scheduling.py @@ -9,6 +9,7 @@ from fastapi import FastAPI, Request, Response from fastapi.testclient import TestClient +from openhound.core.clients import bloodhound_enterprise from openhound.core.clients.bloodhound_enterprise import JobStatus from openhound.core.models.graph import Graph from openhound.scheduler import service as scheduler_service @@ -40,6 +41,7 @@ def mock_bloodhound_api(): app.state.job_ended = False app.state.end_payload = None app.state.start_payload = None + app.state.client_update_payload = None app.state.ingested_edges = 0 app.state.ingested_nodes = 0 @@ -74,6 +76,11 @@ async def ingest(request: Request): app.state.ingested_edges += len(validate_graph.graph.edges) return {"status": "success"} + @app.put("/api/v2/clients/update") + async def update_client(body: dict): + app.state.client_update_payload = body + return {"status": "success"} + return TestClient(app) @@ -104,6 +111,8 @@ def mock_request(method, url, **kwargs): return mock_bloodhound_api.get(path) if method.upper() == "POST": return mock_bloodhound_api.post(path, **kwargs) + if method.upper() == "PUT": + return mock_bloodhound_api.put(path, **kwargs) raise AssertionError(f"Unhandled method: {method}") @@ -118,6 +127,63 @@ def mock_request(method, url, **kwargs): ) +def test_client_update_sends_metadata(mock_service, mock_bloodhound_api, monkeypatch): + monkeypatch.setattr(bloodhound_enterprise.openhound, "__version__", "1.2.3") + monkeypatch.setattr(bloodhound_enterprise.socket, "gethostname", lambda: "test-host") + monkeypatch.setattr( + bloodhound_enterprise.socket, + "gethostbyname", + lambda hostname: "192.0.2.10", + ) + + mock_service.client.update_client_metadata() + + assert mock_bloodhound_api.app.state.client_update_payload == { + "Address": "192.0.2.10", + "Hostname": "test-host", + "Version": "1.2.3", + } + + +def test_client_update_uses_unknown_when_hostname_lookup_fails( + mock_service, mock_bloodhound_api, monkeypatch +): + monkeypatch.setattr(bloodhound_enterprise.openhound, "__version__", "1.2.3") + + def raise_error(): + raise OSError("hostname unavailable") + + monkeypatch.setattr(bloodhound_enterprise.socket, "gethostname", raise_error) + + mock_service.client.update_client_metadata() + + assert mock_bloodhound_api.app.state.client_update_payload == { + "Address": "unknown", + "Hostname": "unknown", + "Version": "1.2.3", + } + + +def test_client_update_uses_unknown_when_ip_lookup_fails( + mock_service, mock_bloodhound_api, monkeypatch +): + monkeypatch.setattr(bloodhound_enterprise.openhound, "__version__", "1.2.3") + monkeypatch.setattr(bloodhound_enterprise.socket, "gethostname", lambda: "test-host") + + def raise_error(hostname: str): + raise OSError(f"{hostname} unavailable") + + monkeypatch.setattr(bloodhound_enterprise.socket, "gethostbyname", raise_error) + + mock_service.client.update_client_metadata() + + assert mock_bloodhound_api.app.state.client_update_payload == { + "Address": "unknown", + "Hostname": "test-host", + "Version": "1.2.3", + } + + def test_jobs_starts_new_job(mock_service, mock_bloodhound_api): """Runs _check_jobs and checks if the new job is started when available."""