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

from typing import Dict, Any, List
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 = dict()
    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')

        # 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})

            # 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
                    fields['cvssvector'] = cvss_v3.get('vectorString')

            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)
            )
        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()
