from intelligence.base import *
import requests
import json

URL = "https://www.virustotal.com/api/v3/"



class VirusTotal(OnlineChecker):

    name = "VirusTotal"
    options = {
        "key": ("", "VirusTotal private or public API key. Note that you will need a private API key to perform vtgrep lookups, but public is fine for hash lookups."),
    }
    
    def check(self, analysis):
        detections = { }
        key = self.options.get("key")
        if not key:
            raise KeyError("No API key")
        message = ""
        session = self.preconfigured_session()  # automatically set options such as timeout or ssl_verify
        session.headers.update({'x-apikey': key})
        response = session.get(URL + "files/{}".format(analysis.entropy.sha256))
        if response.status_code == 404:
            return None
        elif response.status_code == 401:
            return KeyError("Bad API key")
        elif not response.ok:
            response.raise_for_status()
            return None
        data = response.json()
        infos = data.get("data", {}).get("attributes", {})
        for av, result in infos.get("last_analysis_results", {}).items():
            category = result.get("category", "clean")
            if category == "timeout":
                continue
            name = result.get("result", "")
            if name is None:
                name = ""
            detections[av] = OnlineDetection(level=DetectionLevel.from_text(category), name=name.strip())
        return OnlineResult(detections=detections, 
            url="https://www.virustotal.com/gui/file/{}/detection".format(analysis.entropy.sha256), 
            message=message
        )

    
