diff --git a/armis_sdk/clients/data_export_client.py b/armis_sdk/clients/data_export_client.py index 9b1589f..2a016be 100644 --- a/armis_sdk/clients/data_export_client.py +++ b/armis_sdk/clients/data_export_client.py @@ -100,8 +100,10 @@ async def main(): 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(): + data_frame: pandas.DataFrame = await asyncio.to_thread( + pandas.read_parquet, url + ) + for _, row in data_frame.iterrows(): yield entity.series_to_model(row) async def get(self, entity: Type[BaseExportedEntity]) -> DataExport: diff --git a/armis_sdk/clients/device_custom_properties_client.py b/armis_sdk/clients/device_custom_properties_client.py new file mode 100644 index 0000000..3d435ec --- /dev/null +++ b/armis_sdk/clients/device_custom_properties_client.py @@ -0,0 +1,230 @@ +from typing import AsyncIterator + +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.device_custom_property import DeviceCustomProperty + + +@universalasync.wrap +class DeviceCustomPropertiesClient(BaseEntityClient): + """ + A client for interacting with device custom properties. + + The primary entity for this client is + [DeviceCustomProperty][armis_sdk.entities.device_custom_property.DeviceCustomProperty]. + """ + + async def create(self, property_: DeviceCustomProperty) -> DeviceCustomProperty: + # pylint: disable=line-too-long + """Create a `DeviceCustomProperty`. + + Args: + property_: The `DeviceCustomProperty` to create. + + Returns: + The same property as the input with the addition of id. + + Example: + Example: + ```python linenums="1" hl_lines="10" + import asyncio + + from armis_sdk.clients.device_custom_properties_client import DeviceCustomPropertiesClient + from armis_sdk.entities.device_custom_property import DeviceCustomProperty + + + async def main(): + client = DeviceCustomPropertiesClient() + property_ = DeviceCustomProperty(name="MyConfig", type="string") + print(await client.create(property_)) + + asyncio.run(main()) + ``` + Will output: + ```python linenums="1" + DeviceCustomPropertiesClient(id=1, name="MyConfig", type="string") + ``` + """ + if property_.id is not None: + raise ArmisError( + "Can't create a property that already has an id. " + "Did you mean to call `.update(property_)`?", + ) + + if not property_.name: + raise ArmisError("Can't create a property without a name.") + + if not property_.type: + raise ArmisError("Can't create a property without a type.") + + payload = property_.model_dump(exclude_none=True) + + async with self._armis_client.client() as client: + response = await client.post( + "/v3/settings/device-custom-properties", json=payload + ) + data = response_utils.get_data_dict(response) + return DeviceCustomProperty.model_validate(data) + + async def delete(self, property_: DeviceCustomProperty): + # pylint: disable=line-too-long + """Delete a `DeviceCustomProperty`. + + Args: + property_: The `DeviceCustomProperty` to delete. + + Example: + Example: + ```python linenums="1" hl_lines="10" + import asyncio + + from armis_sdk.clients.device_custom_properties_client import DeviceCustomPropertiesClient + from armis_sdk.entities.device_custom_property import DeviceCustomProperty + + + async def main(): + client = DeviceCustomPropertiesClient() + property_ = DeviceCustomProperty(id=1, name="MyConfig", type="string") + await client.delete(property_) + + asyncio.run(main()) + ``` + """ + if property_.id is None: + raise ArmisError("Can't delete a property without an id.") + + async with self._armis_client.client() as client: + response = await client.delete( + f"/v3/settings/device-custom-properties/{property_.id}" + ) + response_utils.raise_for_status(response) + + async def get(self, property_id: int) -> DeviceCustomProperty: + # pylint: disable=line-too-long + """Get a `DeviceCustomProperty` by its ID. + + Args: + property_id: The ID of the `DeviceCustomProperty` to get. + + Returns: + A `DeviceCustomProperty` object. + + Example: + Example: + ```python linenums="1" hl_lines="8" + import asyncio + + from armis_sdk.clients.device_custom_properties_client import DeviceCustomPropertiesClient + + + async def main(): + client = DeviceCustomPropertiesClient() + print(await client.get(1)) + + asyncio.run(main()) + ``` + Will output: + ```python linenums="1" + DeviceCustomPropertiesClient(id=1, name="MyConfig", type="string") + ``` + """ + async with self._armis_client.client() as client: + response = await client.get( + f"/v3/settings/device-custom-properties/{property_id}" + ) + data = response_utils.get_data_dict(response) + return DeviceCustomProperty.model_validate(data) + + async def list(self) -> AsyncIterator[DeviceCustomProperty]: + # pylint: disable=line-too-long + """List all the tenant's `DeviceCustomProperty`s. + This method takes care of pagination, so you don't have to deal with it. + + Returns: + An (async) iterator of `DeviceCustomProperty` object. + + Example: + ```python linenums="1" hl_lines="8" + import asyncio + + from armis_sdk.clients.device_custom_properties_client import DeviceCustomPropertiesClient + + + async def main(): + client = DeviceCustomPropertiesClient() + async for property_ in client.list() + print(property_) + + asyncio.run(main()) + ``` + Will output: + ```python linenums="1" + DeviceCustomPropertiesClient(id=1, name="MyConfig", type="string") + DeviceCustomPropertiesClient(id=1, name="MyOtherConfig", type="integer") + ``` + """ + + async with self._armis_client.client() as client: + # endpoint doesn't support paging + response = await client.get("/v3/settings/device-custom-properties") + data = response_utils.get_data_dict(response) + for item in data["items"]: + yield DeviceCustomProperty.model_validate(item) + + async def update(self, property_: DeviceCustomProperty) -> DeviceCustomProperty: + # pylint: disable=line-too-long + """Update a `DeviceCustomProperty`. + Only `description` and `allowed_values` are updatable. + + Args: + property_: The `DeviceCustomProperty` to update. + + Raises: + ResponseError: If an error occurs while communicating with the API. + + Example: + ```python linenums="1" hl_lines="15" + import asyncio + + from armis_sdk.clients.device_custom_properties_client import DeviceCustomPropertiesClient + from armis_sdk.entities.device_custom_property import DeviceCustomProperty + + + async def main(): + client = DeviceCustomPropertiesClient() + property_ = DeviceCustomProperty( + id=1, + name="MyConfig", + type="string", + description="New description", + ) + await client.update(property_) + + asyncio.run(main()) + ``` + """ + if property_.id is None: + raise ArmisError( + "Can't update a property without an id. " + "Did you mean to call `.create(property_)`?", + ) + + data = property_.model_dump( + exclude={"id", "name", "creation_time", "created_by", "type"}, + exclude_none=True, + ) + + if not data: + return property_ + + async with self._armis_client.client() as client: + response = await client.patch( + f"/v3/settings/device-custom-properties/{property_.id}", + json=data, + ) + response_utils.raise_for_status(response) + + return property_ diff --git a/armis_sdk/clients/sites_client.py b/armis_sdk/clients/sites_client.py index b6a10f7..def82a1 100644 --- a/armis_sdk/clients/sites_client.py +++ b/armis_sdk/clients/sites_client.py @@ -22,7 +22,7 @@ async def create(self, site: Site) -> Site: """Create a `Site`. Args: - site: The site to delete. + site: The site to create. Returns: The same site as the input with the addition of id. diff --git a/armis_sdk/core/armis_sdk.py b/armis_sdk/core/armis_sdk.py index e376f6a..ce78abc 100644 --- a/armis_sdk/core/armis_sdk.py +++ b/armis_sdk/core/armis_sdk.py @@ -1,6 +1,9 @@ from typing import Optional from armis_sdk.clients.data_export_client import DataExportClient +from armis_sdk.clients.device_custom_properties_client import ( + DeviceCustomPropertiesClient, +) from armis_sdk.clients.sites_client import SitesClient from armis_sdk.core.armis_client import ArmisClient from armis_sdk.core.client_credentials import ClientCredentials @@ -15,6 +18,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] + device_custom_properties (DeviceCustomPropertiesClient): An instance of [DeviceCustomPropertiesClient][armis_sdk.clients.device_custom_properties_client.DeviceCustomPropertiesClient] sites (SitesClient): An instance of [SitesClient][armis_sdk.clients.sites_client.SitesClient] Example: @@ -36,4 +40,7 @@ async def main(): def __init__(self, credentials: Optional[ClientCredentials] = None): self.client: ArmisClient = ArmisClient(credentials=credentials) self.data_export: DataExportClient = DataExportClient(self.client) + self.device_custom_properties: DeviceCustomPropertiesClient = ( + DeviceCustomPropertiesClient(self.client) + ) self.sites: SitesClient = SitesClient(self.client) diff --git a/armis_sdk/entities/data_export/data_export.py b/armis_sdk/entities/data_export/data_export.py index 29b2301..e46476d 100644 --- a/armis_sdk/entities/data_export/data_export.py +++ b/armis_sdk/entities/data_export/data_export.py @@ -3,6 +3,7 @@ from typing import Optional from pydantic import BaseModel +from pydantic import Field class DataExport(BaseModel): @@ -23,5 +24,5 @@ class DataExport(BaseModel): urls: list[str] """URLs to the files that contain the exported data.""" - urls_creation_time: Optional[datetime.datetime] + urls_creation_time: Optional[datetime.datetime] = Field(strict=False) """The creation time of the URLs.""" diff --git a/armis_sdk/entities/device_custom_property.py b/armis_sdk/entities/device_custom_property.py new file mode 100644 index 0000000..dec7f24 --- /dev/null +++ b/armis_sdk/entities/device_custom_property.py @@ -0,0 +1,53 @@ +import datetime +from typing import Literal +from typing import Optional + +from pydantic import Field + +from armis_sdk.core.base_entity import BaseEntity + + +class DeviceCustomProperty(BaseEntity): + id: Optional[int] = None + """The id of the property.""" + + name: str = Field(max_length=40, pattern=r"^[\w_]*$") + """ + The name of the property. + + Example: `Size` + """ + + description: Optional[str] = Field(max_length=250, default=None) + """ + The description of the property. + + Example: `The size of the device` + """ + + type: Literal[ + "boolean", + "enum", + "externalLink", + "integer", + "string", + "timestamp", + ] + """ + The type of the property. + + Example: `enum` + """ + + allowed_values: Optional[list[str]] = None + """ + The allowed values of the property when the 'type' is 'enum'. + + Example: `["s", "m", "l"]` + """ + + created_by: Optional[str] = Field(max_length=50, default=None) + """Who / what created the property.""" + + creation_time: Optional[datetime.datetime] = Field(strict=False, default=None) + """The creation time of the property.""" diff --git a/docs/clients/DeviceCustomPropertiesClient.md b/docs/clients/DeviceCustomPropertiesClient.md new file mode 100644 index 0000000..33f2d44 --- /dev/null +++ b/docs/clients/DeviceCustomPropertiesClient.md @@ -0,0 +1 @@ +::: armis_sdk.clients.device_custom_properties_client.DeviceCustomPropertiesClient diff --git a/docs/entities/DeviceCustomProperty.md b/docs/entities/DeviceCustomProperty.md new file mode 100644 index 0000000..f477e86 --- /dev/null +++ b/docs/entities/DeviceCustomProperty.md @@ -0,0 +1 @@ +::: armis_sdk.entities.device_custom_property.DeviceCustomProperty diff --git a/mkdocs.yml b/mkdocs.yml index 3fd0874..e2133a1 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -11,9 +11,11 @@ nav: - DataExport: entities/data_export/DataExport.md - RiskFactor: entities/data_export/RiskFactor.md - Vulnerability: entities/data_export/Vulnerability.md + - DeviceCustomProperty: entities/DeviceCustomProperty.md - Site: entities/Site.md - Clients: - DataExportClient: clients/DataExportClient.md + - DeviceCustomPropertiesClient: clients/DeviceCustomPropertiesClient.md - SitesClient: clients/SitesClient.md - Core: - ArmisClient: core/ArmisClient.md diff --git a/tests/armis_sdk/clients/device_custom_properties_client_test.py b/tests/armis_sdk/clients/device_custom_properties_client_test.py new file mode 100644 index 0000000..d3ad0af --- /dev/null +++ b/tests/armis_sdk/clients/device_custom_properties_client_test.py @@ -0,0 +1,234 @@ +import datetime + +import pytest +import pytest_httpx + +from armis_sdk.clients.device_custom_properties_client import ( + DeviceCustomPropertiesClient, +) +from armis_sdk.core.armis_error import ArmisError +from armis_sdk.entities.device_custom_property import DeviceCustomProperty + +pytest_plugins = ["tests.plugins.auto_setup_plugin"] + + +async def test_create(httpx_mock: pytest_httpx.HTTPXMock): + httpx_mock.add_response( + url="https://api.armis.com/v3/settings/device-custom-properties", + method="POST", + match_json={ + "name": "mock_name", + "description": "mock_description", + "type": "enum", + "allowed_values": ["a", "b", "c"], + }, + json={ + "id": 1, + "name": "mock_name", + "description": "mock_description", + "type": "enum", + "allowed_values": ["a", "b", "c"], + "created_by": "mock_created_by", + "creation_time": "2025-11-20T00:00:00", + }, + ) + + client = DeviceCustomPropertiesClient() + property_to_create = DeviceCustomProperty( + name="mock_name", + description="mock_description", + type="enum", + allowed_values=["a", "b", "c"], + ) + + created_property = await client.create(property_to_create) + + assert created_property == DeviceCustomProperty( + id=1, + name="mock_name", + description="mock_description", + type="enum", + allowed_values=["a", "b", "c"], + created_by="mock_created_by", + creation_time=datetime.datetime(2025, 11, 20), + ) + + +async def test_create_with_id(httpx_mock: pytest_httpx.HTTPXMock): + httpx_mock.reset() + + client = DeviceCustomPropertiesClient() + property_ = DeviceCustomProperty(id=1, name="mock_name", type="string") + + with pytest.raises( + ArmisError, + match=( + r"Can't create a property that already has an id. " + r"Did you mean to call `\.update\(property_\)`?" + ), + ): + await client.create(property_) + + +async def test_delete(httpx_mock: pytest_httpx.HTTPXMock): + httpx_mock.add_response( + url="https://api.armis.com/v3/settings/device-custom-properties/1", + method="DELETE", + ) + + client = DeviceCustomPropertiesClient() + property_ = DeviceCustomProperty(id=1, name="mock_name", type="string") + + await client.delete(property_) + + +async def test_delete_without_id(httpx_mock: pytest_httpx.HTTPXMock): + httpx_mock.reset() + + client = DeviceCustomPropertiesClient() + property_ = DeviceCustomProperty(name="mock_name", type="string") + + with pytest.raises(ArmisError, match=r"Can't delete a property without an id."): + await client.delete(property_) + + +async def test_get(httpx_mock: pytest_httpx.HTTPXMock): + httpx_mock.add_response( + url="https://api.armis.com/v3/settings/device-custom-properties/1", + method="GET", + json={ + "id": 1, + "name": "mock_name", + "description": "mock_description", + "type": "enum", + "allowed_values": ["a", "b", "c"], + "created_by": "mock_created_by", + "creation_time": "2025-11-20T00:00:00", + }, + ) + + client = DeviceCustomPropertiesClient() + + property_ = await client.get(1) + + assert property_ == DeviceCustomProperty( + id=1, + name="mock_name", + description="mock_description", + type="enum", + allowed_values=["a", "b", "c"], + created_by="mock_created_by", + creation_time=datetime.datetime(2025, 11, 20), + ) + + +@pytest.mark.parametrize( + ["from_response", "expected"], + [ + pytest.param( + { + "id": 1, + "name": "mock_name1", + "type": "string", + "created_by": "mock_created_by1", + "creation_time": "2025-11-20T00:00:00", + }, + DeviceCustomProperty( + id=1, + name="mock_name1", + type="string", + created_by="mock_created_by1", + creation_time=datetime.datetime(2025, 11, 20), + ), + id="Only mandatory fields", + ), + pytest.param( + { + "id": 2, + "name": "mock_name2", + "description": "mock_description2", + "type": "enum", + "allowed_values": ["a", "b", "c"], + "created_by": "mock_created_by2", + "creation_time": "2025-11-20T00:00:00", + }, + DeviceCustomProperty( + id=2, + name="mock_name2", + description="mock_description2", + type="enum", + allowed_values=["a", "b", "c"], + created_by="mock_created_by2", + creation_time=datetime.datetime(2025, 11, 20), + ), + id="All fields", + ), + ], +) +async def test_list_properties( + from_response, expected, httpx_mock: pytest_httpx.HTTPXMock +): + httpx_mock.add_response( + url="https://api.armis.com/v3/settings/device-custom-properties", + method="GET", + json={"items": [from_response]}, + ) + + client = DeviceCustomPropertiesClient() + properties = [property_ async for property_ in client.list()] + + assert properties == [expected] + + +async def test_update(httpx_mock: pytest_httpx.HTTPXMock): + httpx_mock.add_response( + url="https://api.armis.com/v3/settings/device-custom-properties/1", + method="PATCH", + match_json={"description": "new_description"}, + json={ + "id": 1, + "name": "mock_name", + "description": "new_description", + "type": "string", + }, + ) + + client = DeviceCustomPropertiesClient() + property_ = DeviceCustomProperty( + id=1, name="mock_name", type="string", description="new_description" + ) + + updated_property = await client.update(property_) + assert updated_property == DeviceCustomProperty( + id=1, + name="mock_name", + description="new_description", + type="string", + ) + + +async def test_update_with_nothing_to_change(httpx_mock: pytest_httpx.HTTPXMock): + httpx_mock.reset() + + client = DeviceCustomPropertiesClient() + property_ = DeviceCustomProperty(id=1, name="mock_name", type="string") + + updated_property = await client.update(property_) + + assert updated_property == property_ + + +async def test_update_without_id(httpx_mock: pytest_httpx.HTTPXMock): + httpx_mock.reset() + + client = DeviceCustomPropertiesClient() + property_ = DeviceCustomProperty(name="mock_name", type="string") + + with pytest.raises( + ArmisError, + match=( + r"Can't update a property without an id. " + r"Did you mean to call `\.create\(property_\)`?" + ), + ): + await client.update(property_)