FlashpointVulnerabilityDetails

import demistomock as demisto
from CommonServerPython import *

VULNERABILITY_LIST_COMMAND = "flashpoint-ignite-vulnerability-list"
SET_INCIDENT_COMMAND = "setIncident"

VULNERABILITY_DETAILS_FIELD = "flashpointvulnerabilitydetails"

MESSAGES = {
    "NO_RECORDS_FOUND": "No vulnerability data found.",
}

ERROR_MESSAGES = {
    "FAILED_COMMAND": "Failed to execute '{}' command. Error: {}",
}

""" HELPER FUNCTIONS """


def get_command_result(command_results: list) -> dict:
    """
    Return the first non-error result from an executeCommand output list.

    :type command_results: list
    :param command_results: Raw list returned by demisto.executeCommand.

    :return: First successful result entry, or an empty dict if all entries are errors.
    :rtype: dict
    """
    for result in command_results:
        if not isError(result):
            return result
    return {}


def execute_command_safe(command: str, args: dict) -> tuple[dict, Any]:
    """
    Execute a demisto command and return a (result, error) tuple.

    :type command: str
    :param command: Name of the demisto command to execute.

    :type args: dict
    :param args: Arguments to pass to the command.

    :return: (result_dict, None) on success; ({}, error_contents) on failure.
             The error value is the raw Contents field from the error entry, which may be str, dict, or list.
    :rtype: tuple[dict, Any]
    """
    raw = demisto.executeCommand(command, args)
    if not isinstance(raw, list):
        raw = [raw]
    result = get_command_result(raw)
    if not result:
        error = raw[0].get("Contents", "Unknown error") if raw else "Unknown error"
        return {}, error
    return result, None


def get_vulnerability_rows() -> list[dict]:
    """
    Get the vulnerability details of the alert from the incident data.

    The vulnerability details are populated on the incident by the incoming mapper, which maps the 'resource.vulns'
    field of the alert to the 'flashpointvulnerabilitydetails' grid field. Each row of the grid holds the details of a
    single vulnerability, out of which the 'vuln_id' key is used to look up the CVE information.

    :return: List of the vulnerability details present in the incident, or an empty list when the incident has no
        vulnerability details.
    :rtype: list[dict]
    """
    incident_info = demisto.incident()
    rows = demisto.get(incident_info, f"CustomFields.{VULNERABILITY_DETAILS_FIELD}") or []

    return [row for row in rows if isinstance(row, dict)]


def get_vulnerabilities(vulnerability_ids: list[str]) -> dict:
    """
    Retrieve the vulnerabilities of the given vulnerability IDs using the vulnerability list command.

    All IDs are requested in a single call, so a large grid does not result in one API call per row.

    :type vulnerability_ids: list[str]
    :param vulnerability_ids: Vulnerability IDs to look up.

    :return: Raw entry returned by the vulnerability list command.
    :rtype: dict

    :raises ValueError: If the vulnerability list command fails.
    """
    args = {
        "vulnerability_ids": ",".join(vulnerability_ids),
        "size": len(vulnerability_ids),
    }

    result, err = execute_command_safe(VULNERABILITY_LIST_COMMAND, args)
    if err:
        raise ValueError(ERROR_MESSAGES["FAILED_COMMAND"].format(VULNERABILITY_LIST_COMMAND, err))

    return result


def get_cve_ids_by_vulnerability(result: dict) -> dict[str, str]:
    """
    Prepare the mapping of the vulnerability ID to its CVE IDs from the vulnerability list command result.

    :type result: dict
    :param result: Raw entry returned by the vulnerability list command.

    :return: Mapping of the vulnerability ID to its comma-separated CVE IDs. IDs with no CVE information are omitted
        from the mapping.
    :rtype: dict[str, str]
    """
    vulnerabilities = demisto.get(result, "Contents.results") or []

    cve_ids_by_vulnerability = {}
    for vulnerability in vulnerabilities:
        vulnerability_id = str(vulnerability.get("id", ""))
        cve_ids = vulnerability.get("cve_ids", [])
        if vulnerability_id and cve_ids:
            cve_ids_by_vulnerability[vulnerability_id] = ", ".join(cve_ids)
    return cve_ids_by_vulnerability


def build_grid_rows(rows: list[dict], cve_ids_by_vulnerability: dict[str, str]) -> list[dict]:
    """
    Add the 'cve_ids' key to the existing vulnerability details rows.

    :type rows: list[dict]
    :param rows: Existing vulnerability grid rows.

    :type cve_ids_by_vulnerability: dict[str, str]
    :param cve_ids_by_vulnerability: Mapping of the vulnerability ID to its comma-separated CVE IDs.

    :return: The same rows with the 'cve_ids' key populated where CVE information is available.
    :rtype: list[dict]
    """
    for row in rows:
        vulnerability_id = row.get("vuln_id", "")
        row["cve_ids"] = cve_ids_by_vulnerability.get(vulnerability_id, row.get("cve_ids", ""))
    return rows


""" COMMAND FUNCTION """


def get_vulnerability_details() -> dict | CommandResults:
    """
    Populate the CVE IDs of the vulnerability details incident field.

    The vulnerability IDs are read from the vulnerability details of the incident.

    :return: Raw entry returned by the vulnerability list command, or a no records found message when the incident has
        no vulnerability details. The CVE IDs are written to the 'flashpointvulnerabilitydetails' incident field.
    :rtype: dict | CommandResults

    :raises ValueError: If the vulnerability list command fails.
    """
    rows = get_vulnerability_rows()

    vulnerability_ids = [row.get("vuln_id", "") for row in rows if row.get("vuln_id")]

    if not vulnerability_ids:
        return CommandResults(readable_output=MESSAGES["NO_RECORDS_FOUND"])

    result = get_vulnerabilities(vulnerability_ids)
    cve_ids_by_vulnerability = get_cve_ids_by_vulnerability(result)

    grid_rows = build_grid_rows(rows, cve_ids_by_vulnerability)
    execute_command_safe(SET_INCIDENT_COMMAND, {VULNERABILITY_DETAILS_FIELD: grid_rows})

    return result


""" MAIN FUNCTION """


def main():
    """
    Entry point. Executes get_vulnerability_details and returns results.
    Catches all exceptions and surfaces them via return_error.
    """
    try:
        return_results(get_vulnerability_details())
    except Exception as ex:
        demisto.error(traceback.format_exc())
        return_error(f"Failed to execute FlashpointVulnerabilityDetails. Error: {ex!s}")


""" ENTRY POINT """


if __name__ in ("__main__", "__builtin__", "builtins"):
    main()