[client] Show expectation traces

Author: johanah29

pyobas/apis/__init__.py

@@ -4,6 +4,7 @@
 from .endpoint import *  # noqa: F401,F403
 from .inject import *  # noqa: F401,F403
 from .inject_expectation import *  # noqa: F401,F403
+from .inject_expectation_trace import *  # noqa: F401,F403
 from .injector import *  # noqa: F401,F403
 from .kill_chain_phase import *  # noqa: F401,F403
 from .me import *  # noqa: F401,F403

pyobas/apis/inject_expectation/inject_expectation.py

@@ -20,18 +20,18 @@ class InjectExpectationManager(ListMixin, UpdateMixin, RESTManager):
     _obj_cls = InjectExpectation
     _update_attrs = RequiredOptional(required=("collector_id", "result", "is_success"))
 
-    @exc.on_http_error(exc.OpenBASUpdateError)
-    def expectations_for_source(self, source_id: str, **kwargs: Any) -> Dict[str, Any]:
-        path = f"{self.path}/" + source_id
-        result = self.openbas.http_get(path, **kwargs)
-        return result
-
     @exc.on_http_error(exc.OpenBASUpdateError)
     def expectations_assets_for_source(
-        self, source_id: str, **kwargs: Any
+        self, source_id: str, expiration_time: int = None, **kwargs: Any
     ) -> Dict[str, Any]:
         path = f"{self.path}/assets/" + source_id
-        result = self.openbas.http_get(path, **kwargs)
+        result = self.openbas.http_get(
+            path,
+            query_data=(
+                {"expiration_time": expiration_time} if expiration_time else None
+            ),
+            **kwargs,
+        )
         return result
 
     def expectations_models_for_source(self, source_id: str, **kwargs: Any):
@@ -83,10 +83,16 @@ def prevention_expectations_for_source(
 
     @exc.on_http_error(exc.OpenBASUpdateError)
     def detection_expectations_for_source(
-        self, source_id: str, **kwargs: Any
+        self, source_id: str, expiration_time: int = None, **kwargs: Any
     ) -> Dict[str, Any]:
         path = f"{self.path}/detection/" + source_id
-        result = self.openbas.http_get(path, **kwargs)
+        result = self.openbas.http_get(
+            path,
+            query_data=(
+                {"expiration_time": expiration_time} if expiration_time else None
+            ),
+            **kwargs,
+        )
         return result
 
     @exc.on_http_error(exc.OpenBASUpdateError)

pyobas/apis/inject_expectation_trace.py

@@ -0,0 +1,21 @@
+from pyobas.base import RESTManager, RESTObject
+from pyobas.mixins import CreateMixin
+from pyobas.utils import RequiredOptional
+
+
+class InjectExpectationTrace(RESTObject):
+    _id_attr = "inject_expectation_trace_id"
+
+
+class InjectExpectationTraceManager(CreateMixin, RESTManager):
+    _path = "/inject-expectations-traces"
+    _obj_cls = InjectExpectationTrace
+    _create_attrs = RequiredOptional(
+        required=(
+            "inject_expectation_trace_expectation",
+            "inject_expectation_trace_source_id",
+            "inject_expectation_trace_alert_name",
+            "inject_expectation_trace_alert_link",
+            "inject_expectation_trace_date",
+        ),
+    )

pyobas/client.py

@@ -69,6 +69,7 @@ def __init__(
         self.inject_expectation = apis.InjectExpectationManager(self)
         self.payload = apis.PayloadManager(self)
         self.security_platform = apis.SecurityPlatformManager(self)
+        self.inject_expectation_trace = apis.InjectExpectationTraceManager(self)
 
     @staticmethod
     def _check_redirects(result: requests.Response) -> None: