ACTI Vulnerability Query

import demistomock as demisto
from CommonServerPython import *  # noqa # pylint: disable=unused-wildcard-import
import urllib3
import traceback

# Disable insecure warnings
urllib3.disable_warnings()  # pylint: disable=no-member


""" CONSTANTS """
DATE_FORMAT = "%Y-%m-%d %H:%M:%S"
ENDPOINTS = {"vulnerability": "/rest/vulnerability"}


""" CLIENT CLASS """


class Client(BaseClient):
    def __init__(self, input_url: str, api_key: str, verify_certificate: bool, proxy: bool, endpoint="/rest/vulnerability"):
        base_url = urljoin(input_url, endpoint)
        PACK_VERSION = get_pack_version()
        DEMISTO_VERSION = demisto.demistoVersion()
        DEMISTO_VERSION = f'{DEMISTO_VERSION["version"]}.{DEMISTO_VERSION["buildNumber"]}'
        headers = {
            "Content-Type": "application/json",
            "auth-token": api_key,
            "User-Agent": f"AccentureCTI Pack/{PACK_VERSION} Palo Alto XSOAR/{DEMISTO_VERSION}",
        }
        super().__init__(base_url=base_url, verify=verify_certificate, headers=headers, proxy=proxy)

    def vulnerability_search(self, url_suffix: str, data: dict = {}) -> dict:
        return self._http_request(method="GET", url_suffix=url_suffix, params=data)


class CustomCVE(Common.CVE):
    def __init__(self, id, cvss3, cvss2, published, modified, description, dbot, relationships=None):
        # type (str, str, str, str, str, str, str) -> None
        super().__init__(
            id=id, cvss="", published=published, modified=modified, description=description, relationships=relationships
        )
        self.cvss3 = cvss3
        self.cvss2 = cvss2
        self.dbot_score = dbot

    def to_context(self):
        context_values = super().to_context()
        context_values[self.CONTEXT_PATH]["CVSS2"] = self.cvss2
        context_values[self.CONTEXT_PATH]["CVSS3"] = self.cvss3
        context_values[self.CONTEXT_PATH]["CVSS"]["CVSS2"] = self.cvss2
        context_values[self.CONTEXT_PATH]["CVSS"]["CVSS3"] = self.cvss3
        if self.dbot_score:
            context_values.update(self.dbot_score.to_context())
        return context_values


""" HELPER FUNCTIONS """


def _calculate_dbot_score(severity: int) -> int:
    """
    Calculates Dbot score according to table:
    Dbot Score   | severity
     0           | 0
     1           | 1,2
     2           | 3,4
     3           | 5,6,7
    Args:
        severity: value from 1 to 5, determined by iDefense threat indicator

    Returns:
        Calculated score
    """
    dbot_score = Common.DBotScore.NONE

    if severity > 4:
        dbot_score = Common.DBotScore.BAD
    elif severity > 2:
        dbot_score = Common.DBotScore.SUSPICIOUS
    elif severity > 0:
        dbot_score = Common.DBotScore.GOOD

    return dbot_score


def _extract_result(res: dict, dbot_score_type: str, reliability: DBotScoreReliability) -> list[dict]:
    analysis_results = []
    if res.get("total_size"):
        results_array = res.get("results", [])
        if len(results_array):
            for result in results_array:
                cpes = []
                indicator_value = result.get("key", "")
                dbot_score: int = _calculate_dbot_score(result.get("severity", 0))
                desc = "Match found in AccentureCTI database"
                dbot = Common.DBotScore(indicator_value, dbot_score_type, "ACTIVulnerabilityQuery", dbot_score, desc, reliability)

                vuln_techs = result.get("affects", "").get("vuln_techs", "")
                for vuln in vuln_techs:
                    cpes.append(vuln.get("display_text", ""))

                last_published = result.get("last_published", "")
                if last_published:
                    last_published_format = parse_date_string(last_published, DATE_FORMAT)
                else:
                    last_published_format = "NA"

                last_modified = result.get("last_modified", "")
                if last_modified:
                    last_modified_format = parse_date_string(last_modified, DATE_FORMAT)
                else:
                    last_modified_format = "NA"
                display_text = result.get("display_text", "NA")
                analysis_info = {
                    "Name": result.get("key", "NA"),
                    "DbotReputation": dbot_score,
                    "CVSS2": result.get("cvss2_temporal_score", "NA"),
                    "CVSS3": result.get("cvss3_temporal_score", "NA"),
                    "UUID": result.get("uuid", "NA"),
                    "CPEs": cpes,
                    "LastPublished": str(last_published_format),
                    "LastModified": str(last_modified_format),
                    "Description": result.get("description", ""),
                }
                analysis_results.append({"analysis_info": analysis_info, "dbot": dbot, "display_text": display_text})

    return analysis_results


def _returned_response(res: dict) -> list[str]:
    """
    Checks which indicator value found on AccentureCTI database.
    Args:
        res: api response

    Returns: list of indicator values that returned from api request

    """
    returned_res = []
    if res.get("total_size"):
        results_array = res.get("results", [])
        if len(results_array):
            for result_content in results_array:
                returned_res.append(result_content.get("key", ""))
    return returned_res


def _match_not_found(all_inputs: list, res: list) -> list[str]:
    """

    Args:
        all_inputs: all indicator values received from the user
        res: list of all indicator values that returned from api request

    Returns: Which indicator has no match on AccentureCTI database

    """
    value_not_found = []

    for val in all_inputs:
        if val not in res:
            value_not_found.append(val)

    return value_not_found


""" COMMAND FUNCTIONS """


def test_module(client: Client) -> str:  # type: ignore
    """Tests API connectivity and authentication'

    Returning 'ok' indicates that the integration works like it is supposed to.
    Connection to the service is successful.
    Raises exceptions if something goes wrong.

    :type client: ``Client``
    :param Client: client to use

    :return: 'ok' if test passed, anything else will fail the test.
    :rtype: ``str``
    """

    try:
        client.vulnerability_search(url_suffix="/v0")
        return "ok"
    except Exception as e:
        if "Error in API call [403]" in e.args[0]:
            return_results(f"This API token doesn't have permission for accessing vulnerability API!.\n Error: {str(e)}")
            demisto.debug(e.args[0])
        else:
            raise DemistoException(f"Error in API call - check the input parameters and the API Key. Error: {e}.")


def vuln_command(client: Client, args: dict, reliability: DBotScoreReliability) -> list[CommandResults]:
    vulns: list = argToList(args.get("cve", None))
    if not vulns:
        raise ValueError("vuln not specified")

    result = client.vulnerability_search(url_suffix="/v0", data={"key.values": vulns})
    analysis_results = _extract_result(result, DBotScoreType.CVE, reliability)
    returned_cves = _returned_response(result)
    cves_not_found = _match_not_found(vulns, returned_cves)

    command_results = []
    for analysis_result in analysis_results:
        analysis_info: dict = analysis_result.get("analysis_info", {})
        dbot = analysis_result.get("dbot")
        cve_name = analysis_info.get("Name", "")
        display_name = analysis_result.get("display_text", "")
        readable_output = tableToMarkdown(f"{display_name}", analysis_info)
        indicator = CustomCVE(
            cve_name,
            analysis_info.get("CVSS3"),
            analysis_info.get("CVSS2"),
            analysis_info.get("LastPublished"),
            analysis_info.get("LastModified"),
            analysis_info.get("Description"),
            dbot,
        )  # noqa: E501
        command_results.append(CommandResults(indicator=indicator, raw_response=result, readable_output=readable_output))  # type: ignore  # noqa: E501

    for cve in cves_not_found:
        desc = "No results were found on AccentureCTI database"
        dbot = Common.DBotScore(cve, DBotScoreType.CVE, "ACTIVulnerabilityQuery", 0, desc)
        indicator = CustomCVE(cve, None, None, None, None, None, dbot)
        readable_output = f"No results were found for cve {cve}"
        command_results.append(CommandResults(indicator=indicator, readable_output=readable_output))  # type: ignore  # noqa: E501

    return command_results


""" MAIN FUNCTION """


def main():
    params = demisto.params()
    api_key = params.get("api_token")
    if isinstance(api_key, dict):
        api_key = api_key.get("password")

    base_url = urljoin(params.get("url", ""))
    reliability = params.get("integrationReliability", "B - Usually reliable")

    if DBotScoreReliability.is_valid_type(reliability):
        reliability = DBotScoreReliability.get_dbot_score_reliability_from_str(reliability)
    else:
        Exception("Accenture CTI error: Please provide a valid value for the Source Reliability parameter")

    commands = {
        "acti-vuln": vuln_command,
    }

    verify_certificate = not params.get("insecure", False)
    proxy = params.get("use_proxy", False)

    try:
        command = demisto.command()
        client = Client(base_url, api_key, verify_certificate, proxy, endpoint=ENDPOINTS["vulnerability"])
        demisto.debug(f"Command being called is {command}")

        if command == "test-module":
            return_results(test_module(client))
        elif command in commands:
            return_results(commands[command](client, demisto.args(), reliability))

    # Log exceptions and return errors   # noqa: E303
    except Exception as e:
        if "Error in API call [403]" in e.args[0]:
            return_error(f"This API token doesn't have permission for accessing vulnerability API!.\n Error: {str(e)}")
        else:
            demisto.error(traceback.format_exc())
            return_error(f"Failed to execute {demisto.command()} command.\nError:\n{str(e)}")


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