"""
VirusTotal Resolver
===================

Resolver leveraging the VirusTotal JSON API.

This is a third-party data source that provides GeoIP, ASN, associated detected file samples,
detections on web hosted content, known resolutions, and other resolution information for IPv4 and IPv6 addresses.

Data Source Information
-----------------------

This data source is hosted at virustotal.com and requires an internet connection
to use. It offers a free API with registration and approval, with terms described
on their site. Please refer to the service's website for an authoritative source
on API specifications.

This documentation summarizes a point in time understanding of the data source
though since it is a third party service, it may change in a manner that breaks
this tool or causes this documentation to become inaccurate. In no way is
inclusion of a data source in libchickadee an endorsement of the data source.

**Data source documentation:** https://developers.virustotal.com/reference

Endpoints
^^^^^^^^^

The API supports a number of reports, though for IP addresses only allows
the submission of a single IP address per request. This means bulk lookups
will take at most 1 minute per 4 IP addresses, due to rate limiting.

* ``https://www.virustotal.com/vtapi/v2/ip-address/report?apikey=<apikey>&ip=<ip>``

Fields
^^^^^^

These fields are in no particular order.

* query
* count
* asn
* country
* subnet
* resolution_count
* detected_sample_count
* detected_samples
* undetected_sample_count
* undetected_samples
* detected_url_count
* detected_urls
* undetected_url_count
* undetected_urls
* status
* message

Limitations
^^^^^^^^^^^

This service has a free tier for non-commercial use, and is rate limited to:

* 4 requests per minute
* 5760 requests per day
* 172800 requests per month

A 204 response code indicates that you have exceeded your rate limit. The script will
sleep for 15 seconds if a 204 is returned.

Module Documentation
--------------------

"""
import collections
import functools
import logging
import operator
import time
from datetime import datetime

import requests
from tqdm import tqdm

from . import ResolverBase

logger = logging.getLogger(__name__)

__author__ = "Chapin Bryce"
__date__ = 20200805
__license__ = "MIT Copyright 2020 Chapin Bryce"
__desc__ = "Resolver for VirusTotal"

FIELDS = [
    "query",
    "count",
    "asn",
    "country",
    "subnet",
    "resolution_count",
    "detected_sample_count",
    "undetected_sample_count",
    "detected_url_count",
    "undetected_url_count",
    "status",
    "message",
]

NON_DEFAULT_FIELDS = [
    "resolutions",
    "detected_samples",
    "undetected_samples",
    "detected_urls",
    "undetected_urls",
]


class ProResolver(ResolverBase):
    """VirusTotal authenticated API resolver. Requires an API key.
    Full API documentation here: https://developers.virustotal.com/reference"""

    def __init__(self, api_key, fields=None, lang="en"):
        """Initialize class object and configure default values.

        Args:
            api_key (str): API key to authenticate with
            fields (list): Fields to filter responses
            lang (str): Language to use in the API request
        """
        super().__init__()

        self.supported_langs = ["en"]
        self.lang = "en" if lang not in self.supported_langs else lang
        self.fields = FIELDS if not fields else fields
        self.uri = "https://www.virustotal.com/vtapi/v2/ip-address/report"
        self.api_key = api_key
        self.enable_sleep = True
        self.last_request = datetime.now()
        logger.debug("API key found")

    def sleeper(self):
        """Method to sleep operations for rate limiting.

        Will ensure that 4 requests per minute limit is not exceeded if a 204 encountered.

        Return:
            None
        """
        current_request = datetime.now()
        time_since_last_request = current_request - self.last_request
        if time_since_last_request.total_seconds() > 15:
            return
        time_to_sleep = (
            15 - time_since_last_request.total_seconds()
        ) + 0.25  # Add padding
        logger.info(f"Sleeping for {time_to_sleep} seconds due to rate limiting.")
        time.sleep(time_to_sleep)

    def batch(self):
        """Resolve multiple IP addresses.

        Due to API limitations, each IP address must be resolved individually.

        Returns:
            (list): List of resolved results
        """
        # Not supported
        all_ips = self.data
        records = []
        if self.pbar:
            all_ips = tqdm(self.data)
        for x in all_ips:
            self.data = x
            if resp := self.single():
                records.append(resp[0])
        return records

    def single(self):
        """Gathers VirusTotal report data for a single IP address.

        Returns:
            (list): Report information
        """
        params = {"apikey": self.api_key, "ip": self.data}

        self.last_request = datetime.now()
        rdata = requests.get(
            self.uri,
            params=params,
            timeout=60,
        )

        if rdata.status_code == 200:
            return [self.parse_vt_resp(self.data, rdata.json())]
        if rdata.status_code == 204:
            # Rate limit
            self.sleeper()
            # Try again
            return self.single()
        elif rdata.status_code == 400:
            logger.error("Incorrect request. Please check input data")
        elif rdata.status_code == 403:
            logger.error("Authorization error. Please check API key")
        else:
            logger.error(
                f"Unknown error occurred, status code {rdata.status_code}, please report"
            )

    def parse_vt_resp(self, query, vt_resp):
        """Transform the raw response from VirusTotal in to a dictionary easier for analysis

        Args:
            query (str): Queried value
            vt_resp (dict): Raw API response

        Returns:
            (dict): Formatted response for analysis
        """
        attributes = dict.fromkeys(self.fields.copy())
        attributes.update(
            {
                "query": query,
                "asn": "",
                "continent": vt_resp.get("continent"),
                "country": vt_resp.get("country"),
                "subnet": vt_resp.get("network"),
                "status": vt_resp.get("response_code"),
                "message": vt_resp.get("IP address in dataset"),
                "whois": "",
            }
        )

        # ASNs
        if vt_resp.get("asn"):
            attributes["asn"] = f'AS{vt_resp.get("asn")} {vt_resp.get("as_owner")}'

        # Parse WhoIs
        self._extract_whois(attributes, vt_resp)

        # Parse Resolutions
        self._extract_resolutions(attributes, vt_resp)

        # Parse Detected Communicating, Download, and Referrer Samples
        self._extract_detected_samples(attributes, vt_resp)

        # Parse Undetected Communicating Samples
        self._extract_undetected_samples(attributes, vt_resp)

        # Parse Detected URLs
        self._extract_detected_urls(attributes, vt_resp)

        # Parse Undetected URLs
        self._extract_undetected_urls(attributes, vt_resp)

        return attributes

    @staticmethod
    def _extract_resolutions(attributes, vt_resp):
        """Extract information about associated domain resolutions

        Args:
            attributes (dict): Formatted data to present for analysis
            vt_resp (dict): Raw information from the VT API

        Returns:
            None. Updates attributes dictionary
        """
        attributes["resolution_count"] = len(vt_resp.get("resolutions", []))
        hostname_set = {x.get("hostname") for x in vt_resp.get("resolutions", [])}
        attributes["resolutions"] = sorted(list(hostname_set))

    def _extract_undetected_urls(self, attributes, vt_resp):
        """Extract information about undetected URLs

        Args:
            attributes (dict): Formatted data to present for analysis
            vt_resp (dict): Raw information from the VT API

        Returns:
            None. Updates attributes dictionary
        """
        # * Get count
        # * Get all defanged URLs
        # The API does not expose this as a key/value object
        # but instead as a list where the elements are:
        # * 0: URL
        # * 1: SHA256 Hash value
        # * 2: Number of positive detections
        # * 3: Number of scanners
        # * 4: Time of last scan
        attributes["undetected_url_count"] = len(vt_resp.get("undetected_urls", []))
        detected_urls = {
            self.defang_ioc(x[0]) for x in vt_resp.get("undetected_urls", [])
        }
        attributes["undetected_urls"] = sorted(list(detected_urls))

    def _extract_detected_urls(self, attributes, vt_resp):
        """Extract information about detected URLs

        Args:
            attributes (dict): Formatted data to present for analysis
            vt_resp (dict): Raw information from the VT API

        Returns:
            None. Updates attributes dictionary
        """
        # * Get count
        # * Get all defanged URLs
        attributes["detected_url_count"] = len(vt_resp.get("detected_urls", []))
        detected_urls = {
            self.defang_ioc(x.get("url")): x.get("positives")
            for x in vt_resp.get("detected_urls", [])
        }
        attributes["detected_urls"] = detected_urls

    @staticmethod
    def _extract_undetected_samples(attributes, vt_resp):
        """Extract information about undetected samples

        Args:
            attributes (dict): Formatted data to present for analysis
            vt_resp (dict): Raw information from the VT API

        Returns:
            None. Updates attributes dictionary
        """
        # * Get count
        # * Get all hashes
        attributes["undetected_sample_count"] = (
            len(vt_resp.get("undetected_communicating_samples", []))
            + len(vt_resp.get("undetected_downloaded_samples", []))
            + len(vt_resp.get("undetected_referrer_samples", []))
        )

        undetected_samples = {
            x.get("sha256") for x in vt_resp.get("undetected_communicating_samples", [])
        }
        undetected_samples = undetected_samples.union(
            {x.get("sha256") for x in vt_resp.get("undetected_downloaded_samples", [])}
        )
        undetected_samples = undetected_samples.union(
            {x.get("sha256") for x in vt_resp.get("undetected_referrer_samples", [])}
        )
        attributes["undetected_samples"] = sorted(list(undetected_samples))

    @staticmethod
    def _extract_detected_samples(attributes, vt_resp):
        """Extract information about detected samples

        Args:
            attributes (dict): Formatted data to present for analysis
            vt_resp (dict): Raw information from the VT API

        Returns:
            None. Updates attributes dictionary
        """
        # * Get count
        # * Get all hashes
        attributes["detected_sample_count"] = (
            len(vt_resp.get("detected_communicating_samples", []))
            + len(vt_resp.get("detected_downloaded_samples", []))
            + len(vt_resp.get("detected_referrer_samples", []))
        )

        detected_communicating_samples = {
            x.get("sha256"): x.get("positives")
            for x in vt_resp.get("detected_communicating_samples", [])
        }
        detected_downloaded_samples = {
            x.get("sha256"): x.get("positives")
            for x in vt_resp.get("detected_downloaded_samples", [])
        }
        detected_referrer_samples = {
            x.get("sha256"): x.get("positives")
            for x in vt_resp.get("detected_referrer_samples", [])
        }
        # Sum up the counts across categories for the same samples
        attributes["detected_samples"] = dict(
            functools.reduce(
                operator.add,
                map(
                    collections.Counter,
                    [
                        detected_communicating_samples,
                        detected_downloaded_samples,
                        detected_referrer_samples,
                    ],
                ),
            )
        )

    @staticmethod
    def _extract_whois(attributes, vt_resp):
        """Extract whois data from the VirusTotal API response

        Args:
            attributes (dict): Formatted data to present for analysis
            vt_resp (dict): Raw information from the VT API

        Returns:
            None. Updates attributes dictionary
        """
        # OtherRemarks holds info not stored in key/value config
        whois = {"OtherRemarks": ""}
        for line in vt_resp.get("whois", "").split("\n"):
            if len(line) == 0:
                continue
            if ":" not in line:
                whois["OtherRemarks"] += line + "\n"
                continue
            split_line = line.split(":", 1)
            whois[split_line[0].strip()] = split_line[1].strip()
        attributes["whois"] = whois
