diff --git a/armis_sdk/clients/__init__.py b/armis_sdk/clients/__init__.py index e69de29..ba5b6f9 100644 --- a/armis_sdk/clients/__init__.py +++ b/armis_sdk/clients/__init__.py @@ -0,0 +1 @@ +# Required by mkdocs diff --git a/armis_sdk/clients/assets_client.py b/armis_sdk/clients/assets_client.py index fda06e4..1872bb0 100644 --- a/armis_sdk/clients/assets_client.py +++ b/armis_sdk/clients/assets_client.py @@ -1,6 +1,5 @@ import datetime from typing import AsyncIterator -from typing import Literal from typing import Optional from typing import Type from typing import Union @@ -15,14 +14,7 @@ from armis_sdk.entities.asset import Asset from armis_sdk.entities.asset import AssetT from armis_sdk.entities.device import Device - -AssetIdSource = Literal[ - "ASSET_ID", - "IPV4_ADDRESS", - "IPV6_ADDRESS", - "MAC_ADDRESS", - "SERIAL_NUMBER", -] +from armis_sdk.types.asset_id_source import AssetIdSource @universalasync.wrap diff --git a/armis_sdk/clients/collectors_client.py b/armis_sdk/clients/collectors_client.py index 2c48cd5..612a39a 100644 --- a/armis_sdk/clients/collectors_client.py +++ b/armis_sdk/clients/collectors_client.py @@ -11,7 +11,7 @@ from armis_sdk.core.base_entity_client import BaseEntityClient from armis_sdk.entities.collector_image import CollectorImage from armis_sdk.entities.download_progress import DownloadProgress -from armis_sdk.enums.collector_image_type import CollectorImageType +from armis_sdk.types.collector_image_type import CollectorImageType @universalasync.wrap @@ -26,7 +26,7 @@ class CollectorsClient(BaseEntityClient): async def download_image( self, destination: Union[str, IO[bytes]], - image_type: CollectorImageType = CollectorImageType.OVA, + image_type: CollectorImageType = "OVA", ) -> AsyncIterator[DownloadProgress]: """Download a collector image to a specified destination path / file. @@ -77,9 +77,7 @@ async def main(): file.write(chunk) yield DownloadProgress(downloaded=file.tell(), total=total_size) - async def get_image( - self, image_type: CollectorImageType = CollectorImageType.OVA - ) -> CollectorImage: + async def get_image(self, image_type: CollectorImageType = "OVA") -> CollectorImage: """Get collector image information including download URL and credentials. Args: @@ -108,7 +106,7 @@ async def main(): """ async with self._armis_client.client() as client: response = await client.get( - "/v3/collectors/_image", params={"image_type": image_type.value} + "/v3/collectors/_image", params={"image_type": image_type} ) data = response_utils.get_data_dict(response) return CollectorImage.model_validate(data) diff --git a/armis_sdk/core/__init__.py b/armis_sdk/core/__init__.py index e69de29..ba5b6f9 100644 --- a/armis_sdk/core/__init__.py +++ b/armis_sdk/core/__init__.py @@ -0,0 +1 @@ +# Required by mkdocs diff --git a/armis_sdk/entities/__init__.py b/armis_sdk/entities/__init__.py index e69de29..ba5b6f9 100644 --- a/armis_sdk/entities/__init__.py +++ b/armis_sdk/entities/__init__.py @@ -0,0 +1 @@ +# Required by mkdocs diff --git a/armis_sdk/entities/collector_image.py b/armis_sdk/entities/collector_image.py index bbb0b3c..6fb7541 100644 --- a/armis_sdk/entities/collector_image.py +++ b/armis_sdk/entities/collector_image.py @@ -3,7 +3,7 @@ from pydantic import Field from armis_sdk.core.base_entity import BaseEntity -from armis_sdk.enums.collector_image_type import CollectorImageType +from armis_sdk.types.collector_image_type import CollectorImageType class CollectorImage(BaseEntity): @@ -11,7 +11,7 @@ class CollectorImage(BaseEntity): An entity that represents the details required to download and run a collector image. """ - image_type: CollectorImageType = Field(strict=False) + image_type: CollectorImageType """The type of the image.""" image_password: str diff --git a/armis_sdk/entities/data_export/__init__.py b/armis_sdk/entities/data_export/__init__.py index e69de29..ba5b6f9 100644 --- a/armis_sdk/entities/data_export/__init__.py +++ b/armis_sdk/entities/data_export/__init__.py @@ -0,0 +1 @@ +# Required by mkdocs diff --git a/armis_sdk/enums/collector_image_type.py b/armis_sdk/enums/collector_image_type.py deleted file mode 100644 index d6c73b5..0000000 --- a/armis_sdk/enums/collector_image_type.py +++ /dev/null @@ -1,36 +0,0 @@ -from enum import Enum - - -class CollectorImageType(Enum): - DARWIN_AMD64_BROKER = "DARWIN_AMD64_BROKER" - """""" - - DARWIN_ARM64_BROKER = "DARWIN_ARM64_BROKER" - """""" - - DEB = "DEB" - """""" - - LINUX_AMD64_BROKER = "LINUX_AMD64_BROKER" - """""" - - LINUX_ARM64_BROKER = "LINUX_ARM64_BROKER" - """""" - - OVA = "OVA" - """""" - - QCOW2 = "QCOW2" - """""" - - RPM = "RPM" - """""" - - VHD = "VHD" - """""" - - VHDX = "VHDX" - """""" - - WINDOWS_BROKER = "WINDOWS_BROKER" - """""" diff --git a/armis_sdk/types/__init__.py b/armis_sdk/types/__init__.py new file mode 100644 index 0000000..ba5b6f9 --- /dev/null +++ b/armis_sdk/types/__init__.py @@ -0,0 +1 @@ +# Required by mkdocs diff --git a/armis_sdk/types/asset_id_source.py b/armis_sdk/types/asset_id_source.py new file mode 100644 index 0000000..3ae0ebd --- /dev/null +++ b/armis_sdk/types/asset_id_source.py @@ -0,0 +1,9 @@ +from typing import Literal + +AssetIdSource = Literal[ + "ASSET_ID", + "IPV4_ADDRESS", + "IPV6_ADDRESS", + "MAC_ADDRESS", + "SERIAL_NUMBER", +] diff --git a/armis_sdk/types/collector_image_type.py b/armis_sdk/types/collector_image_type.py new file mode 100644 index 0000000..eaf844f --- /dev/null +++ b/armis_sdk/types/collector_image_type.py @@ -0,0 +1,15 @@ +from typing import Literal + +CollectorImageType = Literal[ + "DARWIN_AMD64_BROKER", + "DARWIN_ARM64_BROKER", + "DEB", + "LINUX_AMD64_BROKER", + "LINUX_ARM64_BROKER", + "OVA", + "QCOW2", + "RPM", + "VHD", + "VHDX", + "WINDOWS_BROKER", +] diff --git a/docs/clients/assets_client/AssetIdSource.md b/docs/clients/assets_client/AssetIdSource.md new file mode 100644 index 0000000..8162bb0 --- /dev/null +++ b/docs/clients/assets_client/AssetIdSource.md @@ -0,0 +1 @@ +::: armis_sdk.types.asset_id_source.AssetIdSource diff --git a/docs/clients/AssetsClient.md b/docs/clients/assets_client/index.md similarity index 100% rename from docs/clients/AssetsClient.md rename to docs/clients/assets_client/index.md diff --git a/docs/entities/DownloadProgress.md b/docs/clients/collectors_client/DownloadProgress.md similarity index 100% rename from docs/entities/DownloadProgress.md rename to docs/clients/collectors_client/DownloadProgress.md diff --git a/docs/clients/CollectorsClient.md b/docs/clients/collectors_client/index.md similarity index 100% rename from docs/clients/CollectorsClient.md rename to docs/clients/collectors_client/index.md diff --git a/docs/entities/AsqRule.md b/docs/entities/AsqRule.md deleted file mode 100644 index a246566..0000000 --- a/docs/entities/AsqRule.md +++ /dev/null @@ -1 +0,0 @@ -::: armis_sdk.entities.asq_rule.AsqRule diff --git a/docs/entities/Site.md b/docs/entities/Site.md deleted file mode 100644 index 2c929ab..0000000 --- a/docs/entities/Site.md +++ /dev/null @@ -1 +0,0 @@ -::: armis_sdk.entities.site.Site diff --git a/docs/entities/Boundary.md b/docs/entities/asset/device/Boundary.md similarity index 100% rename from docs/entities/Boundary.md rename to docs/entities/asset/device/Boundary.md diff --git a/docs/entities/NetworkInterface.md b/docs/entities/asset/device/NetworkInterface.md similarity index 100% rename from docs/entities/NetworkInterface.md rename to docs/entities/asset/device/NetworkInterface.md diff --git a/docs/entities/Device.md b/docs/entities/asset/device/index.md similarity index 100% rename from docs/entities/Device.md rename to docs/entities/asset/device/index.md diff --git a/docs/entities/Asset.md b/docs/entities/asset/index.md similarity index 100% rename from docs/entities/Asset.md rename to docs/entities/asset/index.md diff --git a/docs/entities/collector_image/CollectorImageType.md b/docs/entities/collector_image/CollectorImageType.md new file mode 100644 index 0000000..ba02419 --- /dev/null +++ b/docs/entities/collector_image/CollectorImageType.md @@ -0,0 +1 @@ +::: armis_sdk.types.collector_image_type.CollectorImageType diff --git a/docs/entities/CollectorImage.md b/docs/entities/collector_image/index.md similarity index 100% rename from docs/entities/CollectorImage.md rename to docs/entities/collector_image/index.md diff --git a/docs/entities/data_export/DataExport.md b/docs/entities/data_export/index.md similarity index 100% rename from docs/entities/data_export/DataExport.md rename to docs/entities/data_export/index.md diff --git a/docs/enums/CollectorImageType.md b/docs/enums/CollectorImageType.md deleted file mode 100644 index 6c351b3..0000000 --- a/docs/enums/CollectorImageType.md +++ /dev/null @@ -1 +0,0 @@ -::: armis_sdk.enums.collector_image_type.CollectorImageType diff --git a/mkdocs.yml b/mkdocs.yml index db56f97..748daf9 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -5,23 +5,31 @@ repo_url: https://github.com/ArmisSecurity/armis-sdk-python nav: - Getting started: index.md - Entities: - - Asset: entities/Asset.md - - AsqRule: entities/AsqRule.md + - Asset: + - entities/asset/index.md + - Device: + - entities/asset/device/index.md + - entities/asset/device/Boundary.md + - entities/asset/device/NetworkInterface.md - DataExport: + - entities/data_export/index.md - Application: entities/data_export/Application.md - - DataExport: entities/data_export/DataExport.md - RiskFactor: entities/data_export/RiskFactor.md - Vulnerability: entities/data_export/Vulnerability.md - - Boundary: entities/Boundary.md - - CollectorImage: entities/CollectorImage.md - - Device: entities/Device.md + - CollectorImage: + - entities/collector_image/index.md + - entities/collector_image/CollectorImageType.md - DeviceCustomProperty: entities/DeviceCustomProperty.md - - DownloadProgress: entities/DownloadProgress.md - - NetworkInterface: entities/NetworkInterface.md - - Site: entities/Site.md + - Site: + - entities/site/index.md + - entities/site/AsqRule.md - Clients: - - AssetsClient: clients/AssetsClient.md - - CollectorsClient: clients/CollectorsClient.md + - AssetsClient: + - clients/assets_client/index.md + - clients/assets_client/AssetIdSource.md + - CollectorsClient: + - clients/collectors_client/index.md + - clients/collectors_client/DownloadProgress.md - DataExportClient: clients/DataExportClient.md - DeviceCustomPropertiesClient: clients/DeviceCustomPropertiesClient.md - SitesClient: clients/SitesClient.md @@ -29,8 +37,6 @@ nav: - ArmisClient: core/ArmisClient.md - ArmisSdk: core/ArmisSdk.md - Errors: core/errors.md - - Enums: - - CollectorImageType: enums/CollectorImageType.md - About Armis: about.md theme: @@ -39,10 +45,10 @@ theme: logo: images/logo.png palette: primary: custom - theme: - features: - - content.code.annotate - - content.code.copy + features: + - content.code.annotate + - content.code.copy + - navigation.indexes plugins: - mkdocstrings: diff --git a/pylintrc b/pylintrc index b6cfa48..b0fe259 100644 --- a/pylintrc +++ b/pylintrc @@ -5,6 +5,7 @@ ignore-paths= .venv disable= + contextmanager-generator-missing-cleanup, duplicate-code, missing-class-docstring, missing-module-docstring, diff --git a/tests/armis_sdk/clients/collectors_client_test.py b/tests/armis_sdk/clients/collectors_client_test.py index 61ea7f9..5e0bb3e 100644 --- a/tests/armis_sdk/clients/collectors_client_test.py +++ b/tests/armis_sdk/clients/collectors_client_test.py @@ -6,7 +6,6 @@ from armis_sdk.clients.collectors_client import CollectorsClient from armis_sdk.entities.collector_image import CollectorImage from armis_sdk.entities.download_progress import DownloadProgress -from armis_sdk.enums.collector_image_type import CollectorImageType pytest_plugins = ["tests.plugins.auto_setup_plugin"] @@ -26,7 +25,7 @@ async def test_get_image(httpx_mock: pytest_httpx.HTTPXMock): collector_image = await collectors_client.get_image() assert collector_image == CollectorImage( - image_type=CollectorImageType.OVA, + image_type="OVA", image_password="test_password", url="https://example.com/collector.ova", url_expiration_date=datetime.datetime(2025, 12, 10), @@ -45,12 +44,10 @@ async def test_get_with_explicit_image_type(httpx_mock: pytest_httpx.HTTPXMock): ) collectors_client = CollectorsClient() - collector_image = await collectors_client.get_image( - image_type=CollectorImageType.QCOW2 - ) + collector_image = await collectors_client.get_image(image_type="QCOW2") assert collector_image == CollectorImage( - image_type=CollectorImageType.QCOW2, + image_type="QCOW2", image_password="test_password_qcow2", url="https://example.com/collector.qcow2", url_expiration_date=datetime.datetime(2025, 12, 11),