import logging
from datetime import datetime, timedelta
from io import StringIO
from typing import ClassVar

import pandas as pd

from core import taskmanager
from core.schemas import task
from core.schemas.observables import (
    file,
    imphash,
    md5,
    sha1,
    sha256,
    ssdeep,
    tlsh,
)


class AbuseCHMalwareBazaaar(task.FeedTask):
    _defaults = {
        "frequency": timedelta(hours=1),
        "name": "AbuseCHMalwareBazaaar",
        "description": "This feed contains md5/sha1/sha256",
    }

    _SOURCE: ClassVar["str"] = "https://bazaar.abuse.ch/export/csv/recent/"
    _NAMES = [
        "first_seen_utc",
        "sha256_hash",
        "md5_hash",
        "sha1_hash",
        "reporter",
        "file_name",
        "file_type_guess",
        "mime_type",
        "signature",
        "clamav",
        "vtpercent",
        "imphash",
        "ssdeep",
        "tlsh",
    ]

    def run(self):
        data = {"query": "get_recent", "selector": "time"}

        response = self._make_request(self._SOURCE, auth=None, verify=True)
        if response:
            data = StringIO(response.text)

            df = pd.read_csv(
                data,
                comment="#",
                delimiter=",",
                names=self._NAMES,
                quotechar='"',
                quoting=1,
                skipinitialspace=True,
                parse_dates=["first_seen_utc"],
            )

            df = self._filter_observables_by_time(df, "first_seen_utc")
            df = df.fillna("")
            for _, line in df.iterrows():
                self.analyze(line)

            return True

    def analyze(self, block):
        if not block.get("sha256_hash"):
            return

        context = {}

        context["first_seen"] = block["first_seen_utc"]
        context["source"] = self.name
        context["filename"] = block["file_name"]
        context["md5"] = block["md5_hash"]
        context["sha1"] = block["sha1_hash"]
        context["sha256"] = block["sha256_hash"]
        context["imphash"] = block["imphash"]
        context["ssdeep"] = block["ssdeep"]
        context["date_added"] = datetime.utcnow()
        context["tlsh"] = block["tlsh"]
        if block["reporter"]:
            context["reporter"] = block["reporter"]
        else:
            context["reporter"] = "N/A"

        tags = []

        file_value = "FILE:{}".format(block["sha256_hash"])

        malware_file = file.File(value=file_value).save()
        logging.debug(
            f"Malware file: {malware_file} source: {self.name} context: {context}"
        )
        malware_file.name = context["filename"]
        if block["signature"]:
            tags.append(block["signature"])
        malware_file.add_context(self.name, context, skip_compare={"date_added"})
        logging.debug(block)
        tags.append(block["mime_type"])

        if "tags" in block:
            tags.extend(block["tags"].split(","))

        malware_file.tag(tags)

        # md5, sha1, sha256, imphash, ssdeep, tlsh

        sha256_data = sha256.SHA256(value=context["sha256"]).save()
        sha256_data.tag(tags)

        malware_file.link_to(sha256_data, "sha256", self.name)

        sha1_data = sha1.SHA1(value=context["sha1"]).save()
        sha1_data.tag(tags)

        malware_file.link_to(sha1_data, "sha1", self.name)

        md5_data = md5.MD5(value=context["md5"]).save()
        md5_data.tag(tags)

        malware_file.link_to(md5_data, "md5", self.name)

        if context["imphash"]:
            imphash_data = imphash.Imphash(value=context["imphash"]).save()
            imphash_data.tag(tags)
            malware_file.link_to(imphash_data, "imphash", self.name)

        ssdeep_data = ssdeep.Ssdeep(value=context["ssdeep"]).save()
        ssdeep_data.tag(tags)
        malware_file.link_to(ssdeep_data, "ssdeep", self.name)

        tls_data = tlsh.TLSH(value=context["tlsh"]).save()
        tls_data.tag(tags)
        malware_file.link_to(tls_data, "tlsh", self.name)


taskmanager.TaskManager.register_task(AbuseCHMalwareBazaaar)
