[client] Enhance inject expectation APIs

Author: SamuelHassine

pyobas/apis/__init__.py

@@ -1,6 +1,7 @@
 from .attack_pattern import *  # noqa: F401,F403
 from .collector import *  # noqa: F401,F403
 from .document import *  # noqa: F401,F403
+from .endpoint import *  # noqa: F401,F403
 from .inject import *  # noqa: F401,F403
 from .inject_expectation import *  # noqa: F401,F403
 from .injector import *  # noqa: F401,F403

pyobas/apis/endpoint.py

@@ -0,0 +1,19 @@
+from typing import Any, Dict, List
+
+from pyobas import exceptions as exc
+from pyobas.base import RESTManager, RESTObject
+
+
+class Endpoint(RESTObject):
+    _id_attr = "asset_id"
+
+
+class EndpointManager(RESTManager):
+    _path = "/endpoints"
+    _obj_cls = Endpoint
+
+    @exc.on_http_error(exc.OpenBASUpdateError)
+    def get(self, asset_id: str, **kwargs: Any) -> Dict[str, Any]:
+        path = f"{self.path}/" + asset_id
+        result = self.openbas.http_get(path, **kwargs)
+        return result

pyobas/apis/inject_expectation.py

@@ -3,18 +3,31 @@
 from pyobas import exceptions as exc
 from pyobas.base import RESTManager, RESTObject
 from pyobas.mixins import ListMixin, UpdateMixin
+from pyobas.utils import RequiredOptional
 
 
 class InjectExpectation(RESTObject):
-    pass
+    _id_attr = "inject_expectation_id"
 
 
 class InjectExpectationManager(ListMixin, UpdateMixin, RESTManager):
     _path = "/injects/expectations"
     _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 update(
+        self,
+        inject_expectation_id: str,
+        inject_expectation: Dict[str, Any],
+        **kwargs: Any,
+    ) -> Dict[str, Any]:
+        path = f"{self.path}/{inject_expectation_id}"
+        result = self.openbas.http_put(path, post_data=inject_expectation, **kwargs)
+        return result

pyobas/client.py

@@ -64,6 +64,7 @@ def __init__(
         self.kill_chain_phase = apis.KillChainPhaseManager(self)
         self.attack_pattern = apis.AttackPatternManager(self)
         self.team = apis.TeamManager(self)
+        self.endpoint = apis.EndpointManager(self)
         self.user = apis.UserManager(self)
         self.inject_expectation = apis.InjectExpectationManager(self)
 
@@ -294,7 +295,6 @@ def http_put(
     ) -> Union[Dict[str, Any], requests.Response]:
         query_data = query_data or {}
         post_data = post_data or {}
-
         result = self.http_request(
             "put",
             path,

requirements.txt

@@ -1,5 +1,6 @@
 PyYAML~=6.0
 requests~=2.32.2
+datefinder~=0.7.3
 requests-toolbelt~=1.0.0
 dataclasses-json~=0.6.4
 pika~=1.3.0