refactor(TKT-9082): add environment attribute to Alert and Incident classes

Author: Yonatan-Nethanel

intezer_sdk/_api.py

@@ -726,8 +726,10 @@ def get_alert_by_alert_id(self, alert_id: str, environment: Optional[str] = None
 
         return data_response['result'], data_response['status']
 
-    def get_incident_by_id(self, incident_id: str) -> dict:
+    def get_incident_by_id(self, incident_id: str, environment: Optional[str] = None) -> dict:
         data = dict(incident_id=incident_id)
+        if environment:
+            data['environment'] = environment
         response = self.api.request_with_refresh_expired_access_token(method='GET',
                                                                       path='/incidents/get-by-id',
                                                                       data=data)

intezer_sdk/alerts.py

@@ -270,6 +270,7 @@ class Alert:
 
     def __init__(self,
                  alert_id: Optional[str] = None,
+                 environment: Optional[str] = None,
                  alert_stream: Optional[BinaryIO] = None,
                  api: IntezerApiClient = None):
         """
@@ -278,6 +279,7 @@ def __init__(self,
         instance with the given alert id.
 
         :param alert_id: The alert id.
+        :param environment: The environment of the alert.
         :param api: The API connection to Intezer.
         """
         if alert_stream and alert_id:
@@ -294,6 +296,7 @@ def __init__(self,
         else:
             self.alert_id: str = alert_id
 
+        self.environment = environment
         self._intezer_api_client = api
         self._api = IntezerApi(api or get_global_api())
         self._report: Optional[Dict] = None
@@ -320,12 +323,14 @@ def check_status(self) -> AlertStatusCode:
 
         """
         try:
-            alert, status = self._api.get_alert_by_alert_id(alert_id=self.alert_id)
+            alert, status = self._api.get_alert_by_alert_id(alert_id=self.alert_id, environment=self.environment)
         except requests.HTTPError:
             self.status = AlertStatusCode.NOT_FOUND
             raise errors.AlertNotFoundError(self.alert_id)
 
         self._report = alert
+        if not self.environment:
+            self.environment = alert['environment']
 
         if status in (AlertStatusCode.IN_PROGRESS.value, AlertStatusCode.QUEUED.value):
             self.status = AlertStatusCode.IN_PROGRESS
@@ -359,6 +364,7 @@ def result(self) -> dict:
     @classmethod
     def from_id(cls,
                 alert_id: str,
+                environment: Optional[str] = None,
                 api: IntezerApiClient = None,
                 fetch_scans: bool = False,
                 wait: bool = False,
@@ -367,6 +373,7 @@ def from_id(cls,
         Create a new Alert instance, and fetch the alert data from the Intezer Analyze API.
 
         :param alert_id: The alert id.
+        :param environment: The environment of the alert.
         :param api: The API connection to Intezer.
         :param fetch_scans: Whether to fetch the scans for the alert - this could take some time.
         :param wait: Wait for the alert to finish processing before returning.
@@ -375,7 +382,7 @@ def from_id(cls,
         :raises intezer_sdk.errors.AlertInProgressError: If the alert is still being processed.
         :return: The Alert instance, with the updated alert data.
         """
-        new_alert = cls(alert_id=alert_id, api=api)
+        new_alert = cls(alert_id=alert_id, environment=environment, api=api)
         status = new_alert.check_status()
         if status == AlertStatusCode.IN_PROGRESS and not wait:
             raise errors.AlertInProgressError(alert_id)
@@ -488,7 +495,7 @@ def send_phishing_email(cls,
         send_alert_params = {key: value for key, value in send_alert_params.items() if value is not None}
         alert_id = _api.send_binary_alert(**send_alert_params)
 
-        alert = cls(alert_id=alert_id, api=api)
+        alert = cls(alert_id=alert_id, environment=environment, api=api)
         if wait:
             alert.wait_for_completion(timeout=timeout)
         return alert
@@ -550,17 +557,20 @@ def _fetch_scan(scan_: dict,
                 _fetch_scan(scan, 'url_analysis', UrlAnalysis)
 
     def get_raw_data(self, 
-                     environment: str, 
+                     environment: Optional[str] = None, 
                      raw_data_type: str = 'raw_alert') -> dict:
         """
         Get raw alert data.
 
-        :param environment: The environment to get raw data from.
+        :param environment: The environment to get raw data from. If not provided, the environment will be taken from the alert.
         :param raw_data_type: The type of raw data to retrieve. Defaults to 'raw_alert'.
         :return: The raw alert data.
         """
+        if not environment and not self.environment:
+            raise ValueError('Environment is required to get raw data.')
+        
         return self._api.get_raw_alert_data(
             alert_id=self.alert_id,
-            environment=environment,
+            environment=environment or self.environment,
             raw_data_type=raw_data_type
         )

intezer_sdk/incidents.py

@@ -159,7 +159,12 @@ class Incident:
     :vartype intezer_incident_url: str
     """
 
-    def __init__(self, incident_id: Optional[str] = None, api: IntezerApiClient = None):
+    def __init__(
+        self,
+        incident_id: Optional[str] = None,
+        environment: Optional[str] = None,
+        api: IntezerApiClient = None,
+    ):
         """
         Create a new Incident instance with the given incident id.
         Please note that this does not query the Intezer Analyze API for the incident data, but rather creates an Incident
@@ -168,10 +173,11 @@ def __init__(self, incident_id: Optional[str] = None, api: IntezerApiClient = No
         If you wish to fetch the incident data from the Intezer Analyze API, use the `from_id` class method.
 
         :param incident_id: The incident id.
+        :param environment: The environment of the incident.
         :param api: The API connection to Intezer.
         """
         self.incident_id = incident_id
-
+        self.environment = environment
         self._intezer_api_client = api
         self._api = IntezerApi(api or get_global_api())
         self._result: Optional[Dict] = None
@@ -192,7 +198,9 @@ def fetch_info(self):
             raise ValueError("Incident ID is required to fetch incident info.")
 
         try:
-            self._result = self._api.get_incident_by_id(self.incident_id)
+            self._result = self._api.get_incident_by_id(
+                self.incident_id, self.environment
+            )
         except HTTPError as e:
             if e.response.status_code == 404:
                 raise errors.IncidentNotFoundError(self.incident_id)
@@ -215,34 +223,35 @@ def result(self) -> Optional[dict]:
         return self._result
 
     @classmethod
-    def from_id(cls, incident_id: str, api: IntezerApiClient = None) -> 'Incident':
+    def from_id(cls, incident_id: str, environment: Optional[str] = None, api: IntezerApiClient = None) -> 'Incident':
         """
         Create a new Incident instance, and fetch the incident data from the Intezer Analyze API.
 
         :param incident_id: The incident id.
+        :param environment: The environment of the incident.
         :param api: The API connection to Intezer.
         :raises intezer_sdk.errors.IncidentNotFound: If the incident was not found.
         :return: The Incident instance, with the updated incident data.
         """
-        new_incident = cls(incident_id, api=api)
+        new_incident = cls(incident_id, environment=environment, api=api)
         new_incident.fetch_info()
         return new_incident
 
-    def get_raw_data(self, 
-                     environment: str, 
-                     raw_data_type: str = 'raw_incident') -> dict:
+    def get_raw_data(
+        self, environment: Optional[str] = None, raw_data_type: str = 'raw_incident'
+    ) -> dict:
         """
         Get raw incident data.
 
-        :param environment: The environment to get raw data from.
+        :param environment: The environment to get raw data from. If not provided, the environment will be taken from the incident.
         :param raw_data_type: The type of raw data to retrieve. Defaults to 'raw_incident'.
         :return: The raw incident data.
         """
-        if not self.incident_id:
-            raise ValueError('Incident ID is required to get raw data.')
-        
+        if not self.incident_id and not self.environment:
+            raise ValueError('Incident ID and environment are required to get raw data.')
+
         return self._api.get_raw_incident_data(
             incident_id=self.incident_id,
-            environment=environment,
-            raw_data_type=raw_data_type
+            environment=environment or self.environment,
+            raw_data_type=raw_data_type,
         )