GetIdsFromCustomContent

Source

import json
import os
import tarfile
import tempfile
import traceback
from collections import defaultdict
from pathlib import Path
from typing import Any

import demistomock as demisto
from CommonServerPython import *
from demisto_sdk.commands.common.tools import _get_file_id, find_type, get_file, get_file_displayed_name

from CommonServerUserPython import *


def update_file_prefix(file_name: str) -> str:
    """
    Custom content scripts are prefixed with automation instead of script.
    Removing the "playbook-" prefix from files name.
    """
    if file_name.startswith("playbook-"):
        return file_name[len("playbook-") :]
    if file_name.startswith("automation-"):
        return file_name.replace("automation-", "script-")
    return file_name


def get_content_details(tar_file_handler: Any, member_file: Any) -> tuple[str, dict]:
    """Get content id from tar member file.

    Args:
        tar_file_handler: Tarfile open handler that contains the member file to inspect.
        member_file: The member file in the tar file to inspect.

    Return:
        (entity, file_id_name) of the member file.
    """
    file_type_str = ""
    file_id = None
    with tempfile.TemporaryDirectory() as tmp_dir_name:
        file_name = update_file_prefix(member_file.name.strip("/"))
        file_path = os.path.join(tmp_dir_name, file_name)
        with open(file_path, "w") as file_desc:
            if extracted_file := tar_file_handler.extractfile(member_file):
                file_desc.write(extracted_file.read().decode("utf-8"))
            else:
                raise Exception(f"Could not extract file {file_name} from tar: {file_path}")

        if not os.path.isfile(file_path):
            raise Exception(f"Could not create file {file_path}")

        file_type = find_type(path=file_path)
        file_type_str = file_type.value if file_type else ""
        if file_type_str == "automation":
            file_type_str = "script"

        file_type_str = "list" if not file_type_str and file_name.startswith("list-") else file_type_str
        file_dict = get_file(file_path, Path(file_name).suffix[1:])
        file_id = _get_file_id(file_type_str, file_dict)
        file_id = file_id if file_id else file_dict.get("id")
        file_name = get_file_displayed_name(file_path)

    file_id_name = {"id": file_id, "name": file_name}
    return file_type_str, file_id_name


def get_custom_content_ids(file_entry_id: Any) -> dict:
    """Get custom content ids from custom content bundle.

    Args:
         file_entry_id (str): The entry id of the custom content zip file.

    Return:
        A dict of custom content ids.
    """
    custom_content_ids: defaultdict[Any, list] = defaultdict(list)
    get_file_path_res = demisto.getFilePath(file_entry_id)
    custom_content_file_path = get_file_path_res.get("path")
    if not custom_content_file_path:
        raise ValueError(f"Could not find file path for entry id {file_entry_id}")
    custom_content_tar_file = tarfile.open(custom_content_file_path)
    custom_content_members = custom_content_tar_file.getmembers()

    for custom_content_member in custom_content_members:
        entity, entity_id_name = get_content_details(custom_content_tar_file, custom_content_member)
        if entity and entity_id_name.get("id"):
            custom_content_ids[entity].append(entity_id_name)
        else:
            raise Exception(f"Could not parse content type and id from file name {custom_content_member.name}")

    return custom_content_ids


""" COMMAND FUNCTION """


def filter_lists(include: list, exclude: list) -> list:
    return [item for item in include if item.get("id") not in exclude]


def get_included_ids_command(args: dict[str, Any]) -> CommandResults:
    """Get included ids from installed custom content unless id is excluded.

    Args:
        exclude_ids_list (List[Dict[str:List[str]]]): A list of dicts, each specifies entity ids to exclude.
            (example: [{'integration': ['HelloWorld', 'MyIntegration']}, {'script': ['say_hello']}]
        file_entry_id (str): The entry id of the custom content zip file.

    Return:
        CommandResults Outputs with included ids dict and excluded ids dict, ready to pass to DeleteContent script.
    """
    if (excluded_ids_dicts := args.get("exclude_ids_list", [])) and not isinstance(excluded_ids_dicts, list):
        try:
            excluded_ids_dicts = json.loads(str(args.get("exclude_ids_list")))
        except json.JSONDecodeError as err:
            raise ValueError(f"Failed decoding excluded_ids_list as json: {err!s}")

    custom_content_ids = get_custom_content_ids(file_entry_id=args.get("file_entry_id"))

    included_custom_ids_names = {}
    excluded_ids: defaultdict[Any, list] = defaultdict(list)
    if excluded_ids_dicts:
        # Merge exclusion dicts
        for excluded_ids_dict in excluded_ids_dicts:
            for excluded_entity in excluded_ids_dict:
                excluded_ids[excluded_entity] += excluded_ids_dict.get(excluded_entity, [])

        # Exclude what is relevant
        for custom_entity in custom_content_ids:
            included_custom_ids_names[custom_entity] = filter_lists(
                include=custom_content_ids.get(custom_entity, []), exclude=excluded_ids.get(custom_entity, [])
            )

        # Remove included entities from excluded dict
        for entity in included_custom_ids_names:
            if entity in excluded_ids:
                excluded_ids.pop(entity)

    else:
        included_custom_ids_names = custom_content_ids

    included_custom_ids = {key: [value["id"] for value in lst] for key, lst in included_custom_ids_names.items() if lst}
    included_custom_name = {key: [value["name"] for value in lst] for key, lst in included_custom_ids_names.items() if lst}
    return CommandResults(
        outputs_prefix="GetIdsFromCustomContent",
        outputs_key_field="",
        outputs={
            "included_ids": included_custom_ids,
            "excluded_ids": {key: value for key, value in excluded_ids.items() if value},
        },
        readable_output=tableToMarkdown("Included ids", included_custom_name) + tableToMarkdown("Excluded ids", excluded_ids),
    )


def main():  # pragma: no cover
    try:
        return_results(get_included_ids_command(demisto.args()))
    except Exception as ex:
        return_error(f"Failed to execute GetIdsFromCustomContent. Error: {ex!s}\n Traceback: {traceback.format_exc()}")


""" ENTRY POINT """


if __name__ in ("__main__", "__builtin__", "builtins"):  # pragma: no cover
    main()

README