import abc import enum import demistomock as demisto import pandas import pandas as pd from CommonServerPython import * INTEGRATION_NAME = "PrismaCloudCompute" ISSUES_INPUT_PATH = "EnrichedComplianceIssue" class ComplianceObj(enum.Enum): HOST = "host" CONTAINER = "container" IMAGE = "image" class ComplianceObject(abc.ABC): def __init__(self, object_type: ComplianceObj, input_context_path: str, output_context_id: str): self.object_type = object_type self.capitalized_type = object_type.value.capitalize() self.input_context_path = input_context_path self.output_context_path = f"{INTEGRATION_NAME}.ComplianceTable.{self.capitalized_type}" self.output_context_id = output_context_id def get_input_context_id(self, obj: dict): try: return self._get_input_context_id(obj) except AttributeError as err: raise Exception(f"Input context does not match provided type: {self.object_type.value}") from err @abc.abstractmethod def _get_input_context_id(self, obj: dict): pass @abc.abstractmethod def get_data(self, input_data: dict, identifier: str, issues: list) -> dict: pass class Host(ComplianceObject): def __init__(self): super().__init__( object_type=ComplianceObj.HOST, input_context_path=f"{INTEGRATION_NAME}.ReportHostScan", output_context_id="Hostname" ) def _get_input_context_id(self, obj: dict): return obj.get("hostname") def get_data(self, input_data: dict, identifier: str, issues: list) -> dict: """Get the host data as needed in the table. Args: input_data (dict): The input context data containing information about the host. identifier (str): The host name. issues (list): List of issues the host appeared in. Returns: (dict) The host data as needed in the table. """ compliance_distribution = input_data.get("complianceDistribution") cloud_metadata = input_data.get("cloudMetadata") cloud_metadata.pop("labels", None) if cloud_metadata else None host_data = { self.output_context_id: identifier, "ComplianceIssues": issues, "ComplianceDistribution": compliance_distribution, "CloudMetaData": cloud_metadata, } return host_data class Container(ComplianceObject): def __init__(self): super().__init__( object_type=ComplianceObj.CONTAINER, input_context_path=f"{INTEGRATION_NAME}.ContainersScanResults", output_context_id="ContainerID", ) def _get_input_context_id(self, obj: dict): return obj.get("info", {}).get("id") def get_data(self, input_data: dict, identifier: str, issues: list) -> dict: """Get the container data as needed in the table. Args: input_data (dict): The input context data containing information about the container. identifier (str): The container ID. issues (list): List of issues the container appeared in. Returns: (dict) The container data as needed in the table. """ container_info = input_data.get("info", {}) compliance_dist = container_info.get("complianceDistribution") image_name = container_info.get("imageName") cloud_metadata = container_info.get("cloudMetadata", {}) cloud_metadata.pop("labels", None) hostname = input_data.get("hostname", {}) container_data = { self.output_context_id: identifier, "ComplianceIssues": issues, "ComplianceDistribution": compliance_dist, "Hostname": hostname, "ImageName": image_name, "CloudMetaData": cloud_metadata, } return container_data class Image(ComplianceObject): def __init__(self): super().__init__( object_type=ComplianceObj.IMAGE, input_context_path=f"{INTEGRATION_NAME}.ReportsImagesScan", output_context_id="ImageID", ) def _get_input_context_id(self, obj: dict): return obj.get("id") def get_data(self, input_data: dict, identifier: str, issues: list) -> dict: """Get the image data as needed in the table. Args: input_data (dict): The input context data containing information about the image. identifier (str): The image id. issues (list): List of issues it appeared in. Returns: (dict) The image data as needed in the table. """ compliance_dist = input_data.get("complianceDistribution") hosts = list(input_data.get("hosts", {}).keys()) instances_data = input_data.get("instances", []) image_instances = [instance_data.get("image") for instance_data in instances_data] cloud_metadata = input_data.get("cloudMetadata", {}) cloud_metadata.pop("labels", {}) image_data = { self.output_context_id: identifier, "ComplianceIssues": issues, "ComplianceDistribution": compliance_dist, "Hosts": hosts, "ImageInstances": image_instances, "CloudMetaData": cloud_metadata, } return image_data COMPLIANCE_OBJ_CLASS = { ComplianceObj.HOST.value: Host(), ComplianceObj.CONTAINER.value: Container(), ComplianceObj.IMAGE.value: Image(), } def get_input_object_list(context_data: dict, compliance_obj: ComplianceObject) -> list: """Get list of the input objects that the table will be updated with. Args: context_data (dict): The context data object the input objects are stored in. compliance_obj (ComplianceObject): The resource type class to get. Returns: (List[dict]) The list of the resource specified. """ input_objects = demisto.get(context_data, compliance_obj.input_context_path) if type(input_objects) is list: return input_objects return [input_objects] def get_output_object_list(compliance_obj: ComplianceObject, grid_id: str = "") -> tuple[list, list]: """Get the already present resource list in the table. Args: compliance_obj (ComplianceObject): The resource type class to get the list of. Returns: (List[dict], List[str]): The list of the specified resource, list of their ids. """ if grid_id: incident = demisto.incident() custom_fields = incident.get("CustomFields", {}) or {} output_objects = custom_fields.get(grid_id) output_id = compliance_obj.output_context_id.lower() else: context_data = demisto.context() compliance_table_context = context_data.get(f"{INTEGRATION_NAME}", {}).get("ComplianceTable", {}) output_objects = compliance_table_context.get(compliance_obj.capitalized_type, []) output_id = compliance_obj.output_context_id if type(output_objects) is list: output_objects_list = output_objects else: output_objects_list = [output_objects] return output_objects_list, [output_obj.get(output_id) for output_obj in output_objects_list] def update_output_obj_with_issues(compliance_obj: ComplianceObject, input_obj_id: str, issues: list): """Update an object in the output table with new issues. Modify the context data. Args: compliance_obj (ComplianceObject): The resource type class to update. input_obj_id (str): The id of the resource to update. issues (List[str]): Issue records to update with. """ context_key = compliance_obj.output_context_path output_objs, output_objs_ids = get_output_object_list(compliance_obj) output_obj_index = output_objs_ids.index(input_obj_id) output_obj = output_objs[output_obj_index] previous_issues = output_obj.get("ComplianceIssues", []) previous_issues = previous_issues if type(previous_issues) is list else [previous_issues] non_duplicated_issues = [issue for issue in issues if issue not in previous_issues] if non_duplicated_issues: demisto.debug( f"Updating {compliance_obj.object_type.value} in id {input_obj_id} with new issues: {non_duplicated_issues}" ) new_issues_list = previous_issues + non_duplicated_issues output_obj.update({"ComplianceIssues": new_issues_list}) output_id_path = compliance_obj.output_context_id demisto.results( { "Type": entryTypes["note"], "ContentsFormat": formats["json"], "Contents": output_obj, "HumanReadable": tableToMarkdown( f"Updating {compliance_obj.object_type.value} ({input_obj_id}) with new compliance issues", non_duplicated_issues, "Compliance Issue", ), "EntryContext": {f"{context_key}(val.{output_id_path} == obj.{output_id_path})": output_obj}, } ) def update_context_data(all_object_type_data: list, output_objs_to_append: dict, compliance_obj: ComplianceObject) -> None: """Update context data with new objects and new compliance issues that need updating. Args: all_object_type_data (list): List of the new data. output_objs_to_append (dict): Dict of (key: object id) and their (value: new compliance issues). compliance_obj (ComplianceObject): The compliance object the data refers to. """ if all_object_type_data: appendContext(compliance_obj.output_context_path, all_object_type_data) demisto.debug(f"The objects to update are {list(output_objs_to_append.keys())}. Updating") for obj_to_update_id in output_objs_to_append: update_output_obj_with_issues(compliance_obj, obj_to_update_id, output_objs_to_append[obj_to_update_id]) def turn_pd_grid_to_context_table(pd_grid: pandas.DataFrame) -> list: """Turn pandas DataFrame object to a list of dicts to save in the context_data. Args: pd_grid (pandas.DataFrame): The pandas dataframe to convert Returns: (list) of dicts of data inside the pd_grid provided. """ context_table = [] for record in pd_grid.to_dict(orient="records"): new_record_dict = {} for record_key, record_value in record.items(): if record_value and (isinstance(record_value, dict | list) or pd.notnull(record_value)): if isinstance(record_value, dict): str_record_value = "\n".join(f"{dict_key}: {dict_value}" for dict_key, dict_value in record_value.items()) elif isinstance(record_value, list): str_record_value = "\n\n".join(record_value) else: str_record_value = str(record_value) new_record_dict.update({record_key.lower(): str_record_value}) if new_record_dict: context_table.append(new_record_dict) return context_table def update_grid_table( all_new_data: list, output_objs_to_append: dict, compliance_obj: ComplianceObject, grid_id: str, current_table: list ) -> None: """Update grid in the incident context data. Args: all_new_data (list): List of the new data. output_objs_to_append (dict): Dict of object ids (keys) and their new compliance issues (values). compliance_obj (ComplianceObject): The compliance object the data refers to. grid_id (str): The grid id to update. current_table (list): The already present table in the incident. """ demisto.debug( f"Updating grid {grid_id} table with all new:\n{all_new_data}\n" f"output_objs:\n{output_objs_to_append}\ncurrent table is:\n{current_table}" ) current_df = pd.DataFrame(current_table) if current_table else pd.DataFrame() demisto.debug(f"Current dataframe {current_df}") out_id = compliance_obj.output_context_id.lower() for id_to_update in output_objs_to_append: previous_issues = current_df.loc[current_df[out_id] == id_to_update, "complianceissues"].iloc[0] issues = output_objs_to_append[id_to_update] non_duplicated_issues = [issue for issue in issues if issue not in previous_issues] if non_duplicated_issues: merged_issue_list = previous_issues + "\n\n" + "\n\n".join(non_duplicated_issues) current_df.loc[current_df[out_id] == id_to_update, "complianceissues"] = merged_issue_list new_grid = pd.concat([pd.DataFrame(all_new_data), current_df]) demisto.debug(f"New dataframe after concat and sort:\n{new_grid}") # filter empty values in the generated table and turn to dict context_table = turn_pd_grid_to_context_table(new_grid) demisto.debug(f"New table for context: {context_table}") # Execute automation 'setIncident` which change the Context data in the incident demisto.executeCommand( "setIncident", { "customFields": { grid_id: context_table, }, }, ) def create_issue_record(issue_obj: dict): """Create a unique issue string from the issue_obj dict.""" return f"{issue_obj.get('id')} ({issue_obj.get('severity')} | {issue_obj.get('type')}) - {issue_obj.get('title')}" def categorize_issue_in_object( issue: str, input_obj: dict, compliance_obj: ComplianceObject, output_objs_ids: list, output_objs_to_create: dict, output_objs_to_append: dict, ) -> None: """Categorize the new issue information based on if the object is new or old. WARNING: Mutates output_objs_to_create and output_objs_to_append. Args: issue (str): The issue to categorize. input_obj (dict): The input object dict. compliance_obj (ComplianceObject): The type of compliance object. output_objs_ids (list): The ids of the objects already present in the output. output_objs_to_append (dict): Append the issues to this dict if they are already present in the output. output_objs_to_create (dict): Add the issues with the corresponding object id to this dict if they are new. """ input_obj_id = compliance_obj.get_input_context_id(input_obj) if input_obj_id not in output_objs_ids: demisto.debug(f"Got new {compliance_obj.object_type.value} with id {input_obj_id} in issue {issue}") if input_obj_id in output_objs_to_create: output_objs_to_create[input_obj_id]["issues"].append(issue) else: output_objs_to_create[input_obj_id] = {"input_obj": input_obj, "issues": [issue]} else: demisto.debug(f"Got old {compliance_obj.object_type.value} with id {input_obj_id} in issue {issue}") if input_obj_id in output_objs_to_append: output_objs_to_append[input_obj_id].append(issue) else: output_objs_to_append[input_obj_id] = [issue] def update_objects_by_issues(compliance_obj: ComplianceObject, root_context_key: str, grid_id: str = ""): """Go over enriched issues in the context data and update the output table. Args: compliance_obj (ComplianceObject): The resource type class to update in the table. root_context_key (str): The context data path to the root, containing the input data. grid_id (str): The grid id to write the outputs to. """ if root_context_key: issues_input_objects = demisto.get(demisto.context(), root_context_key) else: issues_input_objects = demisto.context() issues_input_objects = issues_input_objects if type(issues_input_objects) is list else [issues_input_objects] output_objs, output_objs_ids = get_output_object_list(compliance_obj, grid_id) demisto.debug(f"Starting update, the already present output object ids are {output_objs_ids}") output_objs_to_create: dict = {} output_objs_to_append: dict = {} demisto.debug(f"Starting to go over {len(issues_input_objects)} issues and search for {compliance_obj.object_type.value}") for issue_input_obj in issues_input_objects: issue = create_issue_record(issue_input_obj.get(ISSUES_INPUT_PATH, {})) input_objs = get_input_object_list(issue_input_obj, compliance_obj) demisto.debug(f"Got {len(input_objs)} {compliance_obj.object_type.value} for issue {issue}") for input_obj in input_objs: categorize_issue_in_object( issue, input_obj, compliance_obj, output_objs_ids, output_objs_to_create, output_objs_to_append ) demisto.debug(f"The new objects to create are {list(output_objs_to_create.keys())}. Creating") all_object_type_data = [] for obj_to_create_id in output_objs_to_create: output_context_data = compliance_obj.get_data( output_objs_to_create[obj_to_create_id]["input_obj"], obj_to_create_id, output_objs_to_create[obj_to_create_id]["issues"], ) all_object_type_data.append(output_context_data) # Append after collecting all the new data if grid_id: update_grid_table(all_object_type_data, output_objs_to_append, compliance_obj, grid_id, output_objs) else: update_context_data(all_object_type_data, output_objs_to_append, compliance_obj) def update_context_paths(demisto_args: dict): compliance_obj = COMPLIANCE_OBJ_CLASS[demisto_args.get("resourceType", "").lower()] return update_objects_by_issues(compliance_obj, demisto_args.get("contextPath", ""), demisto_args.get("gridID", "")) def main(): # pragma: no cover try: return_results(update_context_paths(demisto_args=demisto.args())) except Exception as ex: return_error(f"Failed to execute PrismaCloudComputeComplianceTable. Error: {ex!s}") """ ENTRY POINT """ if __name__ in ("__main__", "__builtin__", "builtins"): main()