From 13642111e4d54cfe613cb6a6766bfdb2a554a62e Mon Sep 17 00:00:00 2001 From: Nikita Turchaninov Date: Mon, 11 Dec 2023 17:31:24 +0300 Subject: [PATCH] Add tasks api --- cybsi/cloud/client.py | 11 ++ cybsi/cloud/error.py | 69 +++++-- cybsi/cloud/insight/__init__.py | 15 ++ cybsi/cloud/insight/api.py | 20 ++ cybsi/cloud/insight/tasks.py | 218 +++++++++++++++++++++ cybsi/cloud/ioc/__init__.py | 1 + cybsi/cloud/ioc/object.py | 43 ++++ cybsi/cloud/iocean/__init__.py | 3 - cybsi/cloud/iocean/objects.py | 41 +--- docs/api.rst | 7 + docs/index.rst | 1 + docs/user/enrichment.rst | 19 ++ examples/create_enrichment_task.py | 18 ++ examples/get_collection_objects_chained.py | 3 +- examples/get_enrichment_task.py | 22 +++ 15 files changed, 430 insertions(+), 61 deletions(-) create mode 100644 cybsi/cloud/insight/__init__.py create mode 100644 cybsi/cloud/insight/api.py create mode 100644 cybsi/cloud/insight/tasks.py create mode 100644 cybsi/cloud/ioc/__init__.py create mode 100644 cybsi/cloud/ioc/object.py create mode 100644 docs/user/enrichment.rst create mode 100644 examples/create_enrichment_task.py create mode 100644 examples/get_enrichment_task.py diff --git a/cybsi/cloud/client.py b/cybsi/cloud/client.py index aa0b1a3..54e4f10 100644 --- a/cybsi/cloud/client.py +++ b/cybsi/cloud/client.py @@ -1,5 +1,6 @@ from .auth import APIKeyAuth, AuthAPI from .client_config import Config +from .insight.api import InsightAPI, InsightAsyncAPI from .internal import AsyncHTTPConnector, HTTPConnector from .iocean import IOCeanAPI, IOCeanAsyncAPI @@ -77,6 +78,11 @@ def iocean(self) -> IOCeanAPI: """IOCean API handle.""" return IOCeanAPI(self._connector) + @property + def insight(self) -> InsightAPI: + """Insight API handle.""" + return InsightAPI(self._connector) + class AsyncClient: """The asynchronous analog of :class:`Client`. @@ -120,3 +126,8 @@ async def aclose(self) -> None: def iocean(self) -> IOCeanAsyncAPI: """IOCean asynchronous API handle.""" return IOCeanAsyncAPI(self._connector) + + @property + def insight(self) -> InsightAsyncAPI: + """Insight asynchronous API handle.""" + return InsightAsyncAPI(self._connector) diff --git a/cybsi/cloud/error.py b/cybsi/cloud/error.py index 882ba4b..31a0508 100644 --- a/cybsi/cloud/error.py +++ b/cybsi/cloud/error.py @@ -89,8 +89,8 @@ class InvalidRequestError(APIError): InvalidPathArgument = "InvalidPathArgument" InvalidQueryArgument = "InvalidQueryArgument" - def __init__(self, content: JsonObject) -> None: - super().__init__(400, content, header="invalid request") + def __init__(self, resp: httpx.Response) -> None: + super().__init__(400, resp.json(), header="invalid request") class UnauthorizedError(APIError): @@ -98,15 +98,15 @@ class UnauthorizedError(APIError): Unauthorized = "Unauthorized" - def __init__(self, content: JsonObject) -> None: - super().__init__(401, content, header="operation not authorized") + def __init__(self, resp: httpx.Response) -> None: + super().__init__(401, resp.json(), header="operation not authorized") class ForbiddenError(APIError): """Operation was forbidden. Retry will not work unless system is reconfigured.""" - def __init__(self, content: JsonObject) -> None: - super().__init__(403, content, header="operation forbidden") + def __init__(self, resp: httpx.Response) -> None: + super().__init__(403, resp.json(), header="operation forbidden") @property def code(self) -> "ForbiddenErrorCodes": @@ -117,15 +117,15 @@ def code(self) -> "ForbiddenErrorCodes": class NotFoundError(APIError): """Requested resource not found. Retry will never work.""" - def __init__(self, content: JsonObject) -> None: + def __init__(self, resp: httpx.Response) -> None: super().__init__(404, {}, header="resource not found", suffix="") class ConflictError(APIError): """Resource already exists. Retry will never work.""" - def __init__(self, content: JsonObject) -> None: - super().__init__(409, content, header="resource already exists") + def __init__(self, resp: httpx.Response) -> None: + super().__init__(409, resp.json(), header="resource already exists") @property def code(self) -> "ConflictErrorCodes": @@ -139,17 +139,19 @@ class ResourceModifiedError(APIError): Read the updated resource from API, and apply your modifications again. """ - def __init__(self, content: JsonObject) -> None: + def __init__(self, resp: httpx.Response) -> None: super().__init__( - 412, content, header="resource was modified since last read", suffix="" + 412, resp.json(), header="resource was modified since last read", suffix="" ) class RequestEntityTooLargeError(APIError): """Request content is too large.""" - def __init__(self, content: JsonObject) -> None: - super().__init__(413, content, header="request content is too large", suffix="") + def __init__(self, resp: httpx.Response) -> None: + super().__init__( + 413, resp.json(), header="request content is too large", suffix="" + ) class SemanticError(APIError): @@ -160,8 +162,8 @@ class SemanticError(APIError): For example, we're trying to unpack an artifact, but the artifact is not an archive. """ - def __init__(self, content: JsonObject) -> None: - super().__init__(422, content, header="semantic error") + def __init__(self, resp: httpx.Response) -> None: + super().__init__(422, resp.json(), header="semantic error") @property def code(self) -> "SemanticErrorCodes": @@ -169,6 +171,28 @@ def code(self) -> "SemanticErrorCodes": return SemanticErrorCodes(self._view.code) +class TooManyRequestsError(APIError): + """Too many requests error. + + Retry request after some time. + """ + + def __init__(self, resp: httpx.Response) -> None: + super().__init__(429, resp.json(), header="too many requests error") + retry_after = resp.headers.get("Retry-After") + self._retry_after = int(retry_after) if retry_after is not None else None + + @property + def code(self) -> "TooManyRequestsErrorCodes": + """Error code.""" + return TooManyRequestsErrorCodes(self._view.code) + + @property + def retry_after(self) -> Optional[int]: + """Period in seconds after which request could be repeated.""" + return self._retry_after + + @document_enum class ForbiddenErrorCodes(CybsiAPIEnum): """Possible error codes of :class:`ForbiddenError`.""" @@ -212,6 +236,18 @@ class SemanticErrorCodes(CybsiAPIEnum): CursorOutOfRange = "CursorOutOfRange" """Cursor for collection changes is obsolete.""" + # Filebox + FileNotFound = "FileNotFound" + """File not found.""" + + +@document_enum +class TooManyRequestsErrorCodes(CybsiAPIEnum): + """Too many requests error codes.""" + + LimitExceeded = "LimitExceeded" + """Request limit exceeded.""" + class ErrorView(dict): """Error returned by Cybsi Cloud API.""" @@ -266,13 +302,14 @@ def message(self) -> str: 412: ResourceModifiedError, 413: RequestEntityTooLargeError, 422: SemanticError, + 429: TooManyRequestsError, } def _raise_cybsi_error(resp: httpx.Response) -> None: err_cls = _error_mapping.get(resp.status_code, None) if err_cls is not None: - raise err_cls(resp.json()) + raise err_cls(resp) raise CybsiError( f"unexpected response status code: {resp.status_code}. " diff --git a/cybsi/cloud/insight/__init__.py b/cybsi/cloud/insight/__init__.py new file mode 100644 index 0000000..7359e94 --- /dev/null +++ b/cybsi/cloud/insight/__init__.py @@ -0,0 +1,15 @@ +"""Use this section of API to access Insight schemas. +""" + +from .api import InsightAPI, InsightAsyncAPI +from .tasks import ( + ObjectKeyForm, + TaskAPI, + TaskAsyncAPI, + TaskErrorView, + TaskForm, + TaskParamsView, + TaskRegistrationView, + TaskState, + TaskView, +) diff --git a/cybsi/cloud/insight/api.py b/cybsi/cloud/insight/api.py new file mode 100644 index 0000000..2fad611 --- /dev/null +++ b/cybsi/cloud/insight/api.py @@ -0,0 +1,20 @@ +from ..internal import BaseAPI, BaseAsyncAPI +from .tasks import TaskAPI, TaskAsyncAPI + + +class InsightAPI(BaseAPI): + """Insight API.""" + + @property + def tasks(self) -> TaskAPI: + """Get Insight task handle.""" + return TaskAPI(self._connector) + + +class InsightAsyncAPI(BaseAsyncAPI): + """Insight asynchronous API.""" + + @property + def tasks(self) -> TaskAsyncAPI: + """Tasks asynchronous API handle.""" + return TaskAsyncAPI(self._connector) diff --git a/cybsi/cloud/insight/tasks.py b/cybsi/cloud/insight/tasks.py new file mode 100644 index 0000000..3b410dc --- /dev/null +++ b/cybsi/cloud/insight/tasks.py @@ -0,0 +1,218 @@ +from typing import Iterable, List, Optional, cast + +from enum_tools import document_enum + +from ..enum import CybsiAPIEnum +from ..internal import BaseAPI, BaseAsyncAPI, JsonObject, JsonObjectForm, JsonObjectView +from ..ioc import ObjectKeyType, ObjectKeyView + +_PATH = "/insight/tasks" + + +@document_enum +class TaskState(CybsiAPIEnum): + """Object key type.""" + + Pending = "Pending" + """Task awaits execution.""" + Completed = "Completed" + """Task successfully completed.""" + Failed = "Failed" + """Task completed with error.""" + + +class TaskAPI(BaseAPI): + """Task API.""" + + def register(self, task: "TaskForm") -> "TaskRegistrationView": + """Register new enrichment task. + + Note: + Calls `POST /insight/tasks`. + Args: + task: Task registration form. + Returns: + Task registration view. + Raises: + :class:`~cybsi.cloud.error.InvalidRequestError`: + Provided values are invalid (see form value requirements). + :class:`~cybsi.cloud.error.SemanticError`: Form contains logic errors. + :class:`~cybsi.cloud.error.TooManyRequestsError`: Request limit exceeded. + Note: + Semantic error codes specific for this method: + * :attr:`~cybsi.cloud.error.SemanticErrorCodes.SchemaNotFound` + * :attr:`~cybsi.cloud.error.SemanticErrorCodes.FileNotFound` + * :attr:`~cybsi.cloud.error.SemanticErrorCodes.InvalidKeySet` + * :attr:`~cybsi.cloud.error.SemanticErrorCodes.InvalidKeyFormat` + Too many requests error codes specific for this method: + * :attr:`~cybsi.cloud.error.TooManyRequestsErrorCodes.LimitExceeded` + """ + resp = self._connector.do_post(path=_PATH, json=task.json()) + return TaskRegistrationView(resp.json()) + + def view(self, task_id: str) -> "TaskView": + """Get the enrichment task view. + + Note: + Calls `GET /insight/tasks/{task_id}`. + Args: + task_id: Task identifier. + Returns: + Task view. + Raises: + :class:`~cybsi.cloud.error.NotFoundError`: Task not found. + """ + + path = f"{_PATH}/{task_id}" + resp = self._connector.do_get(path=path) + return TaskView(resp.json()) + + +class TaskAsyncAPI(BaseAsyncAPI): + """Tasks asynchronous API.""" + + async def register(self, task: "TaskForm") -> "TaskRegistrationView": + """Register new enrichment task. + + Note: + Calls `POST /insight/tasks`. + Args: + task: Task registration form. + Returns: + Task registration view. + Raises: + :class:`~cybsi.cloud.error.InvalidRequestError`: + Provided values are invalid (see form value requirements). + :class:`~cybsi.cloud.error.SemanticError`: Form contains logic errors. + :class:`~cybsi.cloud.error.TooManyRequestsError`: Request limit exceeded. + Note: + Semantic error codes specific for this method: + * :attr:`~cybsi.cloud.error.SemanticErrorCodes.SchemaNotFound` + * :attr:`~cybsi.cloud.error.SemanticErrorCodes.FileNotFound` + * :attr:`~cybsi.cloud.error.SemanticErrorCodes.InvalidKeySet` + * :attr:`~cybsi.cloud.error.SemanticErrorCodes.InvalidKeyFormat` + Too many requests error codes specific for this method: + * :attr:`~cybsi.cloud.error.TooManyRequestsErrorCodes.LimitExceeded` + """ + resp = await self._connector.do_post(path=_PATH, json=task.json()) + return TaskRegistrationView(resp.json()) + + async def view(self, task_id: str) -> "TaskView": + """Get the enrichment task view. + + Note: + Calls `GET /insight/tasks/{task_id}`. + Args: + task_id: Task identifier. + Returns: + Task view. + Raises: + :class:`~cybsi.cloud.error.NotFoundError`: Task not found. + """ + + path = f"{_PATH}/{task_id}" + resp = await self._connector.do_get(path=path) + return TaskView(resp.json()) + + +class TaskRegistrationView(JsonObjectView): + """Task registration view""" + + @property + def id(self) -> str: + """Task identifier.""" + return self._get("taskID") + + +class TaskParamsView(JsonObjectView): + """Task params view.""" + + @property + def schema_id(self) -> str: + """URL friendly string, uniquely identifies json schema.""" + return self._get("schemaID") + + @property + def object_keys(self) -> List[ObjectKeyView]: + """List of object keys.""" + return self._map_list_optional("objectKeys", ObjectKeyView) + + @property + def file_id(self) -> Optional[str]: + """File identifier. Use filebox to upload file.""" # TODO: filebox ref + return self._get_optional("fileID") + + +class TaskErrorView(JsonObjectView): + """Task error view.""" + + @property + def code(self) -> str: + """Error code.""" + return self._get("code") + + +class TaskView(JsonObjectView): + """Task view""" + + @property + def id(self) -> str: + """Task identifier.""" + return self._get("taskID") + + @property + def params(self) -> TaskParamsView: + """Task params.""" + return TaskParamsView(self._get("params")) + + @property + def state(self) -> TaskState: + """Task state.""" + return TaskState(self._get("state")) + + @property + def result(self) -> JsonObject: + """Task result. + + Note: + This value is present if task state is :attr:`~.TaskState.Completed` + """ + return cast(JsonObject, self._get("result")) + + @property + def error(self) -> TaskErrorView: + """Task error. + + Note: + This value is present if task state is :attr:`~.TaskState.Failed` + """ + return TaskErrorView(self._get("result")) + + +class ObjectKeyForm(JsonObjectForm): + """Object key form.""" + + def __init__(self, key_type: ObjectKeyType, value: str): + super().__init__() + self._data["type"] = key_type.value + self._data["value"] = value + + +class TaskForm(JsonObjectForm): + """Task form""" + + def __init__( + self, + schema_id: str, + object_keys: Optional[Iterable[ObjectKeyForm]] = None, + file_id: Optional[str] = None, + ): + super().__init__() + params: dict = { + "schemaID": schema_id, + } + if object_keys is not None: + params["objectKeys"] = [key.json() for key in object_keys] + if file_id is not None: + params["fileID"] = file_id + self._data["params"] = params diff --git a/cybsi/cloud/ioc/__init__.py b/cybsi/cloud/ioc/__init__.py new file mode 100644 index 0000000..2345f95 --- /dev/null +++ b/cybsi/cloud/ioc/__init__.py @@ -0,0 +1 @@ +from .object import ObjectKeyType, ObjectKeyView, ObjectType diff --git a/cybsi/cloud/ioc/object.py b/cybsi/cloud/ioc/object.py new file mode 100644 index 0000000..ebbf4d2 --- /dev/null +++ b/cybsi/cloud/ioc/object.py @@ -0,0 +1,43 @@ +from enum_tools import document_enum + +from ..enum import CybsiAPIEnum +from ..internal import JsonObjectView + + +@document_enum +class ObjectKeyType(CybsiAPIEnum): + """Object key type.""" + + MD5Hash = "MD5Hash" + SHA1Hash = "SHA1Hash" + SHA256Hash = "SHA256Hash" + SHA512Hash = "SHA512Hash" + DomainName = "DomainName" + URL = "URL" + IPAddress = "IPAddress" + IPNetwork = "IPNetwork" + + +@document_enum +class ObjectType(CybsiAPIEnum): + """Object type.""" + + File = "File" + DomainName = "DomainName" + URL = "URL" + IPAddress = "IPAddress" + IPNetwork = "IPNetwork" + + +class ObjectKeyView(JsonObjectView): + """Object key view""" + + @property + def type(self) -> ObjectKeyType: + """Object key type""" + return ObjectKeyType.from_string(self._get("type")) + + @property + def value(self) -> str: + """Key value.""" + return self._get("value") diff --git a/cybsi/cloud/iocean/__init__.py b/cybsi/cloud/iocean/__init__.py index ed4cd30..4763f84 100644 --- a/cybsi/cloud/iocean/__init__.py +++ b/cybsi/cloud/iocean/__init__.py @@ -13,10 +13,7 @@ from .objects import ( ObjectAPI, ObjectsAsyncAPI, - ObjectKeyType, - ObjectType, ObjectView, - ObjectKeyView, ObjectChangeView, ObjectOperation, ) diff --git a/cybsi/cloud/iocean/objects.py b/cybsi/cloud/iocean/objects.py index ceb9f3f..2f00320 100644 --- a/cybsi/cloud/iocean/objects.py +++ b/cybsi/cloud/iocean/objects.py @@ -6,34 +6,9 @@ from ..enum import CybsiAPIEnum from ..internal import BaseAPI, BaseAsyncAPI, JsonObject, JsonObjectView +from ..ioc import ObjectKeyType, ObjectKeyView, ObjectType from ..pagination import AsyncPage, Cursor, Page - -@document_enum -class ObjectKeyType(CybsiAPIEnum): - """Object key type.""" - - MD5Hash = "MD5Hash" - SHA1Hash = "SHA1Hash" - SHA256Hash = "SHA256Hash" - SHA512Hash = "SHA512Hash" - DomainName = "DomainName" - URL = "URL" - IPAddress = "IPAddress" - IPNetwork = "IPNetwork" - - -@document_enum -class ObjectType(CybsiAPIEnum): - """Object type.""" - - File = "File" - DomainName = "DomainName" - URL = "URL" - IPAddress = "IPAddress" - IPNetwork = "IPNetwork" - - _PATH = "/iocean/collections/{}/objects" @@ -378,20 +353,6 @@ def _extract_changes_cursor(resp: httpx.Response) -> Optional[Cursor]: return cast(Optional[Cursor], cursor[0]) if cursor is not None else None -class ObjectKeyView(JsonObjectView): - """Object key view""" - - @property - def type(self) -> ObjectKeyType: - """Object key type""" - return ObjectKeyType.from_string(self._get("type")) - - @property - def value(self) -> str: - """Key value.""" - return self._get("value") - - class ObjectView(JsonObjectView): """Object view.""" diff --git a/docs/api.rst b/docs/api.rst index f6d5057..2ed1887 100644 --- a/docs/api.rst +++ b/docs/api.rst @@ -35,6 +35,13 @@ Auth :inherited-members: :exclude-members: APIKeyAuth +IOC +~~~~~~ +.. automodule:: cybsi.cloud.ioc + :members: + :imported-members: + :inherited-members: + IOCean ~~~~~~ .. automodule:: cybsi.cloud.iocean diff --git a/docs/index.rst b/docs/index.rst index 4fd0a46..63ef34a 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -36,6 +36,7 @@ instructions for getting the most out of Cybsi Cloud SDK. user/data-model user/advanced user/authentication + user/enrichment The API Documentation / Guide diff --git a/docs/user/enrichment.rst b/docs/user/enrichment.rst new file mode 100644 index 0000000..2e1408f --- /dev/null +++ b/docs/user/enrichment.rst @@ -0,0 +1,19 @@ +.. _insight: + +Insight Service +======================== + +.. _enrichment_tasks: + +Enrichment tasks +---------------- + +Insight service enriches object data. + +In the example bellow you create new enrichment task. + +.. literalinclude:: ../../examples/create_enrichment_task.py + +And example of getting enrichment task result. + +.. literalinclude:: ../../examples/get_enrichment_task.py diff --git a/examples/create_enrichment_task.py b/examples/create_enrichment_task.py new file mode 100644 index 0000000..5c98063 --- /dev/null +++ b/examples/create_enrichment_task.py @@ -0,0 +1,18 @@ +#!/usr/bin/env python3 +from cybsi.cloud import Client, Config +from cybsi.cloud.insight import ObjectKeyForm, TaskForm +from cybsi.cloud.ioc import ObjectKeyType + +if __name__ == "__main__": + config = Config(api_key="the cryptic string") + + with Client(config) as client: + schema_id = "phishing" + + # define task params + keys = [ObjectKeyForm(ObjectKeyType.DomainName, "example.com")] + task_form = TaskForm(schema_id, keys) + # create task + task = client.insight.tasks.register(task_form) + # store task identifier + print(task.id) diff --git a/examples/get_collection_objects_chained.py b/examples/get_collection_objects_chained.py index 75d06b1..f3b488c 100644 --- a/examples/get_collection_objects_chained.py +++ b/examples/get_collection_objects_chained.py @@ -10,8 +10,7 @@ # Retrieve collection schema, it describes all attributes # of objects you can encounter in the collection. - schema_view = client.iocean.collections.view_schema( - collection_id=collection_id) + schema_view = client.iocean.collections.view_schema(collection_id=collection_id) print(schema_view.schema) # Retrieve first page of collection objects. diff --git a/examples/get_enrichment_task.py b/examples/get_enrichment_task.py new file mode 100644 index 0000000..7fc8195 --- /dev/null +++ b/examples/get_enrichment_task.py @@ -0,0 +1,22 @@ +#!/usr/bin/env python3 + +from cybsi.cloud import Client, Config +from cybsi.cloud.insight import TaskState + +if __name__ == "__main__": + config = Config(api_key="the cryptic string") + + with Client(config) as client: + # task identifier + task_id = "d3bcba79-2a0f-41ec-a3da-52e82eea4b2b" + # get task view + task = client.insight.tasks.view(task_id) + if task.state == TaskState.Completed: + # handle result + print(task.result) + elif task.state == TaskState.Failed: + # handle task error + print(task.error) + else: + # retry after some time + pass