#! /usr/bin/env python
import json
from datetime import datetime

from harpoon.commands.base import Command
from harpoon.lib.koodous import Koodous, KoodousError, KoodousNotFound


class CommandKoodous(Command):
    """
    # Koodous plugin

    Queries the Koodous API https://koodous.com/

    * get info on a hash : `harpoon koodous hash SHA256`
    * Download a file : `harpoon koodous dl SHA256`
    """

    name = "koodous"
    description = "Request Koodous API"
    config = {"Koodous": ["token"]}

    def add_arguments(self, parser):
        subparsers = parser.add_subparsers(help="Subcommand")
        parser_a = subparsers.add_parser("hash", help="Get info on a SHA256 hash")
        parser_a.add_argument("HASH", help="SHA256 hash")
        parser_a.set_defaults(subcommand="hash")
        parser_b = subparsers.add_parser("search", help="Search in Koodous")
        parser_b.add_argument("QUERY", help="Query")
        parser_b.set_defaults(subcommand="search")
        parser_c = subparsers.add_parser("dl", help="Download a sample from Koodous")
        parser_c.add_argument("HASH", help="Sha256")
        parser_c.set_defaults(subcommand="dl")
        parser_d = subparsers.add_parser(
            "analysis", help="Get a full analysis from Koodous"
        )
        parser_d.add_argument("HASH", help="Sha256")
        parser_d.set_defaults(subcommand="analysis")
        parser_e = subparsers.add_parser(
            "account", help="Get information on the account"
        )
        parser_e.set_defaults(subcommand="account")
        self.parser = parser

    def run(self, args, plugins):
        kd = Koodous(token=self._config_data["Koodous"]["token"])
        if "subcommand" in args:
            try:
                if args.subcommand == "hash":
                    res = kd.sha256(args.HASH)
                    print(json.dumps(res, indent=4))
                elif args.subcommand == "search":
                    res = kd.search(args.QUERY)
                    print(json.dumps(res, indent=4))
                elif args.subcommand == "dl":
                    data = kd.download(args.HASH)
                    with open(args.HASH, "wb+") as f:
                        f.write(data)
                    print("File downloaded as {}".format(args.HASH))
                elif args.subcommand == "analysis":
                    res = kd.analysis(args.HASH)
                    print(json.dumps(res, indent=4))
                elif args.subcommand == "account":
                    res = kd.account()
                    print(json.dumps(res, indent=4))
                else:
                    self.parser.print_help()
            except KoodousNotFound:
                print("Not found")
        else:
            self.parser.print_help()

    def intel(self, type, query, data):
        if type == "hash":
            if len(query) == 64:
                try:
                    kd = Koodous(token=self._config_data["Koodous"]["token"])
                    res = kd.sha256(query)
                except KoodousError:
                    pass
                else:
                    data["samples"].append(
                        {
                            "source": "Koodous",
                            "date": datetime.fromtimestamp(res["created_on"]),
                            "url": "https://koodous.com/apks/" + query,
                        }
                    )
