diff --git a/armis_sdk/__init__.py b/armis_sdk/__init__.py index 67329af..6cc01ba 100644 --- a/armis_sdk/__init__.py +++ b/armis_sdk/__init__.py @@ -1,2 +1,2 @@ -from armis_sdk.core.armis_sdk import ArmisSdk -from armis_sdk.core.client_credentials import ClientCredentials +from armis_sdk.core.armis_sdk import ArmisSdk # noqa: F401 +from armis_sdk.core.client_credentials import ClientCredentials # noqa: F401 diff --git a/armis_sdk/clients/sites_client.py b/armis_sdk/clients/sites_client.py index def82a1..eac1e30 100644 --- a/armis_sdk/clients/sites_client.py +++ b/armis_sdk/clients/sites_client.py @@ -1,5 +1,4 @@ from typing import AsyncIterator -from typing import List import universalasync @@ -129,7 +128,7 @@ async def main(): data = response_utils.get_data_dict(response) return Site.model_validate(data) - async def hierarchy(self) -> List[Site]: + async def hierarchy(self) -> list[Site]: """Create a hierarchy of the tenant's sites, taking into account the parent-child relationships. Returns: diff --git a/armis_sdk/core/armis_error.py b/armis_sdk/core/armis_error.py index 6c12b9d..c6aea88 100644 --- a/armis_sdk/core/armis_error.py +++ b/armis_sdk/core/armis_error.py @@ -4,7 +4,6 @@ """ import json -from typing import List from typing import Optional from typing import Union @@ -26,7 +25,7 @@ def __str__(self): class ErrorBody(BaseModel): - detail: Union[str, List[DetailItem]] + detail: Union[str, list[DetailItem]] class ArmisError(Exception): @@ -63,7 +62,7 @@ class ResponseError(ArmisError): def __init__( self, error_body: ErrorBody, - response_errors: Optional[List[HTTPStatusError]] = None, + response_errors: Optional[list[HTTPStatusError]] = None, ): super().__init__(self._get_message(error_body)) self.response_errors = response_errors diff --git a/armis_sdk/entities/asq_rule.py b/armis_sdk/entities/asq_rule.py index d28ad4a..8b2b5a0 100644 --- a/armis_sdk/entities/asq_rule.py +++ b/armis_sdk/entities/asq_rule.py @@ -1,4 +1,3 @@ -from typing import List from typing import Optional from typing import Union @@ -41,10 +40,10 @@ class AsqRule(BaseEntity): ``` """ - and_: Optional[List[Union[str, "AsqRule"]]] = Field(alias="and", default=None) + and_: Optional[list[Union[str, "AsqRule"]]] = Field(alias="and", default=None) """Rules that all must match.""" - or_: Optional[List[Union[str, "AsqRule"]]] = Field(alias="or", default=None) + or_: Optional[list[Union[str, "AsqRule"]]] = Field(alias="or", default=None) """Rules that at least one of them must match.""" @classmethod diff --git a/armis_sdk/entities/data_export/risk_factor.py b/armis_sdk/entities/data_export/risk_factor.py index 9f734f6..4a86a27 100644 --- a/armis_sdk/entities/data_export/risk_factor.py +++ b/armis_sdk/entities/data_export/risk_factor.py @@ -68,7 +68,7 @@ class RiskFactor(BaseExportedEntity): **Example**: `Device Supports SMBv1` """ - score: int + score: Optional[int] """The score of the risk factor""" group: str @@ -78,14 +78,14 @@ class RiskFactor(BaseExportedEntity): **Example**: `INSECURE_TRAFFIC_AND_BEHAVIOR` """ - remediation_type: str + remediation_type: Optional[str] """ The type of the remediation **Example**: `Disable SMBv1 Protocol` """ - remediation_description: str + remediation_description: Optional[str] """ The description of the remediation @@ -130,15 +130,27 @@ def series_to_model(cls, series: pandas.Series) -> "RiskFactor": category=series.loc["category"], type=series.loc["type"], description=series.loc["description"], - score=series.loc["score"], + score=( + int(score) + if (score := cls._value_or_none(series.loc["score"])) + else None + ), status=series.loc["status"], group=series.loc["group"], - remediation_type=series.loc["remediation"], - remediation_description=series.loc["remediation_description"], - remediation_recommended_actions=[ - RiskFactorRecommendedAction(**item) - for item in json.loads(series.loc["remediation_recommended_actions"]) - ], + remediation_type=cls._value_or_none(series.loc["remediation"]), + remediation_description=cls._value_or_none( + series.loc["remediation_description"] + ), + remediation_recommended_actions=( + [ + RiskFactorRecommendedAction(**item) + for item in json.loads( + series.loc["remediation_recommended_actions"] + ) + ] + if series.loc["remediation_recommended_actions"] + else [] + ), first_seen=series.loc["first_seen"].to_pydatetime(), last_seen=series.loc["last_seen"].to_pydatetime(), status_update_time=cls._value_or_none(series.loc["status_update_time"]), diff --git a/armis_sdk/entities/site.py b/armis_sdk/entities/site.py index 2ab5825..9815d4a 100644 --- a/armis_sdk/entities/site.py +++ b/armis_sdk/entities/site.py @@ -1,7 +1,6 @@ import json from typing import Annotated from typing import Any -from typing import List from typing import Optional from pydantic import BeforeValidator @@ -71,16 +70,16 @@ class Site(BaseEntity): """The ASQ rule of the site.""" network_equipment_device_ids: Annotated[ - Optional[List[int]], BeforeValidator(ensure_list_of_ints) + Optional[list[int]], BeforeValidator(ensure_list_of_ints) ] = None """The ids of network equipment devices associated with the site.""" integration_ids: Annotated[ - Optional[List[int]], BeforeValidator(ensure_list_of_ints) + Optional[list[int]], BeforeValidator(ensure_list_of_ints) ] = None """The ids of the integration associated with the site.""" - children: List["Site"] = Field(default_factory=list) + children: list["Site"] = Field(default_factory=list) """The sub-sites that are directly under this site (their `parent_id` will match this site's `id`).""" diff --git a/pyproject.toml b/pyproject.toml index 3bff884..a450881 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "armis_sdk" -version = "1.1.1" +version = "1.1.2" description = "The Armis SDK is a package that encapsulates common use-cases for interacting with the Armis platform." authors = [ { name = "Shai Lachmanovich", email = "shai@armis.com" }, diff --git a/tests/armis_sdk/clients/assets_client_test.py b/tests/armis_sdk/clients/assets_client_test.py index a631325..3a4ab3e 100644 --- a/tests/armis_sdk/clients/assets_client_test.py +++ b/tests/armis_sdk/clients/assets_client_test.py @@ -374,10 +374,10 @@ async def test_update_with_failed_requests(httpx_mock: pytest_httpx.HTTPXMock): with pytest.raises( BulkUpdateError, match=( - "Failed to update item at index 2. " - 'Request: {"asset_id": 2, "key": "custom.MyField1", ' - '"operation": "SET", "value": "value3"}, ' - 'Response: {"status": 400, "error": "Bad Request"}' + r"Failed to update item at index 2\. " + r'Request: \{"asset_id": 2, "key": "custom\.MyField1", ' + r'"operation": "SET", "value": "value3"\}, ' + r'Response: \{"status": 400, "error": "Bad Request"\}' ), ): await assets_client.update(assets, fields) diff --git a/tests/plugins/auto_setup_plugin.py b/tests/plugins/auto_setup_plugin.py index f2aa9e3..c0bc741 100644 --- a/tests/plugins/auto_setup_plugin.py +++ b/tests/plugins/auto_setup_plugin.py @@ -5,4 +5,4 @@ @pytest.fixture(autouse=True) def auto_setup(setup_env_variables, authorized): # pylint: disable=unused-argument - yield + return