diff --git a/censys/asm/clouds.py b/censys/asm/clouds.py index 19c2f90d..06a3b6e0 100644 --- a/censys/asm/clouds.py +++ b/censys/asm/clouds.py @@ -1,25 +1,8 @@ """Interact with the Censys Clouds API.""" -import datetime -from typing import Union - +from ..common.types import Datetime +from ..common.utils import format_iso8601 from .api import CensysAsmAPI -Since = Union[str, datetime.date, datetime.datetime] - - -def format_since_date(since: Since) -> str: - """Formats since date as ISO 8601 format. - - Args: - since (Since): Date. - - Returns: - str: ISO 8601 formatted date string. - """ - if isinstance(since, (datetime.date, datetime.datetime)): - return since.strftime("%Y-%m-%d") - return since - class Clouds(CensysAsmAPI): """Clouds API class.""" @@ -28,55 +11,55 @@ class Clouds(CensysAsmAPI): def get_host_counts( self, - since: Since, + since: Datetime, ) -> dict: """Retrieve host counts by cloud. Hosts found after the date provided in the `since` parameter will be included in the new asset counts. Args: - since (Since): Date to include hosts from. + since (Datetime): Date to include hosts from. Returns: dict: Host count result. """ - since = format_since_date(since) + since = format_iso8601(since) return self._get(f"{self.base_path}/hostCounts/{since}") - def get_domain_counts(self, since: Since) -> dict: + def get_domain_counts(self, since: Datetime) -> dict: """Retrieve domain counts by cloud. Args: - since (Since): Date to include domains from. + since (Datetime): Date to include domains from. Returns: dict: Domain count result. """ - since = format_since_date(since) + since = format_iso8601(since) return self._get(f"{self.base_path}/domainCounts/{since}") - def get_object_store_counts(self, since: Since) -> dict: + def get_object_store_counts(self, since: Datetime) -> dict: """Retrieve object store counts by cloud. Args: - since (Since): Date to include object stores from. + since (Datetime): Date to include object stores from. Returns: dict: Object store count result. """ - since = format_since_date(since) + since = format_iso8601(since) return self._get(f"{self.base_path}/objectStoreCounts/{since}") - def get_subdomain_counts(self, since: Since) -> dict: + def get_subdomain_counts(self, since: Datetime) -> dict: """Retrieve subdomain counts by cloud. Args: - since (Since): Date to include subdomains from. + since (Datetime): Date to include subdomains from. Returns: dict: Subdomain count result. """ - since = format_since_date(since) + since = format_iso8601(since) return self._get(f"{self.base_path}/subdomainCounts/{since}") def get_unknown_counts(self) -> dict: diff --git a/censys/common/utils.py b/censys/common/utils.py index 1c54b729..4099218a 100644 --- a/censys/common/utils.py +++ b/censys/common/utils.py @@ -17,3 +17,17 @@ def format_rfc3339(time: Datetime) -> str: if isinstance(time, (datetime.date, datetime.datetime)): time = time.strftime("%Y-%m-%dT%H:%M:%S.%fZ") return time + + +def format_iso8601(time: Datetime) -> str: + """Formats a datetime object into an ISO8601 string. + + Args: + time (Datetime): Datetime object to format. + + Returns: + str: ISO8601 formatted string. + """ + if isinstance(time, (datetime.date, datetime.datetime)): + return time.strftime("%Y-%m-%d") + return time diff --git a/tests/asm/test_clouds.py b/tests/asm/test_clouds.py index 61c2ddf1..53362218 100644 --- a/tests/asm/test_clouds.py +++ b/tests/asm/test_clouds.py @@ -1,12 +1,8 @@ -import datetime - import responses -from parameterized import parameterized from ..utils import CensysTestCase from .utils import V1_URL from censys.asm.client import AsmClient -from censys.asm.clouds import format_since_date TEST_COUNT_JSON = { "totalAssetCount": 0, @@ -24,17 +20,6 @@ def setUp(self): super().setUp() self.client = AsmClient(self.api_key) - @parameterized.expand( - [ - ["2021-01-01", "2021-01-01"], - [datetime.date(2021, 1, 1), "2021-01-01"], - [datetime.datetime(2021, 1, 1, 12, 15, 20, 40), "2021-01-01"], - ] - ) - def test_format_since_date(self, since, actual): - # Actual call/assertions - assert format_since_date(since) == actual - def test_get_host_counts(self): # Setup response self.responses.add( diff --git a/tests/test_utils.py b/tests/test_utils.py new file mode 100644 index 00000000..23f4f7c2 --- /dev/null +++ b/tests/test_utils.py @@ -0,0 +1,31 @@ +import datetime +import unittest + +from parameterized import parameterized + +from censys.common.utils import format_iso8601, format_rfc3339 + + +class UtilsTest(unittest.TestCase): + @parameterized.expand( + [ + ["2021-01-01", "2021-01-01"], + [datetime.date(2021, 1, 1), "2021-01-01T00:00:00.000000Z"], + [ + datetime.datetime(2021, 1, 1, 12, 15, 20, 40), + "2021-01-01T12:15:20.000040Z", + ], + ] + ) + def test_format_rfc3339(self, since, actual): + assert format_rfc3339(since) == actual + + @parameterized.expand( + [ + ["2021-01-01", "2021-01-01"], + [datetime.date(2021, 1, 1), "2021-01-01"], + [datetime.datetime(2021, 1, 1, 12, 15, 20, 40), "2021-01-01"], + ] + ) + def test_format_iso8601(self, since, actual): + assert format_iso8601(since) == actual