From d21db72fef4edbc9b049dd03c3b2d92ff9e841eb Mon Sep 17 00:00:00 2001 From: ablakley-r7 <96182471+ablakley-r7@users.noreply.github.com> Date: Tue, 4 Feb 2025 07:47:30 +0000 Subject: [PATCH] [SOAR-18657] mimecast v2 (#3068) * Update threads, error handling, custom config, rate limiting * Add connection test * Fix lint * Add unit tests * Add unit tests * Fix requiremnts * Fix requiremnts --- .../icon_mimecast_v2/connection/connection.py | 25 ++ .../tasks/monitor_siem_logs/task.py | 137 +++++++---- .../mimecast_v2/icon_mimecast_v2/util/api.py | 35 +-- .../util/{endpoints.py => constants.py} | 0 plugins/mimecast_v2/requirements.txt | 2 + .../expected/monitor_siem_logs.json.exp | 47 ++++ .../responses/authenticate.json.resp | 6 + .../responses/monitor_siem_logs.json.resp | 1 + .../monitor_siem_logs_batch.json.resp | 11 + .../monitor_siem_logs_error.json.resp | 7 + .../monitor_siem_logs_json_error.json.resp | 1 + .../unit_test/test_monitor_siem_logs.py | 225 ++++++++++++++++++ plugins/mimecast_v2/unit_test/util.py | 110 +++++++++ 13 files changed, 541 insertions(+), 66 deletions(-) rename plugins/mimecast_v2/icon_mimecast_v2/util/{endpoints.py => constants.py} (100%) create mode 100644 plugins/mimecast_v2/unit_test/expected/monitor_siem_logs.json.exp create mode 100644 plugins/mimecast_v2/unit_test/responses/authenticate.json.resp create mode 100644 plugins/mimecast_v2/unit_test/responses/monitor_siem_logs.json.resp create mode 100644 plugins/mimecast_v2/unit_test/responses/monitor_siem_logs_batch.json.resp create mode 100644 plugins/mimecast_v2/unit_test/responses/monitor_siem_logs_error.json.resp create mode 100644 plugins/mimecast_v2/unit_test/responses/monitor_siem_logs_json_error.json.resp create mode 100644 plugins/mimecast_v2/unit_test/test_monitor_siem_logs.py create mode 100644 plugins/mimecast_v2/unit_test/util.py diff --git a/plugins/mimecast_v2/icon_mimecast_v2/connection/connection.py b/plugins/mimecast_v2/icon_mimecast_v2/connection/connection.py index 6e7a1f33aa..06fbc79bd3 100644 --- a/plugins/mimecast_v2/icon_mimecast_v2/connection/connection.py +++ b/plugins/mimecast_v2/icon_mimecast_v2/connection/connection.py @@ -4,6 +4,7 @@ from icon_mimecast_v2.util.api import API # Custom imports below +from datetime import datetime, timezone class Connection(insightconnect_plugin_runtime.Connection): @@ -23,3 +24,27 @@ def test(self): return {"success": True} except PluginException as error: raise ConnectionTestException(cause=error.cause, assistance=error.assistance, data=error.data) + + def test_task(self): + try: + now_date = datetime.now(tz=timezone.utc).date() + self.api.get_siem_logs(log_type="receipt", query_date=now_date, page_size=1, max_threads=1, next_page=None) + self.logger.info("The connection test to Mimecast was successful.") + return {"success": True} + except PluginException as error: + return_message = "" + failed_message = "The connection test to Mimecast for has failed." + self.logger.info(failed_message) + return_message += f"{failed_message}\n" + + cause_message = f"This failure was caused by: '{error.cause}'" + self.logger.info(cause_message) + return_message += f"{cause_message}\n" + + self.logger.info(error.assistance) + return_message += f"{error.assistance}\n" + raise ConnectionTestException( + cause="Configured credentials do not have permission for this API endpoint.", + assistance="Please ensure credentials have required permissions.", + data=return_message, + ) diff --git a/plugins/mimecast_v2/icon_mimecast_v2/tasks/monitor_siem_logs/task.py b/plugins/mimecast_v2/icon_mimecast_v2/tasks/monitor_siem_logs/task.py index e292c73e72..ea628a9876 100644 --- a/plugins/mimecast_v2/icon_mimecast_v2/tasks/monitor_siem_logs/task.py +++ b/plugins/mimecast_v2/icon_mimecast_v2/tasks/monitor_siem_logs/task.py @@ -1,16 +1,32 @@ import insightconnect_plugin_runtime from insightconnect_plugin_runtime.exceptions import APIException, PluginException +from insightconnect_plugin_runtime.helper import compare_and_dedupe_hashes, hash_sha1 from .schema import MonitorSiemLogsInput, MonitorSiemLogsOutput, MonitorSiemLogsState, Input, Output, Component, State from typing import Dict, List, Tuple from datetime import datetime, timezone, timedelta import copy +# Date format for conversion +DATE_FORMAT = "%Y-%m-%d" +# Default and max values LOG_TYPES = ["receipt", "url protect", "attachment protect"] +DEFAULT_THREAD_COUNT = 10 +DEFAULT_PAGE_SIZE = 100 MAX_LOOKBACK_DAYS = 7 INITIAL_MAX_LOOKBACK_DAYS = 1 +# Run type INITIAL_RUN = "initial_run" SUBSEQUENT_RUN = "subsequent_run" PAGINATION_RUN = "pagination_run" +# Access keys for state and custom config +LOG_HASHES = "log_hashes" +QUERY_CONFIG = "query_config" +QUERY_DATE = "query_date" +CAUGHT_UP = "caught_up" +NEXT_PAGE = "next_page" +# Access keys for custom config +THREAD_COUNT = "thread_count" +PAGE_SIZE = "page_size" class MonitorSiemLogs(insightconnect_plugin_runtime.Task): @@ -24,21 +40,21 @@ def __init__(self): ) def run(self, params={}, state={}, custom_config={}): # pylint: disable=unused-argument - self.logger.info(f"TASK: Received State: {state}") + self.logger.info(f"TASK: Received State: {state.get(QUERY_CONFIG)}") existing_state = state.copy() try: - # TODO: Additional error handling - run_condition = self.detect_run_condition(state.get("query_config", {})) - self.logger.info(f"TASK: Current run state is {run_condition}") - state = self.update_state(state, custom_config) - self.logger.info(f"NEW STATE: {state}") now_date = datetime.now(tz=timezone.utc).date() + run_condition = self.detect_run_condition(state.get(QUERY_CONFIG, {}), now_date) + self.logger.info(f"TASK: Run state is {run_condition}") + state = self.update_state(state) + page_size, thead_count = self.apply_custom_config(state, custom_config) max_run_lookback_date = self.get_max_lookback_date(now_date, run_condition, bool(custom_config)) - query_config = self.prepare_query_params(state.get("query_config", {}), max_run_lookback_date, now_date) - logs, query_config = self.get_all_logs(run_condition, query_config) - # TODO: Dedupe + query_config = self.prepare_query_params(state.get(QUERY_CONFIG, {}), max_run_lookback_date, now_date) + logs, query_config = self.get_all_logs(run_condition, query_config, page_size, thead_count) self.logger.info(f"TASK: Total logs collected this run {len(logs)}") - exit_state, has_more_pages = self.prepare_exit_state(state, query_config, now_date) + logs, log_hashes = compare_and_dedupe_hashes(state.get(LOG_HASHES, []), logs) + self.logger.info(f"TASK: Total logs after deduplication {len(logs)}") + exit_state, has_more_pages = self.prepare_exit_state(state, query_config, now_date, log_hashes) return logs, exit_state, has_more_pages, 200, None except APIException as error: self.logger.info( @@ -47,39 +63,40 @@ def run(self, params={}, state={}, custom_config={}): # pylint: disable=unused- return [], existing_state, False, error.status_code, error except PluginException as error: self.logger.info(f"Error: A Plugin exception has occurred. Cause: {error.cause} Error data: {error.data}.") - return [], existing_state, False, error.status_code, error + return [], existing_state, False, 500, error except Exception as error: self.logger.info(f"Error: Unknown exception has occurred. No results returned. Error Data: {error}") return [], existing_state, False, 500, PluginException(preset=PluginException.Preset.UNKNOWN, data=error) - def detect_run_condition(self, query_config: Dict) -> str: + def detect_run_condition(self, query_config: Dict, now_date: datetime) -> str: """ Return runtype based on query configuration :param query_config: + :param now_date: :return: runtype string """ if not query_config: return INITIAL_RUN for log_type_config in query_config.values(): - if not log_type_config.get("caught_up"): + if not log_type_config.get(CAUGHT_UP) or log_type_config.get(QUERY_DATE) not in str(now_date): return PAGINATION_RUN return SUBSEQUENT_RUN - def update_state(self, state: Dict, custom_config: Dict) -> Dict: + def update_state(self, state: Dict) -> Dict: """ Initialise state, validate state, apply custom config :param state: - :param custom_config: - :return: + :return: State """ - initial_log_type_config = {"caught_up": False} + initial_log_type_config = {CAUGHT_UP: False} if not state: - state = {"query_config": {log_type: copy.deepcopy(initial_log_type_config) for log_type in LOG_TYPES}} - self.apply_custom_config(state, custom_config) + self.logger.info("TASK: Initializing first state...") + state = {QUERY_CONFIG: {log_type: copy.deepcopy(initial_log_type_config) for log_type in LOG_TYPES}} else: for log_type in LOG_TYPES: - if log_type not in state.get("query_config", {}).keys(): - state["query_config"][log_type] = copy.deepcopy(initial_log_type_config) + if log_type not in state.get(QUERY_CONFIG, {}).keys(): + self.logger.info(f"TASK: {log_type} missing from state. Initializing...") + state[QUERY_CONFIG][log_type] = copy.deepcopy(initial_log_type_config) return state def get_max_lookback_date(self, now_date: datetime, run_condition: str, custom_config: bool) -> datetime: @@ -97,18 +114,23 @@ def get_max_lookback_date(self, now_date: datetime, run_condition: str, custom_c max_run_lookback_date = now_date - timedelta(days=max_run_lookback_days) return max_run_lookback_date - def apply_custom_config(self, state: Dict, custom_config: Dict) -> None: + def apply_custom_config(self, state: Dict, custom_config: Dict = {}) -> Tuple[int, int]: """ Apply custom configuration for lookback, query date applies to start and end time of query :param current_query_config: :param custom_config: - :return: N/A + :return: """ - # TODO: Additional custom config for page size, thread size, limit - current_query_config = state.get("query_config") - for log_type, lookback_date_string in custom_config.items(): - self.logger.info(f"TASK: Supplied lookback date of {lookback_date_string} for {log_type} log type") - current_query_config[log_type] = {"query_date": lookback_date_string} + if custom_config: + self.logger.info("TASK: Custom config detected") + if not state: + current_query_config = state.get(QUERY_CONFIG) + for log_type, query_date_string in custom_config.items(): + self.logger.info(f"TASK: Supplied lookback date of {query_date_string} for log type {log_type}") + current_query_config[log_type] = {QUERY_DATE: query_date_string} + page_size = max(1, min(custom_config.get(PAGE_SIZE, DEFAULT_PAGE_SIZE), DEFAULT_PAGE_SIZE)) + thread_count = max(1, custom_config.get(THREAD_COUNT, DEFAULT_THREAD_COUNT)) + return page_size, thread_count def prepare_query_params(self, query_config: Dict, max_lookback_date: Dict, now_date: datetime) -> Dict: """ @@ -119,18 +141,19 @@ def prepare_query_params(self, query_config: Dict, max_lookback_date: Dict, now_ :return: """ for log_type, log_type_config in query_config.items(): - query_date_str = log_type_config.get("query_date") - self.logger.info(f"PREPPING {log_type_config}") - self.logger.info(f"{log_type}, {query_date_str}") + query_date_str = log_type_config.get(QUERY_DATE) if query_date_str: - query_date = datetime.strptime(query_date_str, "%Y-%m-%d").date() + query_date = datetime.strptime(query_date_str, DATE_FORMAT).date() if not query_date_str: - log_type_config["query_date"] = max_lookback_date - elif query_date < now_date and log_type_config.get("caught_up") is True: + self.logger.info( + f"TASK: Query date for {log_type} log type is not present. Initializing a {max_lookback_date}" + ) + log_type_config[QUERY_DATE] = max_lookback_date + elif query_date < now_date and log_type_config.get(CAUGHT_UP) is True: self.logger.info(f"TASK: Log type {log_type} has caught up for {query_date}") - log_type_config["query_date"] = query_date + timedelta(days=1) - log_type_config["caught_up"] = False - log_type_config.pop("next_page") + log_type_config[QUERY_DATE] = query_date + timedelta(days=1) + log_type_config[CAUGHT_UP] = False + log_type_config.pop(NEXT_PAGE) query_config[log_type] = self.validate_config_lookback(log_type_config, max_lookback_date, now_date) return query_config @@ -142,51 +165,61 @@ def validate_config_lookback(self, log_type_config: Dict, max_lookback_date: dat :param now_date: :return: log_type_config """ - query_date = log_type_config.get("query_date") + query_date = log_type_config.get(QUERY_DATE) if isinstance(query_date, str): - query_date = datetime.strptime(query_date, "%Y-%m-%d").date() + query_date = datetime.strptime(query_date, DATE_FORMAT).date() if query_date < max_lookback_date: - return {"query_date": max_lookback_date} + return {QUERY_DATE: max_lookback_date} if query_date > now_date: - log_type_config["query_date"] = now_date + log_type_config[QUERY_DATE] = now_date return log_type_config - def get_all_logs(self, run_condition: str, query_config: Dict) -> Tuple[List, Dict]: + def get_all_logs( + self, run_condition: str, query_config: Dict, page_size: int, thead_count: int + ) -> Tuple[List, Dict]: """ Gets all logs of provided log type. First retrieves batch URLs. Then downloads and reads batches, pooling logs. :param run_condition: :param query_config: + :param page_size: + :param thead_count: :return: Logs, updated query configuration (state) """ complete_logs = [] for log_type, log_type_config in query_config.items(): - if (not log_type_config.get("caught_up")) or (run_condition != PAGINATION_RUN): + if (not log_type_config.get(CAUGHT_UP)) or (run_condition != PAGINATION_RUN): logs, results_next_page, caught_up = self.connection.api.get_siem_logs( log_type=log_type, - query_date=log_type_config.get("query_date"), - next_page=log_type_config.get("next_page"), + query_date=log_type_config.get(QUERY_DATE), + next_page=log_type_config.get(NEXT_PAGE), + page_size=page_size, + max_threads=thead_count, ) complete_logs.extend(logs) - log_type_config.update({"next_page": results_next_page, "caught_up": caught_up}) + log_type_config.update({NEXT_PAGE: results_next_page, CAUGHT_UP: caught_up}) else: self.logger.info(f"TASK: Query for {log_type} is caught up. Skipping as we are currently paginating") return complete_logs, query_config - def prepare_exit_state(self, state: dict, query_config: dict, now_date: datetime) -> Tuple[Dict, bool]: + def prepare_exit_state( + self, state: dict, query_config: dict, now_date: datetime, log_hashes: List[str] + ) -> Tuple[Dict, bool]: """ Prepare state and pagination for task completion. Format date time. :param state: :param query_config: :param now_date: + :param log_hashes: :return: state, has_more_pages """ has_more_pages = False for log_type_config in query_config.values(): - query_date = log_type_config.get("query_date") + query_date = log_type_config.get(QUERY_DATE) if isinstance(query_date, str): - query_date = datetime.strptime(query_date, "%Y-%m-%d").date() - if (not log_type_config.get("caught_up")) or query_date < now_date: + query_date = datetime.strptime(query_date, DATE_FORMAT).date() + if (not log_type_config.get(CAUGHT_UP)) or query_date < now_date: has_more_pages = True - log_type_config["query_date"] = query_date.strftime("%Y-%m-%d") - state["query_config"] = query_config + log_type_config[QUERY_DATE] = query_date.strftime(DATE_FORMAT) + state[QUERY_CONFIG] = query_config + state[LOG_HASHES] = log_hashes return state, has_more_pages diff --git a/plugins/mimecast_v2/icon_mimecast_v2/util/api.py b/plugins/mimecast_v2/icon_mimecast_v2/util/api.py index 04a9e81fd3..893ee662be 100644 --- a/plugins/mimecast_v2/icon_mimecast_v2/util/api.py +++ b/plugins/mimecast_v2/icon_mimecast_v2/util/api.py @@ -5,12 +5,13 @@ HTTPStatusCodes, ResponseExceptionData, ) -from insightconnect_plugin_runtime.helper import make_request, extract_json +from insightconnect_plugin_runtime.helper import extract_json, make_request, rate_limiting from logging import Logger from requests import Response, Request from io import BytesIO -from icon_mimecast_v2.util.endpoints import Endpoints +from icon_mimecast_v2.util.constants import Endpoints from typing import Dict, List, Tuple +from multiprocessing.dummy import Pool import gzip import json @@ -39,17 +40,19 @@ def authenticate(self) -> None: self.logger.info("API: Authenticated") def get_siem_logs( - self, log_type: str, query_date: str, next_page: str, page_size: int = 100 + self, log_type: str, query_date: str, next_page: str, page_size: int = 100, max_threads: int = 10 ) -> Tuple[List[str], str, bool]: batch_download_urls, result_next_page, caught_up = self.get_siem_batches( log_type, query_date, next_page, page_size ) logs = [] self.logger.info(f"API: Getting SIEM logs from batches for log type {log_type}...") - for url in batch_download_urls: - batch_logs = self.get_siem_logs_from_batch(url=url) - if isinstance(batch_logs, (List, Dict)): - logs.extend(batch_logs) + self.logger.info(f"API: Applying page size limit of {page_size}") + with Pool(max_threads) as pool: + batch_logs = pool.imap(self.get_siem_logs_from_batch, batch_download_urls) + for result in batch_logs: + if isinstance(result, (List, Dict)): + logs.extend(result) self.logger.info(f"API: Discovered {len(logs)} logs for log type {log_type}") return logs, result_next_page, caught_up @@ -77,7 +80,6 @@ def get_siem_batches( return urls, batch_response.get("@nextPage"), caught_up def get_siem_logs_from_batch(self, url: str): - # TODO: Threading response = requests.request(method=GET, url=url, stream=False) with gzip.GzipFile(fileobj=BytesIO(response.content), mode="rb") as file_: logs = [] @@ -87,6 +89,7 @@ def get_siem_logs_from_batch(self, url: str): logs.append(json.loads(decoded_line)) return logs + @rate_limiting(5) def make_api_request( self, url: str, @@ -101,7 +104,6 @@ def make_api_request( if auth: headers["Authorization"] = f"Bearer {self.access_token}" request = Request(url=url, method=method, headers=headers, params=params, data=data, json=json) - # TODO: Handle rate limit, handle retry backoff try: response = make_request( _request=request, @@ -117,11 +119,16 @@ def make_api_request( status_code=exception.data.status_code, ) raise exception - if ( - response.status_code == HTTPStatusCodes.UNAUTHORIZED - and extract_json(response).get("fail", [{}])[0].get("code") == "token_expired" - ): - self.authenticate() + if response.status_code == HTTPStatusCodes.UNAUTHORIZED: + json_data = extract_json(response) + if json_data.get("fail", [{}])[0].get("code") == "token_expired": + self.authenticate() + self.logger.info("API: Token has expired, attempting re-authentication...") + return self.make_api_request(url, method, headers, json, data, params, return_json, auth) + if response.status_code == HTTPStatusCodes.UNAUTHORIZED: + raise APIException( + preset=PluginException.Preset.API_KEY, data=response.text, status_code=response.status_code + ) if return_json: json_data = extract_json(response) return json_data diff --git a/plugins/mimecast_v2/icon_mimecast_v2/util/endpoints.py b/plugins/mimecast_v2/icon_mimecast_v2/util/constants.py similarity index 100% rename from plugins/mimecast_v2/icon_mimecast_v2/util/endpoints.py rename to plugins/mimecast_v2/icon_mimecast_v2/util/constants.py diff --git a/plugins/mimecast_v2/requirements.txt b/plugins/mimecast_v2/requirements.txt index 33c8572218..de275296d3 100644 --- a/plugins/mimecast_v2/requirements.txt +++ b/plugins/mimecast_v2/requirements.txt @@ -1,3 +1,5 @@ # List third-party dependencies here, separated by newlines. # All dependencies must be version-pinned, eg. requests==1.2.0 # See: https://pip.pypa.io/en/stable/user_guide/#requirements-files +parameterized==0.8.1 +freezegun==1.5.1 diff --git a/plugins/mimecast_v2/unit_test/expected/monitor_siem_logs.json.exp b/plugins/mimecast_v2/unit_test/expected/monitor_siem_logs.json.exp new file mode 100644 index 0000000000..5c70e65a8b --- /dev/null +++ b/plugins/mimecast_v2/unit_test/expected/monitor_siem_logs.json.exp @@ -0,0 +1,47 @@ +[ + { + "_offset": 273017, + "_partition": 125, + "accountId": "PDLW18H03", + "aggregateId": "jdWc03xmmwdc903dnsnALm", + "fileExtension": "pdf", + "fileName": "attach-9fac38ge-8387-57f0-dfd9-d709gb3ce83d.pdf", + "md5": "f34f2eya7e54hhr80d3937hd084898fe", + "processingId": "JDgeKPUdn2pe8Suxow5oplvfpz37i8M1lWNyGD-OUPE_1738206319", + "sha1": "534dbfd55d2f56fd213b3735fa74d2e80ege4127", + "sha256": "d473cgeerf4e5g2s133584188cf263131a37c54e33648c203d31309f73aa3b00", + "subtype": null, + "timestamp": 946684800, + "type": "attachment protect" + }, + { + "_offset": 273017, + "_partition": 125, + "accountId": "PDLW18H03", + "aggregateId": "jdWc03xmmwdc903dnsnALm", + "fileExtension": "pdf", + "fileName": "attach-9fac38ge-8387-57f0-dfd9-d709gb3ce83d.pdf", + "md5": "f34f2eya7e54hhr80d3937hd084898fe", + "processingId": "JDgeKPUdn2pe8Suxow5oplvfpz37i8M1lWNyGD-OUPE_1738206319", + "sha1": "534dbfd55d2f56fd213b3735fa74d2e80ege4127", + "sha256": "d473cgeerf4e5g2s133584188cf263131a37c54e33648c203d31309f73aa3b00", + "subtype": null, + "timestamp": 946684800, + "type": "attachment protect" + }, + { + "_offset": 273017, + "_partition": 125, + "accountId": "PDLW18H03", + "aggregateId": "jdWc03xmmwdc903dnsnALm", + "fileExtension": "pdf", + "fileName": "attach-9fac38ge-8387-57f0-dfd9-d709gb3ce83d.pdf", + "md5": "f34f2eya7e54hhr80d3937hd084898fe", + "processingId": "JDgeKPUdn2pe8Suxow5oplvfpz37i8M1lWNyGD-OUPE_1738206319", + "sha1": "534dbfd55d2f56fd213b3735fa74d2e80ege4127", + "sha256": "d473cgeerf4e5g2s133584188cf263131a37c54e33648c203d31309f73aa3b00", + "subtype": null, + "timestamp": 946684800, + "type": "attachment protect" + } +] \ No newline at end of file diff --git a/plugins/mimecast_v2/unit_test/responses/authenticate.json.resp b/plugins/mimecast_v2/unit_test/responses/authenticate.json.resp new file mode 100644 index 0000000000..b4427808e1 --- /dev/null +++ b/plugins/mimecast_v2/unit_test/responses/authenticate.json.resp @@ -0,0 +1,6 @@ +{ + "access_token": "RPjLow4SKlsPLycjML3dqvIxzi7B", + "token_type": "Bearer", + "expires_in": 1799, + "scope": "" +} \ No newline at end of file diff --git a/plugins/mimecast_v2/unit_test/responses/monitor_siem_logs.json.resp b/plugins/mimecast_v2/unit_test/responses/monitor_siem_logs.json.resp new file mode 100644 index 0000000000..12f2ac2387 --- /dev/null +++ b/plugins/mimecast_v2/unit_test/responses/monitor_siem_logs.json.resp @@ -0,0 +1 @@ +{ "_offset": 273017, "_partition": 125, "accountId": "PDLW18H03", "aggregateId": "jdWc03xmmwdc903dnsnALm", "fileExtension": "pdf", "fileName": "attach-9fac38ge-8387-57f0-dfd9-d709gb3ce83d.pdf", "md5": "f34f2eya7e54hhr80d3937hd084898fe", "processingId": "JDgeKPUdn2pe8Suxow5oplvfpz37i8M1lWNyGD-OUPE_1738206319", "sha1": "534dbfd55d2f56fd213b3735fa74d2e80ege4127", "sha256": "d473cgeerf4e5g2s133584188cf263131a37c54e33648c203d31309f73aa3b00", "subtype": null, "timestamp": 946684800, "type": "attachment protect" } \ No newline at end of file diff --git a/plugins/mimecast_v2/unit_test/responses/monitor_siem_logs_batch.json.resp b/plugins/mimecast_v2/unit_test/responses/monitor_siem_logs_batch.json.resp new file mode 100644 index 0000000000..66ce90286d --- /dev/null +++ b/plugins/mimecast_v2/unit_test/responses/monitor_siem_logs_batch.json.resp @@ -0,0 +1,11 @@ +{ + "value": [ + { + "url": "https://example.com", + "expiry": "2000-01-10T00:10:00.000Z", + "size": 354 + } + ], + "@nextPage": "NDU1NA==", + "isCaughtUp": true +} \ No newline at end of file diff --git a/plugins/mimecast_v2/unit_test/responses/monitor_siem_logs_error.json.resp b/plugins/mimecast_v2/unit_test/responses/monitor_siem_logs_error.json.resp new file mode 100644 index 0000000000..e29af9c064 --- /dev/null +++ b/plugins/mimecast_v2/unit_test/responses/monitor_siem_logs_error.json.resp @@ -0,0 +1,7 @@ +{ + "fail": [ + { + "code": "unauthorized" + } + ] +} \ No newline at end of file diff --git a/plugins/mimecast_v2/unit_test/responses/monitor_siem_logs_json_error.json.resp b/plugins/mimecast_v2/unit_test/responses/monitor_siem_logs_json_error.json.resp new file mode 100644 index 0000000000..f57ca5f651 --- /dev/null +++ b/plugins/mimecast_v2/unit_test/responses/monitor_siem_logs_json_error.json.resp @@ -0,0 +1 @@ +This is not JSON \ No newline at end of file diff --git a/plugins/mimecast_v2/unit_test/test_monitor_siem_logs.py b/plugins/mimecast_v2/unit_test/test_monitor_siem_logs.py new file mode 100644 index 0000000000..c5ab0c5120 --- /dev/null +++ b/plugins/mimecast_v2/unit_test/test_monitor_siem_logs.py @@ -0,0 +1,225 @@ +import sys +import os + +sys.path.append(os.path.abspath("../")) + +from icon_mimecast_v2.tasks.monitor_siem_logs import MonitorSiemLogs +from icon_mimecast_v2.tasks.monitor_siem_logs.schema import MonitorSiemLogsOutput + +from unittest import TestCase +from unittest.mock import patch +from util import Util +from parameterized import parameterized +from jsonschema import validate +from freezegun import freeze_time + + +STUB_STATE_EXPECTED = { + "log_hashes": [ + "d98dafb4f13b3bb70539a6c251a8a9b42ea80de1", + "d98dafb4f13b3bb70539a6c251a8a9b42ea80de1", + "d98dafb4f13b3bb70539a6c251a8a9b42ea80de1", + ], + "query_config": { + "attachment protect": {"caught_up": True, "next_page": "NDU1NA==", "query_date": "2000-01-06"}, + "receipt": {"caught_up": True, "next_page": "NDU1NA==", "query_date": "2000-01-06"}, + "url protect": {"caught_up": True, "next_page": "NDU1NA==", "query_date": "2000-01-06"}, + }, +} + +STUB_STATE_PAGINATING = { + "log_hashes": [], + "query_config": { + "attachment protect": {"caught_up": False, "next_page": "NDU1NA==", "query_date": "2000-01-06"}, + "receipt": {"caught_up": False, "next_page": "NDU1NA==", "query_date": "2000-01-06"}, + "url protect": {"caught_up": False, "next_page": "NDU1NA==", "query_date": "2000-01-06"}, + }, +} + +STUB_STATE_PAGINATING_LAST_PAGE = { + "log_hashes": [], + "query_config": { + "attachment protect": {"caught_up": True, "next_page": "NDU1NA==", "query_date": "2000-01-06"}, + "receipt": {"caught_up": True, "next_page": "NDU1NA==", "query_date": "2000-01-06"}, + "url protect": {"caught_up": True, "next_page": "NDU1NA==", "query_date": "2000-01-06"}, + }, +} + +STUB_STATE_SECOND_RUN_EXPECTED = { + "log_hashes": [ + "d98dafb4f13b3bb70539a6c251a8a9b42ea80de1", + "d98dafb4f13b3bb70539a6c251a8a9b42ea80de1", + "d98dafb4f13b3bb70539a6c251a8a9b42ea80de1", + ], + "query_config": { + "attachment protect": {"caught_up": True, "next_page": "NDU1NA==", "query_date": "2000-01-07"}, + "receipt": {"caught_up": True, "next_page": "NDU1NA==", "query_date": "2000-01-07"}, + "url protect": {"caught_up": True, "next_page": "NDU1NA==", "query_date": "2000-01-07"}, + }, +} + +STUB_STATE_EXPECTED_CUSTOM_CONFIG = { + "log_hashes": [ + "d98dafb4f13b3bb70539a6c251a8a9b42ea80de1", + "d98dafb4f13b3bb70539a6c251a8a9b42ea80de1", + "d98dafb4f13b3bb70539a6c251a8a9b42ea80de1", + ], + "query_config": { + "attachment protect": {"caught_up": True, "next_page": "NDU1NA==", "query_date": "1999-12-31"}, + "receipt": {"caught_up": True, "next_page": "NDU1NA==", "query_date": "1999-12-31"}, + "url protect": {"caught_up": True, "next_page": "NDU1NA==", "query_date": "1999-12-31"}, + }, +} + +STUB_CUSTOM_CONFIG = { + "query_config": { + "attachment protect": {"caught_up": True, "next_page": "NDU1NA==", "query_date": "1999-12-31"}, + "receipt": {"caught_up": True, "next_page": "NDU1NA==", "query_date": "1999-12-31"}, + "url protect": {"caught_up": True, "next_page": "NDU1NA==", "query_date": "1999-12-31"}, + }, + "page_size": 1, + "thread_count": 1, +} + +STUB_CUSTOM_CONFIG_EXCEED_DATE = { + "query_config": { + "attachment protect": {"caught_up": True, "next_page": "NDU1NA==", "query_date": "1999-12-30"}, + "receipt": {"caught_up": True, "next_page": "NDU1NA==", "query_date": "1999-12-30"}, + "url protect": {"caught_up": True, "next_page": "NDU1NA==", "query_date": "1999-12-30"}, + }, + "page_size": 1, + "thread_count": 1, +} + + +@freeze_time("2000-01-07T00:00:00.000000Z") +class TestMonitorLogs(TestCase): + @classmethod + @patch("requests.Session.send", side_effect=Util.mocked_request) + def setUpClass(cls, mocked_request) -> None: + cls.task = Util.default_connector(MonitorSiemLogs()) + + @parameterized.expand( + [ + [ + "starting", + {}, + {}, + Util.read_file_to_dict("expected/monitor_siem_logs.json.exp"), + STUB_STATE_EXPECTED, + True, + 200, + None, + ], + [ + "paginating", + STUB_STATE_PAGINATING, + {}, + Util.read_file_to_dict("expected/monitor_siem_logs.json.exp"), + STUB_STATE_EXPECTED, + True, + 200, + None, + ], + [ + "paginating_last_page", + STUB_STATE_PAGINATING_LAST_PAGE, + {}, + Util.read_file_to_dict("expected/monitor_siem_logs.json.exp"), + STUB_STATE_SECOND_RUN_EXPECTED, + False, + 200, + None, + ], + [ + "custom_config", + {}, + STUB_CUSTOM_CONFIG, + Util.read_file_to_dict("expected/monitor_siem_logs.json.exp"), + STUB_STATE_EXPECTED_CUSTOM_CONFIG, + True, + 200, + None, + ], + [ + "custom_config_past_cutoff", + {}, + STUB_CUSTOM_CONFIG_EXCEED_DATE, + Util.read_file_to_dict("expected/monitor_siem_logs.json.exp"), + STUB_STATE_EXPECTED_CUSTOM_CONFIG, + True, + 200, + None, + ], + ] + ) + @patch("requests.Session.send", side_effect=Util.mocked_request) + def test_monitor_logs( + self, + test_name, + state, + custom_config, + expected_output, + expected_state, + expected_has_more_pages, + expected_status_code, + expected_error, + mock_request, + ): + output, state, has_more_pages, status_code, error = self.task.run( + params={}, state=state, custom_config=custom_config + ) + self.assertEqual(expected_output, output) + self.assertEqual(expected_state, state) + self.assertEqual(expected_has_more_pages, has_more_pages) + self.assertEqual(expected_status_code, status_code) + self.assertEqual(expected_error, error) + validate(output, MonitorSiemLogsOutput.schema) + + @parameterized.expand( + [ + [ + "401", + {"query_config": {"receipt": {"caught_up": True, "next_page": "NDU1NA==", "query_date": "2000-01-01"}}}, + "Invalid API key provided.", + "Verify your API key configured in your connection is correct.", + 401, + ], + [ + "500", + {"query_config": {"receipt": {"caught_up": True, "next_page": "NDU1NA==", "query_date": "2000-01-02"}}}, + "Something unexpected occurred.", + "Check the logs and if the issue persists please contact support.", + 500, + ], + [ + "json_decode", + {"query_config": {"receipt": {"caught_up": True, "next_page": "NDU1NA==", "query_date": "2000-01-03"}}}, + "Received an unexpected response from the server.", + "(non-JSON or no response was received).", + 500, + ], + [ + "unknown", + {"query_config": {"receipt": {"caught_up": True, "next_page": "NDU1NA==", "query_date": "2000-01-04"}}}, + "Something unexpected occurred.", + "Check the logs and if the issue persists please contact support.", + 500, + ], + ] + ) + @patch("requests.Session.send", side_effect=Util.mocked_request) + def test_monitor_logs_errors( + self, + test_name, + state, + expected_cause, + expected_assistance, + expected_status_code, + mock_request, + ): + output, state, has_more_pages, status_code, error = self.task.run(params={}, state=state) + self.assertEqual(expected_status_code, status_code) + self.assertEqual(expected_cause, error.cause) + self.assertEqual(expected_assistance, error.assistance) + validate(output, MonitorSiemLogsOutput.schema) diff --git a/plugins/mimecast_v2/unit_test/util.py b/plugins/mimecast_v2/unit_test/util.py new file mode 100644 index 0000000000..e21604458d --- /dev/null +++ b/plugins/mimecast_v2/unit_test/util.py @@ -0,0 +1,110 @@ +import sys +import os +import json +import logging + +sys.path.append(os.path.abspath("../")) + +from requests.exceptions import HTTPError +from icon_mimecast_v2.connection.connection import Connection +from icon_mimecast_v2.connection.schema import Input +from icon_mimecast_v2.util.constants import BASE_URL +import gzip +from io import BytesIO + + +class Util: + @staticmethod + def default_connector(action, connect_params: object = None): + default_connection = Connection() + default_connection.logger = logging.getLogger("connection logger") + if connect_params: + params = connect_params + else: + params = { + Input.CLIENT_ID: {"secretKey": "test"}, + Input.CLIENT_SECRET: {"secretKey": "test"}, + } + default_connection.connect(params) + action.connection = default_connection + action.logger = logging.getLogger("action logger") + return action + + @staticmethod + def read_file_to_string(filename: str) -> str: + with open( + os.path.join(os.path.dirname(os.path.realpath(__file__)), filename), "r", encoding="utf-8" + ) as file_reader: + return file_reader.read() + + @staticmethod + def read_file_to_bytes(filename: str) -> bytes: + with open(os.path.join(os.path.dirname(os.path.realpath(__file__)), filename), "rb") as file_reader: + return file_reader.read() + + @staticmethod + def read_file_to_dict(filename: str) -> dict: + return json.loads(Util.read_file_to_string(filename)) + + @staticmethod + def mocked_request(*args, **kwargs): + class MockResponse: + def __init__(self, status_code: int, filename: str = None, url: str = None, gzip=False): + self.filename = filename + self.status_code = status_code + self.text = "This is some error text" + self.url = url + if filename: + self.text = Util.read_file_to_string(f"responses/{filename}.json.resp") + if gzip: + self.content = self._gzip_compress(self.text) + + def _gzip_compress(self, data): + """Compress content using gzip.""" + buf = BytesIO() # Create a buffer to hold the gzipped content + with gzip.GzipFile(fileobj=buf, mode="wb") as f: + f.write(data.encode("utf-8")) # Write the string data to gzip (must be bytes) + return buf.getvalue() + + def json(self): + return json.loads(self.text) + + def raise_for_status(self): + if self.status_code == 200: + return + raise HTTPError("Bad response", response=self) + + if args[0].url == f"{BASE_URL}oauth/token": + return MockResponse(200, "authenticate") + if args[0].url in [ + f"{BASE_URL}siem/v1/batch/events/cg?type=receipt&dateRangeStartsAt=2000-01-06&dateRangeEndsAt=2000-01-06&pageSize=100", + f"{BASE_URL}siem/v1/batch/events/cg?type=url+protect&dateRangeStartsAt=2000-01-06&dateRangeEndsAt=2000-01-06&pageSize=100", + f"{BASE_URL}siem/v1/batch/events/cg?type=attachment+protect&dateRangeStartsAt=2000-01-06&dateRangeEndsAt=2000-01-06&pageSize=100", + f"{BASE_URL}siem/v1/batch/events/cg?type=receipt&dateRangeStartsAt=2000-01-06&dateRangeEndsAt=2000-01-06&pageSize=100&nextPage=NDU1NA%3D%3D", + f"{BASE_URL}siem/v1/batch/events/cg?type=url+protect&dateRangeStartsAt=2000-01-06&dateRangeEndsAt=2000-01-06&pageSize=100&nextPage=NDU1NA%3D%3D", + f"{BASE_URL}siem/v1/batch/events/cg?type=attachment+protect&dateRangeStartsAt=2000-01-06&dateRangeEndsAt=2000-01-06&pageSize=100&nextPage=NDU1NA%3D%3D", + f"{BASE_URL}siem/v1/batch/events/cg?type=attachment+protect&dateRangeStartsAt=2000-01-07&dateRangeEndsAt=2000-01-07&pageSize=100", + f"{BASE_URL}siem/v1/batch/events/cg?type=receipt&dateRangeStartsAt=2000-01-07&dateRangeEndsAt=2000-01-07&pageSize=100", + f"{BASE_URL}siem/v1/batch/events/cg?type=url+protect&dateRangeStartsAt=2000-01-07&dateRangeEndsAt=2000-01-07&pageSize=100", + f"{BASE_URL}siem/v1/batch/events/cg?type=attachment+protect&dateRangeStartsAt=1999-12-31&dateRangeEndsAt=1999-12-31&pageSize=1", + f"{BASE_URL}siem/v1/batch/events/cg?type=receipt&dateRangeStartsAt=1999-12-31&dateRangeEndsAt=1999-12-31&pageSize=1", + f"{BASE_URL}siem/v1/batch/events/cg?type=url+protect&dateRangeStartsAt=1999-12-31&dateRangeEndsAt=1999-12-31&pageSize=1", + ]: + return MockResponse(200, "monitor_siem_logs_batch") + if ( + args[0].url + == f"{BASE_URL}siem/v1/batch/events/cg?type=receipt&dateRangeStartsAt=2000-01-02&dateRangeEndsAt=2000-01-02&pageSize=100" + ): + return MockResponse(401, "monitor_siem_logs_error") + if ( + args[0].url + == f"{BASE_URL}siem/v1/batch/events/cg?type=receipt&dateRangeStartsAt=2000-01-04&dateRangeEndsAt=2000-01-04&pageSize=100" + ): + return MockResponse(401, "monitor_siem_logs_json_error") + if ( + args[0].url + == f"{BASE_URL}siem/v1/batch/events/cg?type=receipt&dateRangeStartsAt=2000-01-05&dateRangeEndsAt=2000-01-05&pageSize=100" + ): + raise AttributeError + if args[0].url == "https://example.com/": + return MockResponse(200, "monitor_siem_logs", gzip=True)