Resolved issues with nextLink following in Sentinel API calls (#617)

* Made the option to follow nextLinks in Sentinel APIs optional
Added options to pass parameters to Sentinel APIs, specifically for Incidents
Set default limit in incidents returned by list_incidents to 50

* Added return statements to create and update methods for Sentinel APIs

* updated test mocks to reflect newly returned elements

msticpy/_version.py

@@ -1,2 +1,2 @@
 """Version file."""
-VERSION = "2.3.0"
+VERSION = "2.3.1"

msticpy/context/azure/sentinel_incidents.py

@@ -260,17 +260,18 @@ def update_incident(
         if response.status_code not in (200, 201):
             raise CloudError(response=response)
         print("Incident updated.")
+        return response.json().get("name")
 
     def create_incident(  # pylint: disable=too-many-arguments, too-many-locals
         self,
         title: str,
         severity: str,
         status: str = "New",
-        description: str = None,
-        first_activity_time: datetime = None,
-        last_activity_time: datetime = None,
-        labels: List = None,
-        bookmarks: List = None,
+        description: Optional[str] = None,
+        first_activity_time: Optional[datetime] = None,
+        last_activity_time: Optional[datetime] = None,
+        labels: Optional[List] = None,
+        bookmarks: Optional[List] = None,
     ) -> Optional[str]:
         """
         Create a Sentinel Incident.
@@ -433,6 +434,7 @@ def post_comment(
         if response.status_code not in (200, 201):
             raise CloudError(response=response)
         print("Comment posted.")
+        return response.json().get("name")
 
     def add_bookmark_to_incident(self, incident: str, bookmark: str):
         """
@@ -473,11 +475,17 @@ def add_bookmark_to_incident(self, incident: str, bookmark: str):
         if response.status_code not in (200, 201):
             raise CloudError(response=response)
         print("Bookmark added to incident.")
+        return response.json().get("name")
 
-    def list_incidents(self) -> pd.DataFrame:
+    def list_incidents(self, params: Optional[dict] = None) -> pd.DataFrame:
         """
         Get a list of incident for a Sentinel workspace.
 
+        Parameters
+        ----------
+        params : Optional[dict], optional
+            Additional parameters to pass to the API call, by default None
+
         Returns
         -------
         pd.DataFrame
@@ -489,6 +497,8 @@ def list_incidents(self) -> pd.DataFrame:
             If incidents could not be retrieved.
 
         """
-        return self._list_items(item_type="incidents")  # type: ignore
+        if params is None:
+            params = {"$top": 50}
+        return self._list_items(item_type="incidents", params=params)  # type: ignore
 
     get_incidents = list_incidents

msticpy/context/azure/sentinel_utils.py

@@ -39,21 +39,25 @@
 class SentinelUtilsMixin:
     """Mixin class for Sentinel core feature integrations."""
 
-    def _get_items(self, url: str, params: str = "2020-01-01") -> httpx.Response:
+    def _get_items(self, url: str, params: Optional[dict] = None) -> httpx.Response:
         """Get items from the API."""
         self.check_connected()  # type: ignore
+        if params is None:
+            params = {"api-version": "2020-01-01"}
         return httpx.get(
             url,
             headers=get_api_headers(self.token),  # type: ignore
-            params={"api-version": params},
+            params=params,
             timeout=get_http_timeout(),
         )
 
     def _list_items(
         self,
         item_type: str,
         api_version: str = "2020-01-01",
-        appendix: str = None,
+        appendix: Optional[str] = None,
+        next_follow: bool = False,
+        params: Optional[Dict[str, Any]] = None,
     ) -> pd.DataFrame:
         """
         Return lists of core resources from APIs.
@@ -66,7 +70,10 @@ def _list_items(
             The API version to use, by default '2020-01-01'
         appendix: str, optional
             Any appendix that needs adding to the URI, default is None
-
+        next_follow: bool, optional
+            If True, follow the nextLink to get all results, by default False
+        params: Dict, optional
+            Any additional parameters to pass to the API call, by default None
         Returns
         -------
         pd.DataFrame
@@ -81,20 +88,27 @@ def _list_items(
         item_url = self.url + _PATH_MAPPING[item_type]  # type: ignore
         if appendix:
             item_url = item_url + appendix
-        response = self._get_items(item_url, api_version)
+        if params is None:
+            params = {}
+        params["api-version"] = api_version
+        response = self._get_items(item_url, params)
         if response.status_code == 200:
             results_df = _azs_api_result_to_df(response)
         else:
             raise CloudError(response=response)
         j_resp = response.json()
         results = [results_df]
         # If nextLink in response, go get that data as well
-        while "nextLink" in j_resp:
-            next_url = j_resp["nextLink"]
-            next_response = self._get_items(next_url, api_version)
-            next_results_df = _azs_api_result_to_df(next_response)
-            results.append(next_results_df)
-            j_resp = next_response.json()
+        if next_follow:
+            i = 0
+            # Limit to 5 nextLinks to prevent infinite loop
+            while "nextLink" in j_resp and i < 5:
+                next_url = j_resp["nextLink"]
+                next_response = self._get_items(next_url, params)
+                next_results_df = _azs_api_result_to_df(next_response)
+                results.append(next_results_df)
+                j_resp = next_response.json()
+                i += 1
         return pd.concat(results)
 
     def _check_config(self, items: List, workspace_name: Optional[str] = None) -> Dict:
@@ -128,7 +142,10 @@ def _check_config(self, items: List, workspace_name: Optional[str] = None) -> Di
         return config_items
 
     def _build_sent_res_id(
-        self, sub_id: str = None, res_grp: str = None, ws_name: str = None
+        self,
+        sub_id: Optional[str] = None,
+        res_grp: Optional[str] = None,
+        ws_name: Optional[str] = None,
     ) -> str:
         """
         Build a resource ID.
@@ -163,7 +180,7 @@ def _build_sent_res_id(
             ]
         )
 
-    def _build_sent_paths(self, res_id: str, base_url: str = None) -> str:
+    def _build_sent_paths(self, res_id: str, base_url: Optional[str] = None) -> str:
         """
         Build an API URL from an Azure resource ID.
 
@@ -202,7 +219,7 @@ def _build_sent_paths(self, res_id: str, base_url: str = None) -> str:
 
     def check_connected(self):
         """Check that Sentinel workspace is connected."""
-        if not self.connected:
+        if not self.connected:  # type: ignore
             raise MsticpyAzureConnectionError(
                 "Not connected to Sentinel, ensure you run `.connect` before calling functions."
             )

tests/context/azure/test_sentinel_incidents.py

@@ -75,7 +75,7 @@ def test_sent_incidents(sent_loader):
 @respx.mock
 def test_sent_updates(sent_loader):
     """Test Sentinel incident update feature."""
-    respx.put(re.compile(r"https://management\.azure\.com/.*")).respond(201, json="")
+    respx.put(re.compile(r"https://management\.azure\.com/.*")).respond(201, json={'name': '97446b1b-26cf-4034-832b-895da135c535'})
     respx.get(re.compile(r"https://management\.azure\.com/.*")).respond(
         200, json=_INCIDENT
     )
@@ -87,7 +87,7 @@ def test_sent_updates(sent_loader):
 @respx.mock
 def test_sent_comments(sent_loader):
     """Test Sentinel comments feature."""
-    respx.put(re.compile(r"https://management\.azure\.com/.*")).respond(200, json="")
+    respx.put(re.compile(r"https://management\.azure\.com/.*")).respond(200, json={'name': '97446b1b-26cf-4034-832b-895da135c535'})
     respx.get(re.compile(r"https://management\.azure\.com/.*")).respond(
         200, json=_INCIDENT
     )