Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
44 changes: 44 additions & 0 deletions armis_sdk/clients/assets_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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],
Expand Down Expand Up @@ -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],
Expand Down Expand Up @@ -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

Expand Down
8 changes: 7 additions & 1 deletion armis_sdk/entities/asset.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand All @@ -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
7 changes: 7 additions & 0 deletions armis_sdk/entities/asset_field_description.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,7 @@
from armis_sdk.core.base_entity import BaseEntity


class AssetFieldDescription(BaseEntity):
name: str
type: str
is_list: bool = False
32 changes: 32 additions & 0 deletions tests/armis_sdk/clients/assets_client_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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
),
]