Add "DataExportClient"

armis_sdk/clients/data_export_client.py

@@ -0,0 +1,170 @@
+import asyncio
+from typing import AsyncIterator
+from typing import Type
+
+import pandas
+import universalasync
+
+from armis_sdk.core import response_utils
+from armis_sdk.core.armis_error import ArmisError
+from armis_sdk.core.base_entity_client import BaseEntityClient
+from armis_sdk.entities.data_export.base_exported_entity import BaseExportedEntity
+from armis_sdk.entities.data_export.base_exported_entity import T
+from armis_sdk.entities.data_export.data_export import DataExport
+
+
+@universalasync.wrap
+class DataExportClient(BaseEntityClient):
+
+    async def disable(self, entity: Type[BaseExportedEntity]):
+        """Disable data export of the entity.
+
+        Args:
+            entity: The entity to disable exporting for.
+
+        Example:
+            ```python linenums="1" hl_lines="9"
+            import asyncio
+
+            from armis_sdk.clients.data_export_client import DataExportClient
+            from armis_sdk.entities.data_export.application import Application
+
+
+            async def main():
+                data_export_client = DataExportClient()
+                await data_export_client.disable(Application)
+
+            asyncio.run(main())
+            ```
+        """
+        await self.toggle(entity, False)
+
+    async def enable(self, entity: Type[BaseExportedEntity]):
+        """Enable data export of the entity.
+
+        Args:
+            entity: The entity to enable exporting for.
+
+        Example:
+            ```python linenums="1" hl_lines="9"
+            import asyncio
+
+            from armis_sdk.clients.data_export_client import DataExportClient
+            from armis_sdk.entities.data_export.application import Application
+
+
+            async def main():
+                data_export_client = DataExportClient()
+                await data_export_client.enable(Application)
+
+            asyncio.run(main())
+            ```
+        """
+        await self.toggle(entity, True)
+
+    async def iterate(self, entity: Type[T]) -> AsyncIterator[T]:
+        """Iterate over the exported data.
+
+        Returns:
+            An (async) iterator of the underlying entity.
+
+        Example:
+            ```python linenums="1" hl_lines="9"
+            import asyncio
+
+            from armis_sdk.clients.data_export_client import DataExportClient
+            from armis_sdk.entities.data_export.application import Application
+
+
+            async def main():
+                data_export_client = DataExportClient()
+                async for row in data_export_client.iterate(Application):
+                    print(type(row))
+
+            asyncio.run(main())
+            ```
+            Will output:
+            ```python linenums="1"
+            <class 'armis_sdk.entities.data_export.application.Application'>
+            <class 'armis_sdk.entities.data_export.application.Application'>
+            <class 'armis_sdk.entities.data_export.application.Application'>
+            ```
+        """
+        data_export = await self.get(entity)
+        if not data_export.enabled:
+            raise ArmisError(
+                "Data export is disabled for this entity, please enable it first."
+            )
+
+        if data_export.file_format != "parquet":
+            raise ArmisError("Only parquet files supported")
+
+        for url in data_export.urls:
+            df: pandas.DataFrame = await asyncio.to_thread(pandas.read_parquet, url)
+            for _, row in df.iterrows():
+                yield entity.series_to_model(row)
+
+    async def get(self, entity: Type[BaseExportedEntity]) -> DataExport:
+        """Get the `DataExport` of the entity
+
+        Args:
+            entity: The entity to get the data for.
+
+        Returns:
+            A `DataExport` object.
+
+        Example:
+            ```python linenums="1" hl_lines="9"
+            import asyncio
+
+            from armis_sdk.clients.data_export_client import DataExportClient
+            from armis_sdk.entities.data_export.application import Application
+
+
+            async def main():
+                data_export_client = DataExportClient()
+                print(await data_export_client.get(Application))
+
+            asyncio.run(main())
+            ```
+            Will output:
+            ```python linenums="1"
+            DataExport(...)
+            ```
+        """
+        async with self._armis_client.client() as client:
+            response = await client.get(f"/v3/data-export/{entity.entity_name}")
+            data = response_utils.get_data_dict(response)
+            return DataExport.model_validate(data)
+
+    async def toggle(self, entity: Type[BaseExportedEntity], enabled: bool):
+        """Enable / disable export of an entity.
+
+        Args:
+            entity: The entity to enable/disable exporting for.
+            enabled: The new value to set.
+
+        Raises:
+            ResponseError: If an error occurs while communicating with the API.
+
+        Example:
+            ```python linenums="1" hl_lines="9"
+            import asyncio
+
+            from armis_sdk.clients.data_export_client import DataExportClient
+            from armis_sdk.entities.data_export.application import Application
+
+
+            async def main():
+                data_export_client = DataExportClient()
+                await data_export_client.toggle(Application, True)
+
+            asyncio.run(main())
+            ```
+        """
+        data = {"enabled": enabled}
+        async with self._armis_client.client() as client:
+            response = await client.patch(
+                f"/v3/data-export/{entity.entity_name}", json=data
+            )
+            response_utils.raise_for_status(response)

armis_sdk/core/armis_sdk.py

@@ -1,5 +1,6 @@
 from typing import Optional
 
+from armis_sdk.clients.data_export_client import DataExportClient
 from armis_sdk.clients.sites_client import SitesClient
 from armis_sdk.core.armis_client import ArmisClient
 from armis_sdk.core.client_credentials import ClientCredentials
@@ -13,6 +14,7 @@ class ArmisSdk:  # pylint: disable=too-few-public-methods
 
     Attributes:
         client (ArmisClient): An instance of [ArmisClient][armis_sdk.core.armis_client.ArmisClient]
+        data_export (DataExportClient): An instance of [DataExportClient][armis_sdk.clients.data_export_client.DataExportClient]
         sites (SitesClient): An instance of [SitesClient][armis_sdk.clients.sites_client.SitesClient]
 
     Example:
@@ -33,4 +35,5 @@ async def main():
 
     def __init__(self, credentials: Optional[ClientCredentials] = None):
         self.client: ArmisClient = ArmisClient(credentials=credentials)
+        self.data_export: DataExportClient = DataExportClient(self.client)
         self.sites: SitesClient = SitesClient(self.client)

armis_sdk/entities/data_export/application.py

@@ -0,0 +1,66 @@
+import datetime
+from typing import ClassVar
+from typing import Optional
+
+import pandas
+
+from armis_sdk.entities.data_export.base_exported_entity import BaseExportedEntity
+
+
+class Application(BaseExportedEntity):
+    """
+    This class represents an application row that was exported using the data export API.
+    """
+
+    entity_name: ClassVar[str] = "applications"
+
+    device_id: int
+    """The id of the device with the application"""
+
+    vendor: str
+    """
+    The vendor of the application
+
+    **Example**: `Google`
+    """
+
+    name: str
+    """
+    The name of the application
+
+    **Example**: `Chrome`
+    """
+
+    version: str
+    """
+    The version of the application
+
+    **Example**: `30.0.1599.40`
+    """
+
+    cpe: Optional[str]
+    """
+    The CPE (Common Platform Enumeration) of the application
+
+    **Example**: `cpe:2.3:a:google:chrome:30.0.1599.40:*:*:*:*:*:*:*`
+    """
+
+    first_seen: datetime.datetime
+    """When the application was first seen on the device"""
+
+    last_seen: datetime.datetime
+    """When the application was last seen on the device"""
+
+    @classmethod
+    def series_to_model(cls, series: pandas.Series) -> "Application":
+        return Application(
+            device_id=series.loc["device_id"],
+            vendor=series.loc["vendor"],
+            name=series.loc["name"],
+            version=series.loc["version"],
+            cpe=cls._value_or_none(
+                series.loc["cpe"] if "cpe" in series.index else None
+            ),
+            first_seen=series.loc["first_seen"].to_pydatetime(),
+            last_seen=series["last_seen"].to_pydatetime(),
+        )

armis_sdk/entities/data_export/base_exported_entity.py

@@ -0,0 +1,32 @@
+import abc
+from typing import Type
+from typing import TypeVar
+
+import pandas
+from pydantic import BaseModel
+
+T = TypeVar("T", bound="BaseExportedEntity")
+
+
+class BaseExportedEntity(BaseModel, abc.ABC):
+    @classmethod
+    @abc.abstractmethod
+    def series_to_model(cls: Type[T], series: pandas.Series) -> T: ...
+
+    @property
+    @abc.abstractmethod
+    def entity_name(self): ...
+
+    @classmethod
+    def _to_list(cls, value) -> list:
+        return [item for item in value if cls._value_or_none(item)]
+
+    @classmethod
+    def _value_or_none(cls, value):
+        if not value or pandas.isnull(value) or value == "N/A":
+            return None
+
+        if isinstance(value, pandas.Timestamp):
+            return value.to_pydatetime()
+
+        return value

armis_sdk/entities/data_export/data_export.py

@@ -0,0 +1,27 @@
+import datetime
+from typing import Literal
+from typing import Optional
+
+from pydantic import BaseModel
+
+
+class DataExport(BaseModel):
+    """
+    The `DataExport` entity represents the data export of an entity.
+    """
+
+    enabled: bool
+    """Whether the entity is enabled or not for data export."""
+
+    file_format: Literal["parquet"] = "parquet"
+    """
+    The file format of the URLs.
+
+    Currently the only supported format is `parquet`.
+    """
+
+    urls: list[str]
+    """URLs to the files that contain the exported data."""
+
+    urls_creation_time: Optional[datetime.datetime]
+    """The creation time of the URLs."""