diff --git a/armis_sdk/clients/assets_client.py b/armis_sdk/clients/assets_client.py index 1872bb0..21eb22d 100644 --- a/armis_sdk/clients/assets_client.py +++ b/armis_sdk/clients/assets_client.py @@ -13,6 +13,7 @@ from armis_sdk.core.base_entity_client import BaseEntityClient from armis_sdk.entities.asset import Asset from armis_sdk.entities.asset import AssetT +from armis_sdk.entities.asset_field_description import AssetFieldDescription from armis_sdk.entities.device import Device from armis_sdk.types.asset_id_source import AssetIdSource @@ -131,6 +132,42 @@ async def main(): async for item in self._list_assets(asset_class, fields, filter_): yield item + async def list_fields( + self, asset_class: Type[AssetT] + ) -> AsyncIterator[AssetFieldDescription]: + """List all available fields for a given asset class. + + Args: + asset_class: The asset class to list fields for. Must inherit from [Asset][armis_sdk.entities.asset.Asset]. + + Yields: + Field descriptions including field name, type, and other metadata. + + Example: + ```python linenums="1" hl_lines="9" + import asyncio + + from armis_sdk.clients.assets_client import AssetsClient + from armis_sdk.entities.device import Device + + async def main(): + assets_client = AssetsClient() + + async for field in assets_client.list_fields(Device): + print(f"{field.name}: {field.type}") + + asyncio.run(main()) + ``` + """ + async with self._armis_client.client() as client: + response = await client.get( + "/v3/assets/_search/fields", + params={"asset_type": asset_class.asset_type}, + ) + data = response_utils.get_data_dict(response) + for item in data["items"]: + yield AssetFieldDescription.model_validate(item) + async def update( self, assets: list[AssetT], @@ -281,6 +318,10 @@ def _get_device_asset_id( def _is_custom_field(cls, field: str) -> bool: return field.startswith("custom.") + @classmethod + def _is_integration_field(cls, field: str) -> bool: + return field.startswith("integration.") + async def _list_assets( self, asset_class: Type[AssetT], @@ -322,6 +363,9 @@ def _validate_fields( if cls._is_custom_field(field): continue + if cls._is_integration_field(field): + continue + if allow_model_members and field in all_fields: continue diff --git a/armis_sdk/entities/asset.py b/armis_sdk/entities/asset.py index fa6873d..6442bc7 100644 --- a/armis_sdk/entities/asset.py +++ b/armis_sdk/entities/asset.py @@ -22,6 +22,9 @@ class Asset(BaseEntity): custom: dict[str, Any] = Field(default_factory=dict) """Custom properties of the asset. Values can by anything.""" + integration: dict[str, Any] = Field(default_factory=dict) + """Integration properties of the asset. Values can by anything.""" + @classmethod def from_search_result(cls: Type[AssetT], data: dict) -> AssetT: fields: DefaultDict[str, Any] = collections.defaultdict(dict) @@ -38,4 +41,7 @@ def from_search_result(cls: Type[AssetT], data: dict) -> AssetT: def all_fields(cls) -> set[str]: # Pylint doesn't recognize that "cls.model_fields" is a dict and not a method # so it's complaining that the method doesn't have a "keys" attribute. - return set(cls.model_fields.keys()) - {"custom"} # pylint: disable=no-member + return set(cls.model_fields.keys()) - { + "custom", + "integration", + } # pylint: disable=no-member diff --git a/armis_sdk/entities/asset_field_description.py b/armis_sdk/entities/asset_field_description.py new file mode 100644 index 0000000..6d04b44 --- /dev/null +++ b/armis_sdk/entities/asset_field_description.py @@ -0,0 +1,7 @@ +from armis_sdk.core.base_entity import BaseEntity + + +class AssetFieldDescription(BaseEntity): + name: str + type: str + is_list: bool = False diff --git a/tests/armis_sdk/clients/assets_client_test.py b/tests/armis_sdk/clients/assets_client_test.py index 51a5bf7..a631325 100644 --- a/tests/armis_sdk/clients/assets_client_test.py +++ b/tests/armis_sdk/clients/assets_client_test.py @@ -7,6 +7,7 @@ from armis_sdk.core.armis_error import ArmisError from armis_sdk.core.armis_error import BulkUpdateError from armis_sdk.entities.asset import Asset +from armis_sdk.entities.asset_field_description import AssetFieldDescription from armis_sdk.entities.device import Device from tests.armis_sdk.clients import assets_test_data @@ -403,3 +404,34 @@ async def test_update_with_validation_errors(assets, fields, expected_error): with pytest.raises(ArmisError, match=expected_error): await assets_client.update(assets, fields) + + +async def test_list_fields(httpx_mock: pytest_httpx.HTTPXMock): + httpx_mock.add_response( + url="https://api.armis.com/v3/assets/_search/fields?asset_type=DEVICE", + method="GET", + json={ + "items": [ + {"name": "device_id", "type": "integer", "is_list": False}, + {"name": "names", "type": "string", "is_list": True}, + {"name": "custom.Size", "type": "enum", "is_list": False}, + { + "name": "integration.qualys_agent_id", + "type": "string", + "is_list": False, + }, + ] + }, + ) + + assets_client = AssetsClient() + fields = [field async for field in assets_client.list_fields(Device)] + + assert fields == [ + AssetFieldDescription(name="device_id", type="integer", is_list=False), + AssetFieldDescription(name="names", type="string", is_list=True), + AssetFieldDescription(name="custom.Size", type="enum", is_list=False), + AssetFieldDescription( + name="integration.qualys_agent_id", type="string", is_list=False + ), + ]