diff --git a/CHANGES b/CHANGES index 69b18ee..6ab7f01 100644 --- a/CHANGES +++ b/CHANGES @@ -1,3 +1,7 @@ +1.22.2 +------- +- enforce providing an environemnt for alerts and incidents + 1.22.1 ------- - Add Get raw data for alerts and incidents diff --git a/intezer_sdk/__init__.py b/intezer_sdk/__init__.py index 800d519..53bfe2e 100644 --- a/intezer_sdk/__init__.py +++ b/intezer_sdk/__init__.py @@ -1 +1 @@ -__version__ = '1.22.1' +__version__ = '1.22.2' diff --git a/intezer_sdk/_api.py b/intezer_sdk/_api.py index 9c9d2f1..4fe0342 100644 --- a/intezer_sdk/_api.py +++ b/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) diff --git a/intezer_sdk/alerts.py b/intezer_sdk/alerts.py index d599fb8..4b3948a 100644 --- a/intezer_sdk/alerts.py +++ b/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 or None, 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 ) diff --git a/intezer_sdk/incidents.py b/intezer_sdk/incidents.py index 6003b53..94f1c59 100644 --- a/intezer_sdk/incidents.py +++ b/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,12 +198,17 @@ 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) raise - + + if not self.environment: + self.environment = self._result['environment'] + self.source = self._result.get('source') self.sender = self._result.get('sender') self.name = self._result.get('incident', {}).get('name') @@ -215,34 +226,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 or 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, ) diff --git a/tests/unit/test_alerts.py b/tests/unit/test_alerts.py index 733f5b4..64b278f 100644 --- a/tests/unit/test_alerts.py +++ b/tests/unit/test_alerts.py @@ -95,7 +95,7 @@ def test_alert_from_id(self): status=HTTPStatus.OK, json={'result': {}, 'status': 'success'}) # Act - alert = Alert.from_id('alert_id') + alert = Alert.from_id('alert_id', environment='environment') # Assert self.assertEqual(alert.alert_id, 'alert_id') @@ -111,7 +111,7 @@ def test_alert_from_id_waits_from_completion(self): status=HTTPStatus.OK, json={'result': {}, 'status': 'success'}) # Act - alert = Alert.from_id('alert_id', wait=True) + alert = Alert.from_id('alert_id', environment='environment', wait=True) # Assert self.assertEqual(alert.alert_id, 'alert_id') diff --git a/tests/unit/test_incidents.py b/tests/unit/test_incidents.py index 27ffff5..dae8ba4 100644 --- a/tests/unit/test_incidents.py +++ b/tests/unit/test_incidents.py @@ -142,7 +142,7 @@ def test_incident_fetch_info_success(self): } }) - incident = Incident(incident_id) + incident = Incident(incident_id, environment='test_environment') # Act incident.fetch_info() @@ -195,7 +195,7 @@ def test_incident_result_returns_raw_data(self): status=HTTPStatus.OK, json={'result': expected_result}) - incident = Incident(incident_id) + incident = Incident(incident_id, environment='test_environment') incident.fetch_info() # Act @@ -206,7 +206,7 @@ def test_incident_result_returns_raw_data(self): def test_incident_result_returns_none_when_no_data_fetched(self): # Arrange - incident = Incident('test_incident_id') + incident = Incident('test_incident_id', environment='test_environment') # Act result = incident.result() @@ -235,7 +235,7 @@ def test_incident_from_id_success(self): }) # Act - incident = Incident.from_id(incident_id) + incident = Incident.from_id(incident_id, environment='test_environment') # Assert self.assertEqual(incident.incident_id, incident_id) @@ -260,7 +260,7 @@ def test_get_raw_incident_data(self): json=expected_raw_data, status=HTTPStatus.OK) - incident = Incident(incident_id=incident_id) + incident = Incident(incident_id=incident_id, environment=environment) # Act result_data = incident.get_raw_data(environment=environment)