from CommonServerPython import *
from collections import defaultdict
from datetime import date


def header_transform(key: str) -> str:
    """
    Function returns header key in human readable

    :param key: header key
    :return: translated headers
    """
    header_dict = {
        'Cve': 'CVE ID',
        'CVSS': 'CVSS Score',
        'VRR': 'VRR Score',
        'ThreatCount': 'Threat Count',
        'VulnLastTrendingOn': 'Last Trending On Date',
        'Trending': 'Trending',
    }

    return header_dict.get(key, '')


def get_ransomware_trending_cves(ransomware_cves) -> list:
    """
    Filter last 7 days trending reansomware cves from ransomware cves.

    :param ransomware_cves: list of ransomware cves.
    :return: list of trending ransomware.
    """
    ransomware_trending = []

    today_date = date.today()
    week_ago_date = today_date - timedelta(days=7)
    demisto.setContext('Date.CurrentDate', str(today_date))
    demisto.setContext('Date.WeekAgoDate', str(week_ago_date))

    for cve in ransomware_cves:
        if cve.get('VulnLastTrendingOn') and cve.get('VulnLastTrendingOn') != 'Not Found':
            vuln_last_trending_on_date = datetime.strptime(cve.get('VulnLastTrendingOn', '2000-01-01'),
                                                           '%Y-%m-%d').date()
            if (today_date - vuln_last_trending_on_date).days <= 7:
                ransomware_trending.append(cve)

    return ransomware_trending


def get_ransomware_cves(host_findings):
    """
    Retrieves ransomware cves from host findings details.

    :param host_findings: list of host findings details
    :return: list of ransomware cves
    """
    all_cves: dict = defaultdict(dict)

    for findings in host_findings:
        risk_rating = findings.get('RiskRating', 0)
        threat_list = defaultdict(list)

        threats = findings.get('Threat') if findings.get('Threat', []) else []
        vulns = findings.get('Vulnerability') if findings.get('Vulnerability', []) else []

        # if dict response receive from context data then convert it to list
        threats = [threats] if isinstance(threats, dict) else threats
        vulns = [vulns] if isinstance(vulns, dict) else vulns

        for threat in threats:
            for cve in threat.get('Cve', []):
                threat_list[cve].append({
                    'Title': threat.get('Title', ''),
                    'Category': threat.get('Category', ''),
                    'Severity': threat.get('Severity', ''),
                    'Description': threat.get('Description', ''),
                    'Cve': threat.get('Cve', ''),
                    'Source': threat.get('Source', ''),
                    'Published': threat.get('Published', ''),
                    'Updated': threat.get('Updated', ''),
                    'ThreatLastTrendingOn': threat.get('ThreatLastTrendingOn', ''),
                    'Trending': threat.get('Trending', '')
                })

        for vulnerability in vulns:
            last_trending = vulnerability.get('VulnLastTrendingOn', '')
            all_cves[vulnerability.get('Cve', 'No ID Found')] = {
                'Cve': vulnerability.get('Cve', ''),
                'CVSS': vulnerability.get('BaseScore', ''),
                'VRR': risk_rating,
                'ThreatCount': len(threat_list.get(vulnerability.get('Cve', ''), [])),
                'Trending': vulnerability.get('Trending', ''),
                'VulnLastTrendingOn': last_trending if last_trending else 'Not Found',
                'Threats': threat_list.get(vulnerability.get('Cve', ''), []),
                'Description': vulnerability.get('Description', '')
            }
    cve_list = []
    for cve in all_cves.values():
        if isinstance(cve, list):
            cve_list.extend(cve)
        else:
            cve_list.append(cve)
    ransomware_cves = []
    for cve in cve_list:
        for threat in cve.get('Threats', []):
            if threat.get('Category', '').lower() == 'ransomware':
                ransomware_cves.append(cve)
                break
    return ransomware_cves


def display_ransomware_cve_results(ransomware_cves):
    """
    Set context data and human readable for ransomware cves.
    will raise an error message if no ransomware cves found.

    :param ransomware_cves: list of ransomware cves
    :return: standard output
    """
    ransomware_cves_count = len(ransomware_cves)
    if ransomware_cves_count <= 0:
        return_results('No ransomware CVEs found.')

    demisto.setContext('CVECount', ransomware_cves_count)
    readable_output = '### Total CVEs found: ' + str(len(ransomware_cves)) + '\n'
    readable_output += tableToMarkdown('List of CVEs that have ransomware threat', ransomware_cves,
                                       ['Cve', 'CVSS', 'VRR', 'ThreatCount', 'VulnLastTrendingOn', 'Trending'],
                                       header_transform, removeNull=True)
    return CommandResults(
        outputs_prefix='RiskSense.RansomwareCves',
        outputs_key_field='Cve',
        outputs=ransomware_cves,
        readable_output=readable_output,
        raw_response={}
    )


def display_ransomware_trending_cve_results(ransomware_cves):
    """
    Retrieve trending ransomware cves and Set context data and human readable for trending ransomware cves.
    will raise an error message if no trending ransomware cves found.

    :param ransomware_cves: list of ransomware cves.
    :return: standard output.
    """
    ransomware_trending = get_ransomware_trending_cves(ransomware_cves)
    ransomware_trending_cves_count = len(ransomware_trending)
    if ransomware_trending_cves_count <= 0:
        return_results('No CVEs found With trending ransomware.')

    demisto.setContext('TrendingCVECount', ransomware_trending_cves_count)
    readable_output = '### Total CVEs found: ' + str(len(ransomware_trending)) + '\n'
    readable_output += tableToMarkdown('List of CVEs that are ransomware trending', ransomware_trending,
                                       ['Cve', 'CVSS', 'VRR', 'ThreatCount', 'VulnLastTrendingOn', 'Trending'],
                                       header_transform, removeNull=True)
    return CommandResults(
        outputs_prefix='RiskSense.RansomwareTrendingCves',
        outputs_key_field='Cve',
        outputs=ransomware_trending,
        readable_output=readable_output,
        raw_response={}
    )


def main() -> None:
    LOG("Start RiskSense get ransomeware cves script")

    # Fetch Parameter
    trending = demisto.args().get('trending')

    results = demisto.executeCommand('risksense-get-host-findings', {
        'fieldname': 'threat_categories',
        'operator': 'EXACT',
        'value': 'Ransomware',
        'size': 1000,
        'status': 'Open'
    })

    host_findings = results[0].get('EntryContext', {}).get('RiskSense.HostFinding(val.ID == obj.ID)')

    if not host_findings:
        return_results('Key [RiskSense.HostFinding] was not found in context data.')

    host_findings = [host_findings] if isinstance(host_findings, dict) else host_findings

    ransomware_cves = get_ransomware_cves(host_findings)

    if argToBoolean(trending):
        result = display_ransomware_trending_cve_results(ransomware_cves)
        return_results(result)
    else:
        result = display_ransomware_cve_results(ransomware_cves)
        return_results(result)


if __name__ == "__builtin__" or __name__ == "builtins":
    main()
