from intelligence.base import *
import requests
import json


PUBLIC_URL = "https://tria.ge/api/v0/"


class RecordedFutureTriage(OnlineChecker):

    name = "Triage"
    options = {
        "url": ("https://tria.ge/api/v0/", "Endpoint URL (default to public one)"),
        "key": ("", "Triage API key"),
    }

    def send(self, query):
        key = self.options.get("key")
        if not key:
            raise KeyError("No API key")
        url = self.options.get("url", PUBLIC_URL)
        if not url.endswith("/"):
            url += "/"
        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'})
        session.headers.update({'Authorization': f"Bearer {key}"})
        return session.get(f"{url}{query}")

    def find_sample(self, h):
        h = h.strip()
        if len(h) == 32:
            query = f"md5:{h}"
        elif len(h) == 40:
            query = f"sha1:{h}"
        elif len(h) == 64:
            query = f"sha256:{h}"
        else:
            raise ValueError(f"Invalid hash format: {h}")
        r = self.send(f"search?query={query}")
        if r.status_code == 404:
            return []
        else:
            r.raise_for_status()
        data = r.json().get("data", [])
        return [x["id"] for x in data]
    
    def check(self, analysis):
        detections = { }
        samples_ids = self.find_sample(analysis.entropy.sha256)
        for sample_id in samples_ids:
            overview = self.send(f"samples/{sample_id}/overview.json")
            overview.raise_for_status()
            overview = overview.json()
            filename = overview["sample"]["target"]
            sha256 = overview["sample"]["sha256"]
            if sha256.lower() != analysis.entropy.sha256:
                continue
            extracted = overview.get("extracted", {})
            a = overview.get("analysis", {})
            has_config = "config" in extracted
            has_dumps = "dumped_file" in extracted
            
            det = OnlineDetection()
            det.level = DetectionLevel.from_number(a.get("score", 0), 10)
            det.name = " + ".join(a.get("family", []))
            if has_config:
                det.name += " (config extracted)"
            if has_dumps:
                det.name += " (has dumps)"

            detections[f"{sample_id}:{filename}"] = det
        if not detections:
            return None
       
        return OnlineResult(
            detections = detections, 
            url = f"https://tria.ge/s?q={analysis.entropy.sha256}"
        )

    def download(self, query):
        samples_ids = self.find_sample(query)
        for sample_id in samples_ids:
            data = self.send(f"samples/{sample_id}/sample")
            if data.ok:
                return OnlineFile(url=f"https://tria.ge/samples/{sample_id}/sample.zip", content=data.content)
        return None

