fix: Add client version update for BHE

src/openhound/__init__.py

@@ -0,0 +1,6 @@
+from importlib.metadata import PackageNotFoundError, version
+
+try:
+    __version__ = version("openhound")
+except PackageNotFoundError:
+    __version__ = "unknown"

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

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())

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
 
 

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()

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