From fb370149fcc188b5f883e35a54ca57f31c3674e9 Mon Sep 17 00:00:00 2001
From: Thomas Harning Jr <thomas@censys.io>
Date: Mon, 21 Oct 2024 11:31:04 -0400
Subject: [PATCH] fix(risks): adds paging to risk_types API (#656)

* fix(risks): adds paging to risk_types API

* build: bump version from 2.2.15 to 2.2.16
---
 censys/asm/risks.py     | 10 +++++++++-
 pyproject.toml          |  2 +-
 tests/asm/test_risks.py |  2 ++
 3 files changed, 12 insertions(+), 2 deletions(-)

diff --git a/censys/asm/risks.py b/censys/asm/risks.py
index f429665f..8e9e5445 100644
--- a/censys/asm/risks.py
+++ b/censys/asm/risks.py
@@ -126,6 +126,8 @@ def patch_risk_instance(self, risk_instance_id: int, data: dict) -> dict:
 
     def get_risk_types(
         self,
+        limit: Optional[int] = None,
+        page: Optional[int] = None,
         sort: Optional[List[str]] = None,
         include_events: Optional[bool] = None,
         accept: Optional[str] = None,
@@ -133,6 +135,8 @@ def get_risk_types(
         """Retrieve risk types.
 
         Args:
+            limit (int, optional): Maximum number of results to return. Defaults to 1000.
+            page (int, optional): Page number to begin at when searching. Defaults to 1.
             sort (list): Optional; Sort by field(s).
             include_events (bool): Optional; Whether to include events.
             accept (str): Optional; Accept header.
@@ -140,7 +144,11 @@ def get_risk_types(
         Returns:
             dict: Risk types result.
         """
-        args = {"sort": sort, "includeEvents": include_events}
+        args: Dict[str, Any] = {"sort": sort, "includeEvents": include_events}
+        if page:
+            args["page"] = page
+        if limit:
+            args["limit"] = limit
         return self._get(
             self.risk_types_path,
             args=args,
diff --git a/pyproject.toml b/pyproject.toml
index 4bf450f2..97bb829a 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "censys"
-version = "2.2.15"
+version = "2.2.16"
 description = "An easy-to-use and lightweight API wrapper for Censys APIs (censys.io)."
 authors = ["Censys, Inc. <support@censys.io>"]
 license = "Apache-2.0"
diff --git a/tests/asm/test_risks.py b/tests/asm/test_risks.py
index f1aaa56b..dd66d7d7 100644
--- a/tests/asm/test_risks.py
+++ b/tests/asm/test_risks.py
@@ -217,6 +217,8 @@ def test_patch_risk_instance(self):
             ({"include_events": True}, "?includeEvents=True"),
             ({"include_events": False}, "?includeEvents=False"),
             ({"sort": ["severity", "type:asc"]}, "?sort=severity&sort=type:asc"),
+            ({"page": 1, "limit": 10000}, "?page=1&limit=10000"),
+            ({"page": 10}, "?page=10"),
         ]
     )
     def test_get_risk_types(self, kwargs, params):