#!/usr/bin/python3

import matplotlib.pyplot as plt

from package_db import *

def get_virus_total_result(js_total):
    return js_total.get("virustotal", None)

def virus_total_succeeded(js_total):
    return (get_virus_total_result(js_total) is not None) and \
           (js_total["virustotal"].get("positives", None) is not None)

def virus_total_what(js_total):
    vt = get_virus_total_result(js_total)
    if vt is not None and vt.get("positives", None) is None:
        return vt

def get_virus_total_score(js_total):
    vt = get_virus_total_result(js_total)
    if vt is not None:
        return vt.get("positives", None)

def get_info_signature(js_total):
    for signature in js_total.get("signatures", []): 
        yield {"name": signature["name"],
               "severity": signature["severity"],
               "description": signature["description"],
               "markcount": signature["markcount"]}

def get_from_signature(parameter):
    def fun(js_total):
        for signature in js_total.get("signatures", []):
            yield signature.get(parameter, None)
    return fun
        
def has_malware_signature_with_severity(threshold=0):
    get_severity = get_from_signature("severity")
    def fun(js_total): 
        for severity in get_severity(js_total):
            if severity > threshold:
                return True
        return False
    return fun

def has_malware_no_signature_with_severity(threshold=0):
    get_severity = get_from_signature("severity")
    def fun(js_total): 
        for severity in get_severity(js_total):
            if severity > threshold:
                return False
        return True
    return fun

def has_error_in_debug(js_total):
    return len(js_total["debug"]["errors"]) > 0

def has_error_and_zero_score(js_total):
    if js_total["info"].get("score", -1) == 0:
        return has_error_in_debug(js_total)
    return False

def errors_from_zero_score(js_total):
    if has_error_and_zero_score(js_total):
        for l in js_total["debug"]["errors"]:
            yield l

def has_malware_score_min(threshold=0):
    return lambda js: get_score(js) >= threshold

def is_malware_score_between(mini, maxi):
    return lambda js: mini <= get_score(js) < maxi

def get_score(js_total):
    return js_total["info"].get("score", None)

def get_score_if(f):
    def fun(js_total):
        if f(js_total):
            return get_score(js_total)
        return None
    return fun

def is_pattern_in_a_sig(pattern):
    get_all = get_from_signature("name")
    def fun(js_total):
        for name in get_all(js_total):
            if pattern in name:
                return True
        return False
    return fun

def count_pattern_in_a_sig(pattern):
    get_all = get_from_signature("name")
    def fun(js_total):
        c = 0
        for name in get_all(js_total):
            if pattern in name:
                c += 1
        return c
    return fun

if __name__=="__main__":
    add_command("total malware in db = ", lambda js_total: True, COUNT_TRUE)
    add_command("antivm attempt for malwares = ", is_pattern_in_a_sig("antivm"), COUNT_TRUE)
    add_command("antisandbox attempt for malwares = ", is_pattern_in_a_sig("antisandbox"), COUNT_TRUE)
    for threshold in range(11):
        add_command(f"malware score higher or equal {threshold} = ", has_malware_score_min(threshold), COUNT_TRUE)
    add_command(f"malware score = 0 = ", is_malware_score_between(0, 0.0001), COUNT_TRUE)
    add_command(f"malware score = 0 and has error = ", has_error_and_zero_score, COUNT_TRUE)
    for mini in range(10):
        add_command(f"{mini} <= malware score < {mini+1} = ", is_malware_score_between(mini, mini+1), COUNT_TRUE)

    for threshold in range(5):
        add_command(f"malware with at least a sig with severety higher than {threshold} = ", has_malware_signature_with_severity(threshold), COUNT_TRUE)

    add_command(f"antivm device", is_pattern_in_a_sig("antivm_shared_device"), COLLECT_NAMES)
    add_command(f"antivbox device", is_pattern_in_a_sig("antivm_vbox_devices"), COLLECT_NAMES)
    add_command(f"antivbox window", is_pattern_in_a_sig("antivm_vbox_window"), COLLECT_NAMES)
    add_command(f"antivm vmware keys", is_pattern_in_a_sig("antivm_vmware_keys"), COLLECT_NAMES)

    add_command(f"score_malware_sig_5", get_score_if(has_malware_signature_with_severity(4)), COLLECT_AND_COUNT_ANSWERS, display=False)
    add_command(f"score_malware_no_sig_5", get_score_if(has_malware_no_signature_with_severity(4)), COLLECT_AND_COUNT_ANSWERS, display=False)

    add_command(f"virus total succeeded = ", virus_total_succeeded, COUNT_TRUE)
    add_command(f"virus total resource not scanned = ", lambda js: virus_total_what(js) is not None, COUNT_TRUE)

    #add_command("signature names detected with count = ", get_from_signature("name"), COLLECT_AND_COUNT_ANSWERS)
    add_command("all_scores", get_score, COLLECT_AND_COUNT_ANSWERS, display=False)
    add_command("virus_total_scores", get_virus_total_score, COLLECT_AND_COUNT_ANSWERS, display=False)
    #add_command("names malwares not detected = ", lambda x: not has_malware_signature_with_severity(0)(x), COLLECT_NAMES)
    
    #add_command("antivm attempt for malwares = ", lambda js_total: 5 < count_pattern_in_a_sig("antivm")(js_total), COLLECT_NAMES)
    
    #add_command("errors from 0 score = ", errors_from_zero_score, COLLECT_ANSWERS, display=False)
    
    add_command("names scores more 20 = ", lambda js_total: get_score(js_total)>=20, COLLECT_NAMES)
    add_command("number scores more 20 = ", lambda js_total: get_score(js_total)>=20, COUNT_TRUE)


    execute_commands()
    print_results()
    
    #errors_zero = get_result("errors from 0 score = ")
    #for l in errors_zero:
    #    print(l)
    
    #exit()
    
    for s in ["all_scores", "virus_total_scores"]:
        scores = get_result(s)
        keys = scores.keys()
        values = [scores[k] for k in keys]

        density = []
        for k in keys:
            density.extend([k for _ in range(scores[k])])

        plt.figure()
        plt.boxplot(density, vert=0, showfliers=False)
        plt.scatter(keys, values)
        plt.yscale('log')
        plt.ylabel('number of malware')
        plt.xlabel('score')
        plt.title(s)


    plt.figure()
    t = ["score_malware_sig_5", "score_malware_no_sig_5"]
    scores = get_result(t[0]), get_result(t[1])
    keys = scores[0].keys(), scores[1].keys()
    values = [scores[0][k] for k in keys[0]], [scores[1][k] for k in keys[1]]

    #density = [], []
    #for k in keys[0]:
    #    density[0].extend([k for _ in range(scores[k])])
    #for k in keys[1]
    #    density[1].extend([k for _ in range(scores[k])])
    #plt.boxplot(density, vert=0)
    plt.scatter(keys[0], values[0], color='blue')
    plt.scatter(keys[1], values[1], color='red')
    plt.yscale('log')
    plt.ylabel('number of malware')
    plt.xlabel('score')
    plt.title(s)

    plt.show()

