Merge pull request #1191 from SEKOIA-IO/lv/improve_mimecast_connector

Mimecast - add new event types

Mimecast/CHANGELOG.md

@@ -7,6 +7,16 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## Unreleased
 
+## 2024-11-26 - 1.1.0
+
+### Added
+
+- Added `av`, `delivery`, `internal email protect`, `impersonation protect`, `attachment protect`, `spam`, `url protect` logs support
+
+### Changed
+
+- Changed the way rate limiting works - now it is shared across the threads
+
 ## 2024-11-14 - 1.0.2
 
 ### Fixed

Mimecast/manifest.json

@@ -25,7 +25,7 @@
   "name": "Mimecast",
   "slug": "mimecast",
   "uuid": "72af1e06-84db-497d-b4ac-10defb1f265f",
-  "version": "1.0.2",
+  "version": "1.1.0",
   "categories": [
     "Email"
   ]

Mimecast/mimecast_modules/client/__init__.py

@@ -1,5 +1,6 @@
 import requests
 from pyrate_limiter import Duration, Limiter, RequestRate
+from requests.auth import AuthBase
 from requests_ratelimiter import LimiterAdapter
 
 from .auth import ApiKeyAuthentication
@@ -9,37 +10,20 @@
 class ApiClient(requests.Session):
     def __init__(
         self,
-        client_id: str,
-        client_secret: str,
+        auth: AuthBase,
+        limiter_batch: Limiter,
+        limiter_default: Limiter,
         nb_retries: int = 5,
-        ratelimit_per_minute: int = 20,
     ):
         super().__init__()
-        self.auth = ApiKeyAuthentication(
-            client_id=client_id,
-            client_secret=client_secret,
-            ratelimit_per_second=ratelimit_per_minute,
-            nb_retries=nb_retries,
-        )
+        self.auth = auth
+        self.limiter_batch = limiter_batch
+        self.limiter_default = limiter_default
 
-        # 50 times within a 15 minutes fixed window
-        fifteen_minutes_rate = RequestRate(limit=50, interval=Duration.MINUTE * 15)
         self.mount(
             "https://api.services.mimecast.com/siem/v1/batch/events/cg",
             LimiterAdapter(
-                limiter=Limiter(fifteen_minutes_rate),
-                max_retries=Retry(
-                    total=nb_retries,
-                    backoff_factor=1,
-                ),
-            ),
-        )
-
-        # 300 api calls/hour
-        self.mount(
-            "https://api.services.mimecast.com/siem/v1/events/cg",
-            LimiterAdapter(
-                per_hour=300,
+                limiter=self.limiter_batch,
                 max_retries=Retry(
                     total=nb_retries,
                     backoff_factor=1,
@@ -50,7 +34,7 @@ def __init__(
         self.mount(
             "https://",
             LimiterAdapter(
-                per_minute=ratelimit_per_minute,
+                limiter=self.limiter_default,
                 max_retries=Retry(
                     total=nb_retries,
                     backoff_factor=1,

Mimecast/mimecast_modules/client/auth.py

@@ -1,6 +1,8 @@
 from datetime import datetime, timedelta
+from threading import Lock
 
 import requests
+from pyrate_limiter import Limiter
 from requests.auth import AuthBase
 from requests_ratelimiter import LimiterAdapter
 from urllib3 import Retry
@@ -19,9 +21,7 @@ def authorization(self) -> str:
 class ApiKeyAuthentication(AuthBase):
     AUTH_URL = "https://api.services.mimecast.com/oauth/token"
 
-    def __init__(
-        self, client_id: str, client_secret: str, ratelimit_per_second: int = 10, nb_retries: int = 5
-    ) -> None:
+    def __init__(self, client_id: str, client_secret: str, limiter: Limiter, nb_retries: int = 5) -> None:
         self.__client_id = client_id
         self.__client_secret = client_secret
 
@@ -30,15 +30,19 @@ def __init__(
         self.__http_session.mount(
             "https://",
             LimiterAdapter(
-                per_second=ratelimit_per_second,
+                limiter=limiter,
                 max_retries=Retry(
                     total=nb_retries,
                     backoff_factor=1,
                 ),
             ),
         )
 
+        self.__lock = Lock()
+
     def get_credentials(self) -> MimecastCredentials:
+        self.__lock.acquire()
+
         current_dt = datetime.utcnow()
         if self.__credentials is None or current_dt + timedelta(seconds=300) >= self.__credentials.expires_at:
             response = self.__http_session.post(
@@ -57,8 +61,11 @@ def get_credentials(self) -> MimecastCredentials:
             credentials.expires_at = current_dt + timedelta(seconds=api_credentials["expires_in"])
             self.__credentials = credentials
 
+        self.__lock.release()
+
         return self.__credentials
 
     def __call__(self, request):
         request.headers["Authorization"] = self.get_credentials().authorization
+
         return request

Mimecast/mimecast_modules/connector_mimecast_siem.py

@@ -8,14 +8,15 @@
 import orjson
 import requests
 from dateutil.parser import isoparse
+from pyrate_limiter import Duration, Limiter, RequestRate
 from sekoia_automation.connector import Connector, DefaultConnectorConfiguration
 from sekoia_automation.storage import PersistentJSON
 
 from . import MimecastModule
-from .client import ApiClient
+from .client import ApiClient, ApiKeyAuthentication
 from .helpers import download_batches, get_upper_second
-from .metrics import EVENTS_LAG, FORWARD_EVENTS_DURATION, INCOMING_MESSAGES, OUTCOMING_EVENTS
 from .logging import get_logger
+from .metrics import EVENTS_LAG, FORWARD_EVENTS_DURATION, INCOMING_MESSAGES, OUTCOMING_EVENTS
 
 logger = get_logger()
 
@@ -53,8 +54,9 @@ def running(self):
     @cached_property
     def client(self) -> ApiClient:
         return ApiClient(
-            client_id=self.connector.module.configuration.client_id,
-            client_secret=self.connector.module.configuration.client_secret,
+            auth=self.connector.api_auth,
+            limiter_batch=self.connector.limiter_batch,
+            limiter_default=self.connector.limiter_default,
         )
 
     @property
@@ -244,16 +246,38 @@ class MimecastSIEMConnector(Connector):
     module: MimecastModule
     configuration: MimecastSIEMConfiguration
 
-    TYPES_TO_GET = ("process", "journal", "receipt")
+    TYPES_TO_GET = (
+        "attachment protect",
+        "av",
+        "delivery",
+        "impersonation protect",
+        "internal email protect",
+        "journal",
+        "process",
+        "receipt",
+        "spam",
+        "url protect",
+    )
 
     def __init__(self, *args, **kwargs) -> None:
         super().__init__(*args, **kwargs)
 
-        self._stop_event = Event()
-
         self.context = PersistentJSON("context.json", self._data_path)
         self.context_lock = Lock()
 
+        # 50 times within a 15 minutes fixed window
+        batch_rate = RequestRate(limit=50, interval=Duration.MINUTE * 15)
+        self.limiter_batch = Limiter(batch_rate)
+
+        default_rate = RequestRate(limit=20, interval=Duration.MINUTE)
+        self.limiter_default = Limiter(default_rate)
+
+        self.api_auth = ApiKeyAuthentication(
+            client_id=self.module.configuration.client_id,
+            client_secret=self.module.configuration.client_secret,
+            limiter=self.limiter_default,
+        )
+
     def start_consumers(self) -> dict[str, MimecastSIEMWorker]:
         consumers = {}
         for consumer_name in self.TYPES_TO_GET: