Mimecast - use cursor from SDK

Mimecast/mimecast_modules/connector_mimecast_siem.py

@@ -9,6 +9,7 @@
 import requests
 from dateutil.parser import isoparse
 from pyrate_limiter import Duration, Limiter, RequestRate
+from sekoia_automation.checkpoint import CheckpointDatetime
 from sekoia_automation.connector import Connector, DefaultConnectorConfiguration
 from sekoia_automation.storage import PersistentJSON
 
@@ -33,9 +34,15 @@ def __init__(self, connector: "MimecastSIEMConnector", log_type: str):
         self.connector = connector
         self.log_type: str = log_type
 
-        self.context = self.connector.context
-        self.from_date = self.most_recent_date_seen
+        self.cursor = CheckpointDatetime(
+            path=self.connector._data_path,
+            start_at=timedelta(days=1),
+            ignore_older_than=timedelta(days=7),
+            lock=self.connector.context_lock,
+            subkey=self.log_type,
+        )
 
+        self.from_date = self.cursor.offset
         self._stop_event = Event()
 
     def log(self, *args, **kwargs):
@@ -59,40 +66,6 @@ def client(self) -> ApiClient:
             limiter_default=self.connector.limiter_default,
         )
 
-    @property
-    def most_recent_date_seen(self):
-        now = datetime.now(timezone.utc)
-
-        self.connector.context_lock.acquire()
-        with self.context as cache:
-            most_recent_date_seen_str = cache.get(self.log_type, {}).get("most_recent_date_seen")
-        self.connector.context_lock.release()
-
-        # if undefined, retrieve events from the last day
-        if most_recent_date_seen_str is None:
-            return now - timedelta(days=1)
-
-        # parse the most recent date seen
-        most_recent_date_seen = isoparse(most_recent_date_seen_str)
-
-        # We don't retrieve messages older than one day
-        one_day_ago = now - timedelta(days=7)
-        if most_recent_date_seen < one_day_ago:
-            most_recent_date_seen = one_day_ago
-
-        return most_recent_date_seen
-
-    @most_recent_date_seen.setter
-    def most_recent_date_seen(self, dt: datetime) -> None:
-        self.connector.context_lock.acquire()
-        with self.context as cache:
-            if self.log_type not in cache:
-                cache[self.log_type] = {}
-
-            cache[self.log_type]["most_recent_date_seen"] = dt.isoformat()
-
-        self.connector.context_lock.release()
-
     @staticmethod
     def __format_datetime(dt: datetime) -> str:
         base = dt.strftime("%Y-%m-%dT%H:%M:%S")
@@ -185,7 +158,7 @@ def fetch_events(self) -> Generator[list, None, None]:
                 self.from_date = most_recent_date_seen
 
                 # save in context the most recent date seen
-                self.most_recent_date_seen = most_recent_date_seen
+                self.cursor.offset = most_recent_date_seen
 
                 # Update the current lag only if the most_recent_date_seen was updated
                 delta_time = datetime.now(timezone.utc) - most_recent_date_seen

Mimecast/tests/test_mimecast_siem_logs.py

@@ -16,7 +16,7 @@ def fake_time():
 
 @pytest.fixture
 def patch_datetime_now(fake_time):
-    with patch("mimecast_modules.connector_mimecast_siem.datetime") as mock_datetime:
+    with patch("sekoia_automation.checkpoint.datetime") as mock_datetime:
         mock_datetime.now.return_value = fake_time
         mock_datetime.side_effect = lambda *args, **kw: datetime(*args, **kw)
         mock_datetime.fromtimestamp = lambda ts: datetime.fromtimestamp(ts)
@@ -114,11 +114,6 @@ def test_fetch_batches(
         mock_time.sleep.assert_called_once_with(44)
 
 
-def test_most_recent_datetime_seen(trigger, patch_datetime_now, fake_time):
-    consumer = MimecastSIEMWorker(connector=trigger, log_type="process")
-    assert consumer.most_recent_date_seen == fake_time - timedelta(days=1)
-
-
 def test_start_consumers(trigger):
     with patch("mimecast_modules.connector_mimecast_siem.MimecastSIEMWorker.start") as mock_start:
         consumers = trigger.start_consumers()