National Vulnerability Database

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

from typing import Any
from datetime import datetime, timedelta
from time import sleep
import urllib3


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

""" CONSTANTS """

DATE_FORMAT = "%Y-%m-%dT%H:%M:%S:000 UTC-00:00"  # ISO8601 format with UTC, default in XSOAR
DEFAULT_LIMIT = 100


# =========================================== Helper Functions ===========================================#


def parse_cve_data(item=None) -> dict:
    fields = {}
    if item and isinstance(item, dict):
        # Populate common fields
        cve: dict = item.get("cve")  # type: ignore

        # Description
        if "description" in cve:
            description = cve.get("description", {}).get("description_data")
            if "en" in [x.get("lang") for x in description]:
                fields["cvedescription"] = description[([x.get("lang") for x in description]).index("en")].get("value")
                # mirror the description fields so that the default cve layout renders with complete data
                fields["description"] = fields["cvedescription"]

        # References
        if "references" in cve:
            references: list = cve.get("references", {}).get("reference_data")
            fields["publications"] = [
                {"link": x.get("url"), "title": x.get("name"), "source": x.get("refsource")} for x in references
            ]

        # Parse impact data
        impact = item.get("impact")
        if impact:
            # All CVSS data
            all_cvss = []

            # Base Metric V2
            if "impact" in impact:
                impact_v2 = impact.get("impact")
                if "baseMetricV2" in impact_v2:
                    base_metric_v2 = impact_v2.get("baseMetricV2")
                    if "cvssV2" in base_metric_v2:
                        cvss_v2 = base_metric_v2.get("cvssV2")
                        all_cvss.append({"cvssV2": cvss_v2})
                        fields["cvss"] = all_cvss

            # Base Metric V3
            if "baseMetricV3" in impact:
                base_metric = impact.get("baseMetricV3")
                if "cvssV3" in base_metric:
                    cvss_v3 = base_metric.get("cvssV3")
                    all_cvss.append({"cvssV3": cvss_v3})
                    cvss_v3_data = []
                    for k, v in cvss_v3.items():
                        cvss_v3_data.append({"metric": camel_case_to_underscore(k).replace("_", " ").title(), "values": v})
                    fields["cvss3"] = cvss_v3_data

                    cvss_data = []
                    for k, v in cvss_v3.items():
                        cvss_data.append({"metrics": camel_case_to_underscore(k).replace("_", " ").title(), "value": v})
                    fields["cvsstable"] = cvss_data
                    # additional fields to write to CVE default layout
                    fields["cvssvector"] = cvss_v3.get("vectorString")
                    # mirror fields so that default cve layout renders properly
                    fields["cvss"] = cvss_v3.get("baseScore")
                    fields["cvssscore"] = cvss_v3.get("baseScore")
                    fields["cvssversion"] = cvss_v3.get("version")

            # this eneded up being a json blob which should be fixed to the baseScore value above
            # fields['cvss'] = all_cvss

    return fields


def extract_titles(data_item={}) -> list:
    titles = []
    for title in data_item.get("titles"):
        titles.append(title.get("title"))
    return titles


def extract_descriptions(data_item={}) -> list:
    descriptions = []
    for description in data_item.get("cve", {}).get("description", {}).get("description_data"):
        descriptions.append(description.get("value"))
    return descriptions


# ========================================== Generic Query ===============================================#


def test_module(client: BaseClient, params: dict[str, Any]):
    api_key = params.get("apiKey")
    try:
        params = {"cveId": "CVE-2020-22120"}
        if api_key:
            params["apiKey"] = api_key
        res = client._http_request("GET", full_url="https://services.nvd.nist.gov/rest/json/cpes/1.0", params=params)
        if "error" in res:
            return_error(res.get("error"))
        elif "resultsPerPage" in res:
            return_results("ok")
    except Exception as err:
        raise DemistoException(err)


def fetch_indicators_command(client, params):
    command = demisto.command()
    api_key = params.get("apiKey")
    get_type = params.get("type")
    cpe_match_string = params.get("cpeMatchString")
    cpe_keyword = params.get("keyword")
    include_deprecated = params.get("deprecated")
    cvss_v2_metrics = params.get("cvssV2Metrics")
    cvss_v2_severity = params.get("cvssV2Severity")
    cvss_v3_metrics = params.get("cvssV3Metrics")
    cvss_v3_severity = params.get("cvssV3Severity")
    history = int(params.get("history"))
    exceeds_span = True
    urls = {"CPE": "/rest/json/cpes/1.0/", "CVE": "/rest/json/cves/1.0/"}
    url = urls[get_type]
    now = datetime.utcnow()
    startIndex = 0
    resultsPerPage = 2000
    data_items = []
    indicators: list[dict] = []
    last_run_data = demisto.getLastRun()
    run_times: list[datetime] = []
    run_limit = 9

    # If there is no last run date, use the history specified in the params
    if "lastRun" not in last_run_data or command == "nvd-get-indicators":
        last_run = now - timedelta(days=history)
    else:
        last_run = dateparser.parse(last_run_data.get("lastRun"))  # type: ignore

    modStartDate = last_run
    modEndDate = now

    # API calls can only span 120 days, so we should loop if the history
    # parameter is greater than this
    while exceeds_span and modEndDate and modStartDate:
        delta = (modEndDate - modStartDate).days
        if delta > 120:
            modEndDate = modStartDate + timedelta(days=120)
        else:
            exceeds_span = False
        params = {
            "modStartDate": modStartDate.strftime(DATE_FORMAT),
            "modEndDate": modEndDate.strftime(DATE_FORMAT),
            "startIndex": startIndex,
            "resultsPerPage": resultsPerPage,
        }
        if api_key:
            params["apiKey"] = api_key
            run_limit = 99
        if get_type == "CPE":
            params["addOns"] = "cves"
            if include_deprecated:
                params["includeDeprecated"] = include_deprecated
        if get_type == "CVE":
            if cvss_v2_metrics:
                params["cvssV2Metrics"] = cvss_v2_metrics
            if cvss_v2_severity:
                params["cvssV2Severity"] = cvss_v2_severity
            if cvss_v3_metrics:
                params["cvssV3Metrics"] = cvss_v3_metrics
            if cvss_v3_severity:
                params["cvssV3Severity"] = cvss_v3_severity
        if cpe_match_string:
            params["cpeMatchString"] = cpe_match_string
        if cpe_keyword:
            params["keyword"] = cpe_keyword

        total_results = 1
        collection_count = 0

        # Collect all the indicators together
        while collection_count < total_results:
            # Check to ensure no rate limits are hit
            if len(run_times) == run_limit:
                first_time = run_times[0]
                last_time = run_times[(run_limit - 1)]

                if (last_time - first_time).seconds <= 60:
                    demisto.info("Rate limit hit, sleeping for 3 seconds")
                    # We sleep 3 seconds to avoid hitting any rate limits
                    sleep(3)
                del run_times[0]

            run_times.append(datetime.utcnow())

            res = client._http_request("GET", url, params=params, timeout=300)

            # Check to see if there are any errors
            if "error" in res:
                return_error(res.get("error"))
            total_results = res.get("totalResults", 0)
            resultsPerPage = res.get("resultsPerPage", 0)
            result = res.get("result")

            if result:
                if get_type == "CPE":
                    data_items += result.get("cpes")
                else:
                    data_items += result.get("CVE_Items")

            params["startIndex"] += resultsPerPage
            collection_count += resultsPerPage
        modStartDate = modEndDate
        modEndDate = now

    # If this is nvd-get-indicators command:
    if command == "nvd-get-indicators":
        # If they are CPEs
        if get_type == "CPE":
            outputs = [
                {
                    "cpe23Uri": x.get("cpe23Uri"),
                    "titles": ". ".join(extract_titles(data_item=x)),
                    "vulnerabilities": ", ".join(x.get("vulnerabilities")),
                }
                for x in data_items
            ]
            command_results = CommandResults(
                outputs_prefix="CPE",
                outputs_key_field="cpe23Uri",
                outputs=data_items,
                readable_output=tableToMarkdown("National Vulnerability Database CPEs:", outputs),
            )
        # If they are CVEs
        elif get_type == "CVE":
            outputs = [
                {"id": x.get("cve").get("CVE_data_meta").get("ID"), "description": ". ".join(extract_descriptions(data_item=x))}
                for x in data_items
            ]
            command_results = CommandResults(
                outputs_prefix="CVE",
                outputs_key_field="id",
                outputs=data_items,
                readable_output=tableToMarkdown("National Vulnerability Database CVEs:", outputs),
            )
        else:
            command_results = CommandResults("No result was found")
        return_results(command_results)

    # Else if this is fetch-indicators
    elif command == "fetch-indicators":
        indicators = []

        # If they are CPEs
        if get_type == "CPE" and data_items:
            for item in data_items:
                item["type"] = "CPE"
                indicator = {"value": item.get("cpe23Uri"), "rawJSON": item}

                # This is reserved for future use
                if "vulnerabilities" in item:
                    relationships = []
                    for vulnerability in item.get("vulnerabilities", []):
                        relationship = EntityRelationship(
                            name=EntityRelationship.Relationships.RELATED_TO,
                            entity_a=item.get("cpe23Uri"),
                            entity_a_family="Indicator",
                            entity_a_type="CPE",
                            entity_b=vulnerability,
                            entity_b_family="Indicator",
                            entity_b_type="CVE",
                        )
                        relationships.append(relationship.to_indicator())
                    indicator["relationships"] = relationships
                indicators.append(indicator)

        # If they are CVEs
        elif get_type == "CVE" and data_items:
            for item in data_items:
                item["type"] = "CVE"
                fields: dict = parse_cve_data(item)
                indicators.append(
                    {
                        "value": item.get("cve", {}).get("CVE_data_meta", {}).get("ID"),
                        "type": FeedIndicatorType.CVE,
                        "fields": fields,
                        "rawJSON": item,
                    }
                )

        # Create the indicators in a batch, 2000 at a time
        for b in batch(indicators, batch_size=2000):
            demisto.createIndicators(b)

    # Set new integration context
    demisto.setLastRun({"lastRun": now.isoformat()})


# =========================================== Built-In Queries ===========================================#


""" MAIN FUNCTION """

# COMMAND CONSTANTS

commands = {
    "test-module": test_module,
    "fetch-indicators": fetch_indicators_command,
    "nvd-get-indicators": fetch_indicators_command,
}


def main() -> None:
    params = demisto.params()
    base_url = "https://services.nvd.nist.gov"
    verify_cert = not params.get("insecure", False)
    proxy = params.get("proxy", False)
    command = demisto.command()
    demisto.debug(f"Command being called is {command}")
    try:
        client = BaseClient(
            base_url=base_url,
            verify=verify_cert,
            proxy=proxy,
        )
        commands[command](client, params)

    # Log exceptions and return errors
    except Exception as e:
        demisto.error(traceback.format_exc())  # print the traceback
        return_error(f"Failed to execute {command} command.\nError: {str(e)}")


""" ENTRY POINT """

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