from intelligence.base import *
import requests
import json

URL = "https://mb-api.abuse.ch/api/v1/"

VENDOR_MAPPINGS = {
        "ANY.RUN": ("tags", "verdict", None),
        "CERT-PL_MWDB": ("detection", None, None),
        "YOROI_YOMI": ("detection", "score", 1.0),
        "vxCube": (None, "verdict", None),
        "CAPE": ("detection", None, None),
        "Triage": ("malware_family", "score", 10),
        "DocGuard": ("verdict", "alertlevel", 10.0),
        "VMRay": ("malware_family", "verdict", None),
        "FileScan-IO": (None, "verdict", None),
        "InQuest": (None, "verdict", None),
        "Spamhaus_HBL": (None, "detection", None),
        "Intezer": ("family_name", "verdict", None),
        "ReversingLabs": ("threat_name", "status", None),
        "FileScan-IO": ("verdict", "threatlevel", 1.0),
}


class MalwareBazaar(OnlineChecker):

    name = "MalwareBazaar"
    options = {
        "key": ("", "MalwareBazaar API key"),
    }

    def __session(self):
        session = self.preconfigured_session()  # automatically set options such as timeout or ssl_verify
        session.headers.update({'accept': 'application/json'})
        key = self.options.get("key")
        if not key:
            raise KeyError("No API key")
        session.headers.update({'Auth-Key': key})
        session.headers.update({'User-Agent': 'malcat'})
        return session
    
    def check(self, analysis):
        detections = { }
        message = ""
        session = self.__session()
        response = session.post(URL, data = {
                "query": "get_info",
                "hash": analysis.entropy.sha256
        })
        if not response.ok:
            raise response.raise_for_status()
        data = response.json()
        status = data.get("query_status", "")
        if  status == "hash_not_found":
            return None
        elif status and status != "ok":
            raise ValueError("API returned: {}".format(status))
        for data_item in data.get("data", []):
            #print(json.dumps(data_item, indent=4))
            for vendor, infos in data_item.get("vendor_intel", {}).items():
                if not vendor in VENDOR_MAPPINGS or vendor.lower() in detections:
                    continue
                if infos and type(infos) == list:
                    infos = infos[0]
                if not infos:
                    continue
                det = OnlineDetection()
                sig_field, level_field, level_max = VENDOR_MAPPINGS[vendor]
                level = infos.get(level_field, None)
                if level_max is not None and level is not None:
                    level = type(level_max)(level)
                sig = infos.get(sig_field, None)
                #print(vendor, level, type(level), sig)
                if type(level) == str:
                    det.level = DetectionLevel.from_text(level)
                elif level_max is not None:
                    det.level = DetectionLevel.from_number(level, level_max)
                if sig:
                    if level_field is None:
                        det.level = DetectionLevel.MALWARE
                    if type(sig) == list:
                        sig = sig[0]
                    if sig:
                        det.name = sig
                detections[vendor.lower()] = det
        return OnlineResult(detections=detections, 
            url="https://bazaar.abuse.ch/sample/{}/".format(analysis.entropy.sha256), 
            message=message
        )

    def download(self, query):
        if len(query) != 64:
            # malware bazaar only supports sha256 downloads
            return None
        session = self.__session()
        response = session.post(URL, data = {
                "query": "get_file",
                "sha256_hash": query
        })
        if not response.ok:
            raise response.raise_for_status()
        if not response.content.startswith(b"PK"):
            data = response.json()
            status = data.get("query_status", "")
            if status == "file_not_found":
                return None
            else:
                raise ValueError("API returned: {}".format(status))
        return OnlineFile(
            url="https://bazaar.abuse.ch/sample/{}/".format(query), 
            content=response.content
        )
