Merge pull request #87 from GunnerUjjwol/add-ueba-status-endpoints

Add ueba status and internal status endpoints
SSHcom · Dec 2, 2022 · 5c56cba · 5c56cba
2 parents abd8758 + 385d602
commit 5c56cba
Show file tree

Hide file tree

Showing 3 changed files with 27 additions and 1 deletion.
diff --git a/.github/workflows/pythonpackage.yml b/.github/workflows/pythonpackage.yml
@@ -12,7 +12,7 @@ on:
 jobs:
   build:
 
-    runs-on: ubuntu-latest
+    runs-on: ubuntu-20.04
     strategy:
       fail-fast: false
       matrix:

diff --git a/privx_api/connection_manager.py b/privx_api/connection_manager.py
@@ -508,3 +508,25 @@ def get_ueba_connection_count(self, payload) -> PrivXAPIResponse:
             body=get_value(payload, dict()),
         )
         return PrivXAPIResponse(response_status, HTTPStatus.OK, data)
+
+    def get_ueba_status(self) -> PrivXAPIResponse:
+        """
+        Get UEBA microservice status
+
+        Returns:
+            PrivXAPIResponse
+        """
+        response_status, data = self._http_get(UrlEnum.CONNECTION_MANAGER.UEBA_STATUS)
+        return PrivXAPIResponse(response_status, HTTPStatus.OK, data)
+
+    def get_ueba_internal_status(self) -> PrivXAPIResponse:
+        """
+        Get UEBA microservice internal status
+
+        Returns:
+            PrivXAPIResponse
+        """
+        response_status, data = self._http_get(
+            UrlEnum.CONNECTION_MANAGER.UEBA_INTERNAL_STATUS
+        )
+        return PrivXAPIResponse(response_status, HTTPStatus.OK, data)
diff --git a/privx_api/enums.py b/privx_api/enums.py
@@ -184,6 +184,8 @@ class ConnectionManagerEnum:
     UEBA_DATASET = "CONNECTION_MANAGER.UEBA_DATASET"
     UEBA_TRAIN_DATASET = "CONNECTION_MANAGER.UEBA_TRAIN_DATASET"
     UEBA_CONNECTION_COUNT = "CONNECTION_MANAGER.UEBA_CONNECTION_COUNT"
+    UEBA_INTERNAL_STATUS = "CONNECTION_MANAGER.UEBA_INTERNAL_STATUS"
+    UEBA_STATUS = "CONNECTION_MANAGER.UEBA_STATUS"
 
     urls = {
         ACCESS_ROLE: "/connection-manager/api/v1/connections/access_roles/{role_id}",
@@ -218,6 +220,8 @@ class ConnectionManagerEnum:
         UEBA_DATASET: "/connection-manager/api/v1/ueba/datasets/{dataset_id}",
         UEBA_TRAIN_DATASET: "/connection-manager/api/v1/ueba/train/{dataset_id}",
         UEBA_CONNECTION_COUNT: "/connection-manager/api/v1/ueba/query-connection-count",
+        UEBA_STATUS: "/connection-manager/api/v1/ueba/status",
+        UEBA_INTERNAL_STATUS: "/connection-manager/api/v1/ueba/status/internal",
     }