#!/ur/bin/env python3

import requests


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

    def getIPVirusTotal(self, ip, report):
        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'

        report.append(data)

    def getRateFromHash(self, h, report):
        headers = self._headers
        headers['resource'] = h

        res = requests.get(
            f"{self._url}/file/report",
            headers=headers
        ).json()

        data = dict()