[client] add basic injector contract client (OpenAEV-Platform/openaev#3334)

Signed-off-by: Antoine MAZEAS <[email protected]>

Author: antoinemzs

pyobas/apis/__init__.py

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

pyobas/apis/injector_contract.py

@@ -0,0 +1,56 @@
+from typing import Any, Dict
+
+from pyobas import exceptions as exc
+from pyobas.apis.inputs.search import InjectorContractSearchPaginationInput
+from pyobas.base import RESTManager, RESTObject
+from pyobas.mixins import CreateMixin, DeleteMixin, UpdateMixin
+from pyobas.utils import RequiredOptional
+
+
+class InjectorContract(RESTObject):
+    pass
+
+
+class InjectorContractManager(CreateMixin, UpdateMixin, DeleteMixin, RESTManager):
+    _path = "/injector_contracts"
+    _obj_cls = InjectorContract
+    _create_attrs = RequiredOptional(
+        required=(
+            "contract_content",
+            "contract_id",
+            "contract_labels",
+            "injector_id",
+        ),
+        optional=(
+            "contract_attack_patterns_ids",
+            "contract_attack_patterns_external_ids",
+            "contract_vulnerability_external_ids",
+            "contract_manual",
+            "contract_platforms",
+            "external_contract_id",
+            "is_atomic_testing",
+        ),
+    )
+    _update_attrs = RequiredOptional(
+        required=(
+            "contract_content",
+            "contract_labels",
+        ),
+        optional=(
+            "contract_attack_patterns_ids",
+            "contract_vulnerability_ids",
+            "contract_vulnerability_external_ids",
+            "contract_manual",
+            "contract_platforms",
+            "is_atomic_testing",
+        ),
+    )
+
+    @exc.on_http_error(exc.OpenBASUpdateError)
+    def search(
+        self, input: InjectorContractSearchPaginationInput, **kwargs: Any
+    ) -> Dict[str, Any]:
+        path = f"{self.path}/search"
+        # force the serialisation here since we only need a naive serialisation to json
+        result = self.openbas.http_post(path, post_data=input.to_dict(), **kwargs)
+        return result

pyobas/apis/inputs/__init__.py

@@ -0,0 +1,72 @@
+from typing import Any, Dict, List
+
+
+class Filter:
+    def __init__(self, key: str, mode: str, operator: str, values: List[str]):
+        self.key = key
+        self.mode = mode
+        self.operator = operator
+        self.values = values
+
+    def to_dict(self) -> dict[str, Any]:
+        return {k: v for k, v in self.__dict__.items() if v is not None}
+
+
+class FilterGroup:
+    def __init__(self, mode: str, filters: List[Filter]):
+        self.mode = mode
+        self.filters = filters
+
+    def to_dict(self) -> dict[str, Any]:
+        dictionary: dict[str, Any] = {"mode": self.mode}
+        if self.filters:
+            filter_dicts: List[dict[str, Any]] = []
+            for filter_ in self.filters:
+                filter_dicts.append(filter_.to_dict())
+            dictionary["filters"] = filter_dicts
+        return dictionary
+
+
+class SearchPaginationInput:
+    def __init__(
+        self,
+        page: int,
+        size: int,
+        filter_group: FilterGroup,
+        text_search: str,
+        sorts: Dict[str, str],
+    ):
+        self.size = size
+        self.page = page
+        self.filterGroup = filter_group
+        self.text_search = text_search
+        self.sorts = sorts
+
+    def to_dict(self) -> dict[str, Any]:
+        dictionary: dict[str, Any] = {"page": self.page, "size": self.size}
+        if self.sorts:
+            dictionary["sorts"] = self.sorts
+        if self.text_search:
+            dictionary["textSearch"] = self.text_search
+        if self.filterGroup:
+            dictionary["filterGroup"] = self.filterGroup.to_dict()
+        return dictionary
+
+
+class InjectorContractSearchPaginationInput(SearchPaginationInput):
+    def __init__(
+        self,
+        page: int,
+        size: int,
+        filter_group: FilterGroup,
+        text_search: str = None,
+        sorts: Dict[str, str] = None,
+        include_full_details: bool = True,
+    ):
+        super().__init__(page, size, filter_group, text_search, sorts)
+        self.include_full_details = include_full_details
+
+    def to_dict(self) -> dict[str, Any]:
+        dictionary: dict[str, Any] = super().to_dict()
+        dictionary["include_full_details"] = self.include_full_details
+        return dictionary

pyobas/apis/inputs/search.py

@@ -61,6 +61,7 @@ def __init__(
         self.collector = apis.CollectorManager(self)
         self.cve = apis.CveManager(self)
         self.inject = apis.InjectManager(self)
+        self.injector_contract = apis.InjectorContractManager(self)
         self.document = apis.DocumentManager(self)
         self.kill_chain_phase = apis.KillChainPhaseManager(self)
         self.attack_pattern = apis.AttackPatternManager(self)

pyobas/client.py

@@ -132,15 +132,45 @@ class Contract:
         + VariableHelper.uri_variables()
     )
     contract_attack_patterns_external_ids: List[str] = field(default_factory=list)
+    contract_vulnerability_external_ids: List[str] = field(default_factory=list)
     is_atomic_testing: bool = True
     platforms: List[str] = field(default_factory=list)
+    external_id: str = None
 
     def add_attack_pattern(self, var: str):
         self.contract_attack_patterns_external_ids.append(var)
 
+    def add_vulnerability(self, var: str):
+        self.contract_vulnerability_external_ids.append(var)
+
     def add_variable(self, var: ContractVariable):
         self.variables.append(var)
 
+    def to_contract_add_input(self, source_id: str):
+        return {
+            "contract_id": self.contract_id,
+            "external_contract_id": self.external_id,
+            "injector_id": source_id,
+            "contract_manual": self.manual,
+            "contract_labels": self.label,
+            "contract_attack_patterns_external_ids": self.contract_attack_patterns_external_ids,
+            "contract_vulnerability_external_ids": self.contract_vulnerability_external_ids,
+            "contract_content": json.dumps(self, cls=utils.EnhancedJSONEncoder),
+            "is_atomic_testing": self.is_atomic_testing,
+            "contract_platforms": self.platforms,
+        }
+
+    def to_contract_update_input(self):
+        return {
+            "contract_manual": self.manual,
+            "contract_labels": self.label,
+            "contract_attack_patterns_external_ids": self.contract_attack_patterns_external_ids,
+            "contract_vulnerability_external_ids": self.contract_vulnerability_external_ids,
+            "contract_content": json.dumps(self, cls=utils.EnhancedJSONEncoder),
+            "is_atomic_testing": self.is_atomic_testing,
+            "contract_platforms": self.platforms,
+        }
+
 
 @dataclass
 class ContractTeam(ContractCardinalityElement):

pyobas/contracts/contract_config.py

@@ -216,3 +216,27 @@ def create(
             assert not isinstance(server_data, requests.Response)
             assert self._obj_cls is not None
         return self._obj_cls(self, server_data)
+
+
+class DeleteMixin(_RestManagerBase):
+    _computed_path: Optional[str]
+    _from_parent_attrs: Dict[str, Any]
+    _obj_cls: Optional[Type[base.RESTObject]]
+    _parent: Optional[base.RESTObject]
+    _parent_attrs: Dict[str, Any]
+    _path: Optional[str]
+    openbas: pyobas.OpenBAS
+
+    @exc.on_http_error(exc.OpenBASCreateError)
+    def delete(
+        self, id: Optional[Union[str, int]] = None, **kwargs: Any
+    ) -> requests.Response:
+        if id is None:
+            path = self.path
+        else:
+            path = f"{self.path}/{utils.EncodedId(id)}"
+
+        result = self.openbas.http_delete(path, **kwargs)
+        if TYPE_CHECKING:
+            assert isinstance(result, requests.Response)
+        return result

pyobas/mixins.py

@@ -129,8 +129,7 @@ def add_fields(self, log_record, record, message_dict):
             log_record["level"] = record.levelname
 
 
-def logger(level, json_logging=True):
-    # Exceptions
+def setup_logging_config(level, json_logging=True):
     logging.getLogger("urllib3").setLevel(logging.WARNING)
     logging.getLogger("pika").setLevel(logging.ERROR)
     # Exceptions
@@ -143,34 +142,43 @@ def logger(level, json_logging=True):
     else:
         logging.basicConfig(level=level)
 
-    class AppLogger:
-        def __init__(self, name):
-            self.local_logger = logging.getLogger(name)
 
-        @staticmethod
-        def prepare_meta(meta=None):
-            return None if meta is None else {"attributes": meta}
+class AppLogger:
+    def __init__(self, level, json_logging=True, name: str = __name__):
+        self.log_level = level
+        self.json_logging = json_logging
+        setup_logging_config(self.log_level, self.json_logging)
+        self.local_logger = logging.getLogger(name)
+
+    def __call__(self, name):
+        self.local_logger = logging.getLogger(name)
+        return self
 
-        @staticmethod
-        def setup_logger_level(lib, log_level):
-            logging.getLogger(lib).setLevel(log_level)
+    @staticmethod
+    def prepare_meta(meta=None):
+        return None if meta is None else {"attributes": meta}
 
-        def debug(self, message, meta=None):
-            self.local_logger.debug(message, extra=AppLogger.prepare_meta(meta))
+    @staticmethod
+    def setup_logger_level(lib, log_level):
+        logging.getLogger(lib).setLevel(log_level)
 
-        def info(self, message, meta=None):
-            self.local_logger.info(message, extra=AppLogger.prepare_meta(meta))
+    def debug(self, message, meta=None):
+        self.local_logger.debug(message, extra=AppLogger.prepare_meta(meta))
 
-        def warning(self, message, meta=None):
-            self.local_logger.warning(message, extra=AppLogger.prepare_meta(meta))
+    def info(self, message, meta=None):
+        self.local_logger.info(message, extra=AppLogger.prepare_meta(meta))
 
-        def error(self, message, meta=None):
-            # noinspection PyTypeChecker
-            self.local_logger.error(
-                message, exc_info=1, extra=AppLogger.prepare_meta(meta)
-            )
+    def warning(self, message, meta=None):
+        self.local_logger.warning(message, extra=AppLogger.prepare_meta(meta))
 
-    return AppLogger
+    def error(self, message, meta=None):
+        # noinspection PyTypeChecker
+        self.local_logger.error(message, exc_info=1, extra=AppLogger.prepare_meta(meta))
+
+
+# DEPRECATED: compatibility
+def logger(level, json_logging=True):
+    return AppLogger(level, json_logging)
 
 
 class PingAlive(threading.Thread):

pyobas/utils.py

@@ -0,0 +1,58 @@
+from unittest import TestCase, main, mock
+from unittest.mock import ANY
+
+from pyobas import OpenBAS
+from pyobas.apis.inputs.search import (
+    Filter,
+    FilterGroup,
+    InjectorContractSearchPaginationInput,
+)
+
+
+def mock_response(**kwargs):
+    class MockResponse:
+        def __init__(self, json_data, status_code):
+            self.json_data = json_data
+            self.status_code = status_code
+            self.history = None
+            self.content = None
+            self.headers = {"Content-Type": "application/json"}
+
+        def json(self):
+            return self.json_data
+
+    return MockResponse(None, 200)
+
+
+class TestInjectorContract(TestCase):
+    @mock.patch("requests.Session.request", side_effect=mock_response)
+    def test_search_input_correctly_serialised(self, mock_request):
+        api_client = OpenBAS("url", "token")
+
+        search_input = InjectorContractSearchPaginationInput(
+            0,
+            20,
+            FilterGroup("or", [Filter("prop", "and", "eq", ["titi", "toto"])]),
+            None,
+            None,
+        )
+
+        expected_json = search_input.to_dict()
+        api_client.injector_contract.search(search_input)
+
+        mock_request.assert_called_once_with(
+            method="post",
+            url="url/api/injector_contracts/search",
+            params={},
+            data=None,
+            timeout=None,
+            stream=False,
+            verify=True,
+            json=expected_json,
+            headers=ANY,
+            auth=ANY,
+        )
+
+
+if __name__ == "__main__":
+    main()