From c216679ed3347fa41fca6f09b6eaf66bae574c76 Mon Sep 17 00:00:00 2001
From: Sebastien Quioc <sebastien.quioc@sekoia.fr>
Date: Fri, 13 Dec 2024 13:01:51 +0100
Subject: [PATCH] fix(SentinelOne): fix the pagination for the singularity
 connector

---
 .../singularity/connectors.py                 | 25 ++++++++++---------
 .../tests/singularity/test_connectors.py      |  2 +-
 2 files changed, 14 insertions(+), 13 deletions(-)

diff --git a/SentinelOne/sentinelone_module/singularity/connectors.py b/SentinelOne/sentinelone_module/singularity/connectors.py
index f5a090fea..efc322ca2 100644
--- a/SentinelOne/sentinelone_module/singularity/connectors.py
+++ b/SentinelOne/sentinelone_module/singularity/connectors.py
@@ -36,7 +36,6 @@ def __init__(self, *args: Any, **kwargs: Optional[Any]) -> None:
             start_at=timedelta(days=7),
             ignore_older_than=timedelta(days=7),
         )
-        self.last_checkpoint = CheckpointCursor(path=self.data_path)
 
     @cached_property
     def client(self) -> SingularityClient:
@@ -56,36 +55,38 @@ def stop(self, *args: Any, **kwargs: Optional[Any]) -> None:
     async def single_run(self) -> int:
         result = 0
 
+        # Set up parameters
         last_event_date = self.last_event_date.offset
-        while self.running:
-            cursor: str | None = self.last_checkpoint.offset
-            if cursor == "":  # TODO: Fix this in SDK in cursor handling.
-                cursor = None
+        start_time = int(last_event_date.timestamp())
+        cursor: str | None = None
+        has_more_items: bool = True
 
-            start_time = int(last_event_date.timestamp()) if not cursor else None
+        # Iter over the responses
+        while self.running and has_more_items:
+            # Get the next alerts
             data = await self.client.list_alerts(
                 product_name=self.product_name,
+                start_time=start_time if cursor is None else None,
                 after=cursor,
-                start_time=start_time,
             )
 
+            # Push the collected alerts
             pushed_events = await self.push_data_to_intakes(
                 [orjson.dumps(alert).decode("utf-8") for alert in data.alerts]
             )
 
             result += len(pushed_events)
 
+            # Save the most recent date seen
             for alert in data.alerts:
                 alert_detected_at = isoparse(alert["detectedAt"])
                 if alert_detected_at > last_event_date:
                     last_event_date = alert_detected_at
 
-            self.last_checkpoint.offset = data.end_cursor
-
-            if not data.has_next_page:
-                break
+            # Update parameters for the next page (if exists)
+            cursor = data.end_cursor
+            has_more_items = data.has_next_page
 
-        self.last_checkpoint.offset = ""
         self.last_event_date.offset = last_event_date
 
         return result
diff --git a/SentinelOne/tests/singularity/test_connectors.py b/SentinelOne/tests/singularity/test_connectors.py
index 6069b6a83..67f2a0a4d 100644
--- a/SentinelOne/tests/singularity/test_connectors.py
+++ b/SentinelOne/tests/singularity/test_connectors.py
@@ -46,7 +46,7 @@ def gql_return_values(query, variable_values) -> dict[str, Any]:
         return {
             "alerts": {
                 "totalCount": 0,
-                "pageInfo": {"endCursor": "cursor-1", "hasNextPage": False},
+                "pageInfo": {"endCursor": None, "hasNextPage": False},
                 "edges": [{"node": alert}, {"node": alert}],
             }
         }