import pymongo
import types
import json
import sys
import os

PORT=27017
client = pymongo.MongoClient("localhost", PORT)
db = client.cuckoo_db
collection = db.malware_results

commands = {}
COLLECT_NAMES = "collect_names"
COUNT_TRUE = "count_true"
COLLECT_ANSWERS = "collect_answers"
COLLECT_AND_COUNT_ANSWERS = "collect_and_count_answers"
COLLECT_JS = "collect_js"

def execute_commands():
    for js in collection.find():
        for command_name in commands.keys():
            command = commands[command_name]
            response = command["fun"](js)
            if command["do"] == COLLECT_NAMES:
                if response:
                    command["ans"].append(js["malware_name"])
            elif command["do"] == COUNT_TRUE:
                if response:
                    command["ans"] += 1
            elif command["do"] == COLLECT_ANSWERS:
                if type(response) is types.GeneratorType or type(response) is list:
                    for element in response:
                        if element is not None:
                            command["ans"].append(element)
                elif response is not None:
                    command["ans"].append(response)
            elif command["do"] == COLLECT_AND_COUNT_ANSWERS:
                if type(response) is types.GeneratorType or type(response) is list:
                    for element in response:
                        if element is not None:
                            command["ans"][element] = command["ans"].get(element, 0) + 1
                elif response is not None:
                    command["ans"][response] = command["ans"].get(response, 0) + 1
            elif command["do"] == COLLECT_JS:
                if response:
                    command["ans"].append(js)
                
def add_command(name, fun, response_type, display=True):
    commands[name] = {}
    if response_type == COUNT_TRUE:
        commands[name]["ans"] = 0
    elif response_type in [COLLECT_ANSWERS, COLLECT_NAMES, COLLECT_JS]:
        commands[name]["ans"] = []
    elif response_type == COLLECT_AND_COUNT_ANSWERS:
        commands[name]["ans"] = {}
    else:
        commands.pop(name, None)
        return False
    commands[name]["fun"] = fun
    commands[name]["do"] = response_type
    commands[name]["display"] = display
    return True

def get_result(command_name):
    if commands.get(command_name, None) is not None:
        return commands[command_name].get("ans", None)

def print_results():
    for command_name in commands.keys():
        if commands[command_name]["display"]:
            print(command_name, commands[command_name]["ans"])

def get_malware_from_name(name):
    return collection.find_one({ "malware_name": name })
