import json
import logging
from typing import Optional

import requests
from backoff import expo, on_exception
from ratelimit import RateLimitException, limits

from ioccheck.services import Service

logger = logging.getLogger(__name__)

f_handler = logging.FileHandler("ioccheck.log")
f_handler.setLevel(logging.INFO)

f_format = logging.Formatter("%(asctime)s - %(name)s - %(levelname)s - %(message)s")
f_handler.setFormatter(f_format)

logger.addHandler(f_handler)


class MalwareBazaar(Service):
    name = "malwarebazaar"

    def __init__(self, file_hash: str, api_key: str):
        self.url = "https://mb-api.abuse.ch/api/v1/"
        self.response = self._get_api_response(file_hash, api_key)

        if self.response.get("query_status") == "hash_not_found":
            return

        self.hashes = self._get_other_hashes(self.response)

        self.investigation_url = self._get_investigation_url()
        self.is_malicious = None  # self._get_is_malicious(self.response)
        self.signature = self._get_signature(self.response)

        self.tags = self._get_tags(self.response)

        self.file_size = self._get_file_size(self.response)
        self.file_type = self._get_file_type(self.response)
        self.mime_type = self._get_mime_type(self.response)

    @on_exception(expo, RateLimitException, max_tries=10)
    @limits(calls=15, period=60)
    def _get_api_response(self, file_hash: str, api_key: str) -> dict:
        data = {"query": "get_info", "hash": file_hash}

        r = requests.post(self.url, data=data, headers={"API_KEY": api_key})

        return json.loads(r.text)

    def _get_investigation_url(self) -> str:
        # MalwareBazaar indexes URLs by their SHA256 hash
        investigation_url = ""

        sha256_hash = self._get_other_hashes(self.response)

        if sha256_hash.get("SHA256") is not None:
            investigation_url = (
                f"https://bazaar.abuse.ch/sample/{sha256_hash.get('SHA256')}/"
            )

        return investigation_url

    def _get_other_hashes(self, response) -> dict:
        hash_keys = [
            "sha256_hash",
            "sha1_hash",
            "md5_hash",
            "sha3_384_hash",
            "imphash",
            "tlsh",
            "telfhash",
        ]

        return dict(
            (key, value)
            for key, value in response.get("data")[0].items()
            if key in hash_keys
        )

    def _get_signature(self, response) -> Optional[str]:
        return response.get("signature")

    def _get_tags(self, response) -> Optional[list]:
        return response.get("data")[0].get("tags")

    def _get_file_size(self, response) -> Optional[int]:
        return response.get("data")[0].get("file_size")

    def _get_file_type(self, response) -> Optional[str]:
        return response.get("data")[0].get("file_type")

    def _get_mime_type(self, response) -> Optional[str]:
        return response.get("data")[0].get("file_type_mime")
