diff --git a/censys/asm/risks.py b/censys/asm/risks.py index 336e0ec3..f429665f 100644 --- a/censys/asm/risks.py +++ b/censys/asm/risks.py @@ -1,5 +1,5 @@ """Interact with the Censys Risks API.""" -from typing import List, Optional +from typing import Any, Dict, List, Optional from .api import CensysAsmAPI @@ -8,9 +8,49 @@ class Risks(CensysAsmAPI): """Risks API class.""" base_path = "/v2/risk" + risk_events_path = f"{base_path}-events" risk_instances_path = f"{base_path}-instances" risk_types_path = f"{base_path}-types" + def get_risk_events( + self, + start: Optional[str] = None, + end: Optional[str] = None, + after_id: Optional[int] = None, + limit: Optional[int] = None, + cursor: Optional[str] = None, + accept: Optional[str] = None, + ) -> dict: + """Retrieve risk events. + + Args: + start (str): Optional; Starting event time, inclusive (in RFC3339 format). + end (str): Optional; Ending event time, inclusive (in RFC3339 format). + after_id (int): Optional; Risk event ID to query for events after. + limit (int): Optional; Max number of events to return. + cursor (str): Optional; Cursor value to continue collecting events started in a previous request. + accept (str): Optional; Accept header. + + Returns: + dict: Risk events result. + """ + args: Dict[str, Any] = {} + if start: + args["start"] = start + if end: + args["end"] = end + if after_id: + args["afterID"] = after_id + if limit: + args["limit"] = limit + if cursor: + args["cursor"] = cursor + return self._get( + self.risk_events_path, + args=args, + headers={"Accept": accept} if accept else None, + ) + def get_risk_instances( self, include_events: Optional[bool] = None, accept: Optional[str] = None ) -> dict: diff --git a/docs/usage-asm.rst b/docs/usage-asm.rst index 5db04fbf..3f4849bf 100644 --- a/docs/usage-asm.rst +++ b/docs/usage-asm.rst @@ -250,6 +250,10 @@ Below we show an example of **getting risk instances**. r = Risks() + # Get risk events + risk_events = r.get_risk_events() + print(risk_events) + # Get a dict that returns all risk instances risk_instances = r.get_risk_instances() print(risk_instances) diff --git a/tests/asm/test_risks.py b/tests/asm/test_risks.py index 51c7e041..f1aaa56b 100644 --- a/tests/asm/test_risks.py +++ b/tests/asm/test_risks.py @@ -15,6 +15,12 @@ "srcID": "string", "ts": "2022-02-15T20:22:30.262Z", } +TEST_RISK_EVENTS_JSON = { + "total": 1, + "next": "eyJhZnRlcklEIjo3NzQwLCJsaW1pdCI6MTAwfQ==", + "events": [TEST_EVENT_JSON], + "endOfEvents": True, +} TEST_RISK_INSTANCE_JSON = { "id": 0, "events": [TEST_EVENT_JSON], @@ -70,6 +76,33 @@ def setUp(self): super().setUp() self.setUpApi(Risks(self.api_key)) + @parameterized.expand( + [ + ({}, ""), + ( + {"start": "2023-12-14T03:27:00Z", "end": "2022-03-18T09:18:54Z"}, + "?start=2023-12-14T03:27:00Z&end=2022-03-18T09:18:54Z", + ), + ({"after_id": 489, "limit": 100}, "?afterID=489&limit=100"), + ( + {"cursor": "eyJhZnRlcklEIjo3NzQwLCJsaW1pdCI6MTAwfQ=="}, + "?cursor=eyJhZnRlcklEIjo3NzQwLCJsaW1pdCI6MTAwfQ==", + ), + ] + ) + def test_get_risk_events(self, kwargs, params): + # Setup response + self.responses.add( + responses.GET, + V2_URL + f"/risk-events{params}", + status=200, + json=TEST_RISK_EVENTS_JSON, + ) + # Actual call + res = self.api.get_risk_events(**kwargs) + # Assertions + assert res == TEST_RISK_EVENTS_JSON + @parameterized.expand( [ ({}, ""),