from intelligence.base import *
import requests
import json

TAGS = {
        "": DetectionLevel.MALWARE,
        "ripped": DetectionLevel.MALWARE,
        "yara": DetectionLevel.SUSPECT,
        "et": DetectionLevel.SUSPECT,
}


class MWDB(OnlineChecker):

    name = "MWDB"
    options = {
        "url": ("https://mwdb.cert.pl", "Endpoint url"),
        "login": ("", "login"),
        "password": ("", "password"),
    }

    def __init__(self, **options):
        OnlineChecker.__init__(self, **options)
        self.url = self.options.get("url", "https://mwdb.cert.pl").strip()
        if self.url.endswith("/"):
            self.url = self.url[:-1]


    def login(self):
        login = self.options.get("login")
        if not login:
            raise KeyError("No username provided")
        password = self.options.get("password")
        if not password:
            raise KeyError("No password provided")
        session = self.preconfigured_session()  # automatically set options such as timeout or ssl_verify
        session.headers.update({'accept': 'application/json'})
        session.headers.update({'content-type': 'application/json'})
        auth = session.post(self.url + "/api/auth/login", json={"login": login, "password": password})
        auth.raise_for_status()
        token = auth.json()["token"]
        session.headers.update({"Authorization": f"Bearer {token}"})
        return session

    
    def check(self, analysis):
        detections = { }
        session = self.login()
        finalurl = '/'.join([self.url, 'api', 'file', analysis.entropy.sha256])
        response = session.get(finalurl)
        if response.ok:
            data = response.json()
            if not data:
                return None
            if "message" in data:
                raise ValueError(message)
            tags = []
            det = OnlineDetection()
            for tag in data.get("tags") or []:
                tag = tag.get("tag", "")
                if not tag:
                    continue
                if ":" in tag:
                    prefix = tag.split(":")[0]
                    tag = ":".join(tag.split(":")[1:])
                else:
                    prefix = ""
                lvl = TAGS.get(prefix)
                if lvl is not None:
                    tags.append(tag)
                    det.level = max(det.level, lvl)
            det.name = " + ".join(tags)
            if data.get("last_config"):
                det.name += " [CONFIG]"
            det.name += " [{:d} children] [{:d} parents]".format(len(data.get("children", [])), len(data.get("parents", [])))
            
            file_name = data.get("file_name", analysis.entropy.sha256)
            detections[file_name] = det
        elif response.status_code == 404:
            return None
        else:
            raise response.raise_for_status()

        return OnlineResult(detections=detections, 
            url = "{}/file/{}/relations".format(self.url, analysis.entropy.sha256)
        )

    def download(self, query):
        detections = { }
        if len(query) != 64:
            return None
        session = self.login()
        finalurl = '/'.join([self.url, 'api', 'file', query, "download"])
        response = session.get(finalurl)
        if response.ok:
            return OnlineFile(
                url = "{}/file/{}/relations".format(self.url, query),
                content=response.content
            )
        elif response.status_code == 404:
            return None
        else:
            raise response.raise_for_status()
