#!/ur/bin/env python3

import requests
from config import VT_ATTRIBUTES_MAPPING
from datetime import datetime


class VT:
    def __init__(self, api_key):
        self._url = "https://www.virustotal.com/api/v3"
        self._headers = {
            'x-apikey': api_key,
        }

    def getIPVirusTotal(self, ip):
        """
        This function get IP information from VirusTotal
        """
        res = requests.get(
            f"{self._url}/ip_addresses/{ip}",
            headers=self._headers
        ).json()

        data = dict()
        data['ip'] = ip

        if 'error' in res:
            report.append({
                'error': res['error']['message'],
                'ip': ip
            })
            return

        vt = res['data']['attributes']
        for entry in VT_ATTRIBUTES_MAPPING.keys():
            if entry in vt:
                try:
                    data[entry] = vt[entry]
                except KeyError:
                    data[entry] = 'Unknown'
        return data

    def getDomainReport(self, fqdn, report):
        """
        This function get the report for the specific domain 
        """
        res = requests.get(
            f"{self._url}/domains/{fqdn}",
            headers=self._headers
        )

        js = res.json()
        report['reputation'] = js['data']['attributes']['reputation']
        report['last_update'] = \
            datetime.fromtimestamp(js['data']['attributes']['last_update_date'])

        # Get number of security vendors
        report['total_vendors'] = 0
        report['clean'] = 0
        report['unrated'] = 0
        report['malicious'] = 0
        vendors = js['data']['attributes']['last_analysis_results']

        for entry in vendors:
            report['total_vendors'] += 1
            if vendors[entry]['result'] == 'clean':
                report['clean'] += 1
            elif vendors[entry]['result'] == 'unrated':
                report['unrated'] += 1
            elif vendors[entry]['result'] == 'malicious':
                report['malicious'] += 1

    def getRateFromHash(self, h, report):
        """
        This function get the report of the hash specified by the parameter h
        """
        headers = self._headers

        res = requests.get(
            f"{self._url}/files/{h}",
            headers=self._headers
        ).json()

        if 'error' in res:
            report["error"] = "Can not find the result"
            return

        attributes = res['data']['attributes']
        report['results'] = dict()
        report['results']['file'] = dict()
        report['results']['file']['magic'] = attributes['magic']
        report['results']['file']['sha1'] = attributes['sha1']
        report['results']['file']['md5'] = attributes['md5']
        report['results']['file']['filetype'] = attributes['detectiteasy']['filetype']
        report['results']['file']['size'] = attributes['size']
        report['results']['file']['extension'] = attributes['type_extension']
        try:
            report['results']['file']['first_seen'] = \
                datetime.fromtimestamp(attributes['first_seen_itw_date'])
        except KeyError:
            report['results']['file']['first_seen'] = "Unknown"
            
        report['results']['file']['last_analysis'] = \
            datetime.fromtimestamp(attributes['last_analysis_date'])

        # Identify vendors 
        report['results']['vendors'] = dict()
        report['results']['vendors']['total_vendors'] = 0
        report['results']['vendors']['total_malicious'] = 0
        report['results']['vendors']['result'] = dict()
        report['results']['vendors']['result']['undetected'] = 0
        results = res['data']['attributes']['last_analysis_results']
        for entry in results:
            report['results']['vendors']['total_vendors'] += 1
            if results[entry]['category'] == 'undetected':
                report['results']['vendors']['result']['undetected'] += 1
            else:
                result = results[entry]['result']
                if result not in report['results']['vendors']['result']:
                    report['results']['vendors']['result'][result] = 0
                report['results']['vendors']['result'][result] += 1
                report['results']['vendors']['total_malicious'] += 1