from CommonServerPython import * from itertools import zip_longest from enum import StrEnum """ BRANDS ENUM """ class Brands(StrEnum): """ Enum representing different integration brands. """ FIREEYE_HX_V2 = "FireEyeHX v2" CROWDSTRIKE_FALCON = "CrowdstrikeFalcon" CORTEX_CORE_IR = "Cortex Core - IR" BUILTIN = "Builtin" MICROSOFT_DEFENDER_ADVANCED_THREAT_PROTECTION = "Microsoft Defender Advanced Threat Protection" @classmethod def get_all_values(cls) -> list[str]: """ Returns a list of all string values defined in the Enum. """ return [member.value for member in cls] """ COMMAND CLASS """ class Command: def __init__( self, brand: str, name: str, arg_mapping: dict, hard_coded_args: dict = None, ): """ Args: brand (str): The brand associated with the command. name (str): The name of the command. arg_mapping (dict): A dictionary containing the command arguments. The commands in this script must include at least one argument from this dictionary. hard_coded_args (dict): Additional arguments to add for the command, arguments with hard-coded values. """ self.brand = brand self.name = name self.arg_mapping = arg_mapping self.hard_coded_args = hard_coded_args def initialize_commands() -> list: """ Initializes the commands for the Isolate Endpoint. """ commands = [ Command( # On platform, we use the built-in commands and brand Builtin brand=Brands.BUILTIN if is_platform() else Brands.CORTEX_CORE_IR, name="isolateEndpoint" if is_platform() else "core-isolate-endpoint", arg_mapping={"endpoint_id": "endpoint_id"}, ), Command( brand=Brands.CROWDSTRIKE_FALCON, name="cs-falcon-contain-host", arg_mapping={"ids": "endpoint_id"}, ), Command( brand=Brands.FIREEYE_HX_V2, name="fireeye-hx-host-containment", arg_mapping={"agentId": "endpoint_id", "hostName": "endpoint_hostname"}, # command can use agentId or hostName ), Command( brand="Microsoft Defender ATP", # this name is used in get-endpoint-data script, name="microsoft-atp-isolate-machine", arg_mapping={"machine_id": "endpoint_id"}, hard_coded_args={"isolation_type": "Full", "comment": "Isolated endpoint with IsolateEndpoint script."}, ), ] return commands """ HELPER FUNCTIONS """ def is_endpoint_already_isolated(endpoint_data: dict, endpoint_args: dict, endpoint_output: dict) -> bool: """ Checks whether an endpoint is isolated already. Args: endpoint_data (dict): A dictionary containing endpoint details. endpoint_args (dict): The arguments used in the command execution. endpoint_output (dict): A list to store structured output results. Returns: bool: True if the endpoint is isolated, False otherwise. """ demisto.debug(f"Got endpoint {endpoint_data} with field isIsolated{endpoint_data.get('IsIsolated')}") is_isolated = endpoint_data.get("IsIsolated", "No") if is_isolated == "No": return False message = "The endpoint is already isolated." create_message_to_context_and_hr( is_isolated=True, endpoint_args=endpoint_args, result="Success", message=message, endpoint_output=endpoint_output ) return True def create_message_to_context_and_hr( is_isolated: bool, endpoint_args: dict, result: str, message: str, endpoint_output: dict ) -> None: """ Generates a structured message for context and human-readable outputs. Args: is_isolated (bool): Whether the endpoint is isolated. endpoint_args (dict): A dictionary containing endpoint details such as hostname, ID, or IP. result (str): The result status, e.g., "Success" or "Fail". message (str): A message explaining the result. endpoint_output (dict): A list to store the structured output for context. """ endpoint_hostname = endpoint_args.get("endpoint_id") or endpoint_args.get("endpoint_ip") brand = endpoint_args.get("endpoint_brand", "") if brand == "Microsoft Defender ATP": # convert brand brand = Brands.MICROSOFT_DEFENDER_ADVANCED_THREAT_PROTECTION endpoint_output["Endpoint"] = endpoint_hostname endpoint_output["Result"] = result endpoint_output["Source"] = brand endpoint_output["Message"] = message endpoint_output["Isolated"] = "Yes" if is_isolated else "No" def are_there_missing_args(command: Command, endpoint_args: dict, endpoint_output: dict) -> bool: """ Checks if all required arguments are existing in the provided arguments. Args: command (Command): The command to use for checking the required arguments. endpoint_args (dict): A dictionary containing the provided arguments. Returns: bool: True if all expected arguments are missing, False otherwise. """ if not command.arg_mapping: # If there are no expected args, return False return False is_missing_args = all( endpoint_args.get(key, "") == "" for key in command.arg_mapping.values() ) # checks if *all* args are missing if is_missing_args: demisto.debug(f"Missing the next args {endpoint_args} for command.name") create_message_to_context_and_hr( is_isolated=False, endpoint_args=endpoint_args, result="Fail", message=f"Missing args for {command.name}.", endpoint_output=endpoint_output, ) return True return False def map_args(command: Command, args: dict) -> dict: """ Maps provided arguments to their expected keys based on a given mapping. Args: command (Command): The command that its args need to be mapped. args (dict): A dictionary containing the provided arguments. Returns: dict: A dictionary with mapped arguments, using expected keys with corresponding values from args. """ mapped_args = {k: args.get(v, "") for k, v in command.arg_mapping.items()} if command.hard_coded_args: mapped_args.update(command.hard_coded_args) return mapped_args def map_zipped_args(endpoint_ids: list, endpoint_ips: list) -> list: """ Combines agent IDs and IPs into a list of dictionaries. Args: endpoint_ids (list): A list of agent IDs. endpoint_ips (list): A list of agent IPs. Returns: list: A list of dictionaries, each containing 'endpoint_id' and 'endpoint_ip'. """ return [ {"endpoint_id": endpoint_id, "endpoint_ip": endpoint_ip} for endpoint_id, endpoint_ip in zip_longest(endpoint_ids, endpoint_ips, fillvalue="") ] def check_missing_executed_args_in_output(zipped_args: list, valid_args: list, outputs: list) -> None: """ Checks if any of the given agent details (ID, IP) exist in a list of valid arguments. If no match is found, a failure message is added to the context and human-readable outputs. Args: zipped_args (list): A list of dictionaries, each containing 'endpoint_id', 'endpoint_ip'. valid_args (list): A list of dictionaries representing valid agents with corresponding details. outputs (list): A list to store structured output results. """ for args in zipped_args: endpoint_id = args.get("endpoint_id", "") endpoint_ip = args.get("endpoint_ip", "") are_args_found = False for entry in valid_args: demisto.debug(f"Got {entry=}, and comparing it to {endpoint_id=} and {endpoint_ip=}") # Checks if any of the args exists in valid_args if (endpoint_id and entry.get("endpoint_id") == endpoint_id) or ( endpoint_ip and entry.get("endpoint_ip") == endpoint_ip ): are_args_found = True if not are_args_found: endpoint_context_output: dict = {} create_message_to_context_and_hr( is_isolated=False, endpoint_args=args, result="Fail", message="Did not find information on endpoint in any available brand.", endpoint_output=endpoint_context_output, ) outputs.append(endpoint_context_output) def get_args_from_endpoint_data(endpoint_data: dict) -> dict: """ Extracts agent details from endpoint data and maps them to a dictionary. Args: endpoint_data (dict): A dictionary containing endpoint details such as hostname, ID, IP address, and brand. Returns: dict: A dictionary with extracted values, including 'endpoint_id', 'endpoint_hostname', 'endpoint_ip', and 'endpoint_brand'. """ return { "endpoint_id": endpoint_data.get("ID", ""), "endpoint_ip": endpoint_data.get("IPAddress", ""), "endpoint_brand": endpoint_data.get("Brand", ""), "endpoint_hostname": endpoint_data.get("Hostname", ""), "endpoint_message": endpoint_data.get("Message", ""), } def structure_endpoints_data(get_endpoint_data_results: dict | list | None) -> list: """ Structures and filters endpoint data, ensuring it is in list format and contains only the entry of the context. Args: get_endpoint_data_results (dict | list | None): The raw endpoint data, which may be a dictionary, list, or None. Returns: list: A structured list containing the entry of the context, excluding None values. """ if not get_endpoint_data_results: return [] if not isinstance(get_endpoint_data_results, list): get_endpoint_data_results = [get_endpoint_data_results] # Remove None values structured_list = [item for item in get_endpoint_data_results if item is not None] if structured_list and isinstance(structured_list[0], list): return structured_list[0] return structured_list def handle_raw_response_results(command: Command, raw_response: dict, endpoint_args: dict, endpoint_output: dict) -> None: """ Handles the raw response of a command execution by determining success or failure and updating outputs accordingly. Args: command (Command): The executed command object. raw_response (dict): The raw response returned from the command execution. endpoint_args (dict): The arguments used in the command execution. endpoint_output (dict): A list to store structured output results. """ endpoint_id = endpoint_args.get("endpoint_id", "") if is_error(raw_response): demisto.debug(f"Got an error from raw_response with {endpoint_args}") create_message_to_context_and_hr( is_isolated=False, endpoint_args=endpoint_args, result="Fail", message=f"Failed to isolate {endpoint_id} with command {command.name}." f"Error:{get_error(raw_response)}", endpoint_output=endpoint_output, ) else: create_message_to_context_and_hr( is_isolated=True, endpoint_args=endpoint_args, result="Success", message=f"{endpoint_id} was isolated successfully with command {command.name}.", endpoint_output=endpoint_output, ) def find_command_by_brand(commands: list[Command], brand: str): """ Finds and returns the command from the list that matches the specified brand. Args: commands (list): A list of Command objects to search through. brand (str): The brand name to match against the command's brand. Returns: Command: The matching Command object. """ for command in commands: if command.brand == brand: return command return None def run_commands_for_endpoint(commands: list, endpoint_args: dict, endpoint_output: dict) -> None: # type: ignore[arg-type,union-attr] """ Processes an endpoint by executing isolation commands and updating outputs accordingly. Args: commands (list): A list of available commands for isolation. endpoint_args (dict): The arguments provided for the isolation operation. endpoint_output (dict): A dictionary to store structured output results. """ demisto.debug(f"Got into the run_commands_for_endpoint command with {endpoint_args}") command = find_command_by_brand(commands, endpoint_args.get("endpoint_brand", "")) if are_there_missing_args(command, endpoint_args, endpoint_output): # type: ignore[arg-type] return mapped_args = map_args(command, endpoint_args) demisto.debug(f"Executing command {command.name} with {endpoint_args=}") raw_response = demisto.executeCommand(command.name, mapped_args) demisto.debug(f"Got raw response for execute_command {command.name} with {endpoint_args=}: {raw_response=}") handle_raw_response_results(command, raw_response, endpoint_args, endpoint_output) def prepare_args() -> tuple[dict, list]: """ Prepares and validates the script arguments for endpoint data collection. Returns: tuple[dict, list]: - A dictionary containing the processed endpoint arguments, including a default list of brands if not provided. - A list of zipped argument pairs combining endpoint IDs and IPs. """ endpoint_args = demisto.args() endpoint_ids = argToList(endpoint_args.get("endpoint_id", [])) endpoint_ips = argToList(endpoint_args.get("endpoint_ip", [])) brands_to_run = argToList(endpoint_args.get("brands", [])) if not any((endpoint_ids, endpoint_ips)): raise ValueError("At least one of the following arguments must be specified: endpoint_id or endpoint_ip.") if not brands_to_run: # In case no brands selected, the default is all brands. # We want to send to get-endpoint-data only the brands this script supports. endpoint_args["brands"] = Brands.get_all_values() zipped_args = map_zipped_args(endpoint_ids, endpoint_ips) return endpoint_args, zipped_args def process_endpoints(endpoint_data_results: list, commands: list[Command]) -> tuple[list, list, list]: """ Processes endpoint data results and executes the appropriate commands for isolation. Args: endpoint_data_results (list): A list of endpoint data results retrieved from get-endpoint-data. commands (list[Command]): A list of Command objects to run on each endpoint. Returns: tuple[list, list, list]: - results (list): A list of command execution results (currently unused). - context_outputs (list): A list of context output dictionaries for each processed endpoint. - args_from_endpoint_data (list): A list of argument dictionaries built from the endpoint data. """ results: list = [] context_outputs: list = [] args_from_endpoint_data: list = [] for endpoint_data in endpoint_data_results: endpoint_context_output: dict = {} endpoint_args = get_args_from_endpoint_data(endpoint_data) demisto.debug(f"Running with args {endpoint_args=}") # Skip the failing endpoints from get-data-endpoint if "fail" in endpoint_args.get("endpoint_message", "").lower(): demisto.debug(f"Skipping endpoint {endpoint_args} because of a failing error from get-endpoint-data.") continue if is_endpoint_already_isolated(endpoint_data, endpoint_args, endpoint_context_output): demisto.debug(f"Skipping endpoint {endpoint_args} because it is already isolated.") args_from_endpoint_data.append(endpoint_args) context_outputs.append(endpoint_context_output) continue demisto.debug(f"Continue isolating endpoint {endpoint_args}") args_from_endpoint_data.append(endpoint_args) run_commands_for_endpoint(commands, endpoint_args, endpoint_context_output) context_outputs.append(endpoint_context_output) return results, context_outputs, args_from_endpoint_data def main(): # pragma: no cover try: endpoint_args, zipped_args = prepare_args() commands = initialize_commands() executed_command = execute_command(command="get-endpoint-data", args=endpoint_args) demisto.debug(f"Response from get-endpoint-data: {executed_command=}") endpoint_data_results = structure_endpoints_data(executed_command) results, context_outputs, args_from_endpoint_data = process_endpoints(endpoint_data_results, commands) # comparing the executed args for isolated-endpoint with the input args check_missing_executed_args_in_output(zipped_args, args_from_endpoint_data, context_outputs) readable_output = tableToMarkdown(name="IsolateEndpoint Results", t=context_outputs, removeNull=True) results.append( CommandResults( outputs_prefix="IsolateEndpoint", outputs_key_field="Endpoint", outputs=context_outputs, readable_output=readable_output, ) ) return_results(results) except Exception as e: return_error(f"Failed to execute isolate-endpoint. Error: {e!s}") """ ENTRY POINT """ if __name__ in ("__main__", "__builtin__", "builtins"): main()