Merge pull request #1066 from SEKOIA-IO/fix/trellix_edr_alerts_fixes

Fix: Trellix EDR Alerts response changes

Trellix/CHANGELOG.md

@@ -7,6 +7,12 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+## 2024-08-06 - 1.9.0
+
+### Fixed
+
+- add support of changed structure for Trellix alerts response
+
 ## 2024-07-31 - 1.8.6
 
 ### Fixed

Trellix/client/errors.py

@@ -4,8 +4,6 @@
 class Error(Exception):
     """Base class for exceptions in this module"""
 
-    pass
-
 
 class APIError(Error):
     pass

Trellix/client/http_client.py

@@ -4,10 +4,11 @@
 from datetime import datetime
 from typing import Any, AsyncGenerator, Optional, Set
 
+from aiohttp_retry import ExponentialRetry, RetryClient
 from aiolimiter import AsyncLimiter
-from aiohttp_retry import RetryClient, ExponentialRetry, RetryOptionsBase
 from yarl import URL
 
+from .retry import RetryWithRateLimiter
 from .schemas.attributes.edr_affectedhosts import EdrAffectedhostAttributes
 from .schemas.attributes.edr_alerts import EdrAlertAttributes
 from .schemas.attributes.edr_detections import EdrDetectionAttributes
@@ -16,7 +17,6 @@
 from .schemas.token import Scope
 from .schemas.trellix_response import TrellixResponse
 from .token_refresher import TrellixTokenRefresher
-from .retry import RetryWithRateLimiter
 
 
 class TrellixHttpClient(object):
@@ -453,4 +453,9 @@ async def get_edr_alerts(self, start_date: datetime, limit: int = 10) -> list[Tr
 
         data = await self._get_data(url, headers)
 
-        return [TrellixResponse[EdrAlertAttributes](**result) for result in data["data"]]
+        return [
+            TrellixResponse[EdrAlertAttributes](
+                **{**result, "attributes": EdrAlertAttributes.parse_response(result.get("attributes"))}
+            )
+            for result in data["data"]
+        ]

Trellix/client/retry.py

@@ -1,7 +1,7 @@
 from typing import Any
 
-from aiohttp_retry import RetryOptionsBase
 from aiohttp import ClientResponse
+from aiohttp_retry import RetryOptionsBase
 
 
 class RetryWithRateLimiter(RetryOptionsBase):

Trellix/client/schemas/attributes/edr_alerts.py

@@ -2,8 +2,13 @@
 Models for EDR Alerts attributes.
 
 https://developer.manage.trellix.com/mvision/apis/threats
+
+Since some version, alert attributes contains different format.
+For backward compatibility, we need to handle both formats.
 """
 
+from typing import Any
+
 from pydantic import BaseModel
 
 
@@ -33,3 +38,46 @@ class EdrAlertAttributes(BaseModel):
     h_os: str | None = None
     domain: str | None = None
     hostName: str | None = None
+
+    @staticmethod
+    def parse_response(json: dict[str, Any]) -> "EdrAlertAttributes":
+        """
+        Parse the response.
+
+        Args:
+            json: dict[str, str]
+
+        Returns:
+            EdrAlertAttributes
+        """
+        return EdrAlertAttributes(
+            traceId=json.get("Trace_Id"),
+            parentTraceId=json.get("Parent_Trace_Id"),
+            rootTraceId=json.get("Root_Trace_Id"),
+            aGuid=json.get("MAGUID"),
+            detectionDate=json.get("DetectionDate"),
+            eventDate=json.get("Event_Date"),
+            # api docs does not have this field, and seems like we do not parse it in format
+            # but for backward compatibility `eventType` by default is set to `alert`
+            eventType=json.get("Event_Type") or "alert",
+            severity=json.get("Severity"),
+            score=json.get("Score"),
+            detectionTags=json.get("Detection_Tags"),
+            relatedTraceIds=json.get("Related_Trace_Id"),
+            ruleId=json.get("RuleId"),
+            rank=json.get("Rank"),
+            pid=json.get("Pid"),
+            version=json.get("Version"),
+            parentsTraceId=json.get("Parents_Trace_Id"),
+            processName=json.get("ProcessName"),
+            cmdLine=json.get("CommandLine"),
+            hashId=json.get("Hash_Id"),
+            h_os=json.get("Host_OS"),
+            hostName=json.get("Host_Name"),
+            # Api docs does not have these fields in separate keys, however `User` in new version contains:
+            #  - `domain`
+            #  - `Name`
+            # So we split it into two fields: `user` and `domain`
+            user=json.get("User", {}).get("name"),
+            domain=json.get("User", {}).get("domain"),
+        )

Trellix/client/token_refresher.py

@@ -4,16 +4,16 @@
 import time
 from asyncio import Lock, Task
 from contextlib import asynccontextmanager
+from posixpath import join as urljoin
 from typing import AsyncGenerator, Optional, Set
 from urllib.parse import urlencode
-from posixpath import join as urljoin
 
 from aiohttp import BasicAuth, ClientSession
 from loguru import logger
 from yarl import URL
 
-from .schemas.token import HttpToken, Scope, TrellixToken
 from .errors import APIError, AuthenticationFailed
+from .schemas.token import HttpToken, Scope, TrellixToken
 
 
 class TrellixTokenRefresher(object):

Trellix/connectors/metrics.py

@@ -1,4 +1,4 @@
-from prometheus_client import Counter, Histogram, Gauge
+from prometheus_client import Counter, Gauge, Histogram
 
 # Declare prometheus metrics
 prom_namespace = "symphony_module_trellix_edr"

Trellix/manifest.json

@@ -46,5 +46,5 @@
   "name": "Trellix",
   "uuid": "888071f8-1456-11ee-be56-0242ac120002",
   "slug": "trellix",
-  "version": "1.8.6"
+  "version": "1.9.0"
 }