import os
os.environ.setdefault('DJANGO_SETTINGS_MODULE', 'malware_toolkit_project.settings')
import re

import django
django.setup()

from malware_toolkit.models import Feature, Function, Body, \
    Function_Global_Variables, Function_Import, Header, Definition, \
    Feature_Function, Toolkit_Option, Feature_Type, Feature_Import, \
    Feature_Global_Variables
from django.core.exceptions import ObjectDoesNotExist

def add_new_feature(feat):
    print(feat['name'])
    print(feat['body'])
    f = Feature.objects.get_or_create(name=feat['name'])[0]
    f.priority = feat['priority']
    f.feature_type = feat['type']
    f.toolkit_prompt = feat['prompt']
    add_update_table(Body, "feature", f, feat['body'])
    add_update_table(Feature_Global_Variables, "feature", f, feat['globals'])
    add_update_table(Feature_Import, "feature", f, feat['import'])
    add_update_table(Toolkit_Option, "feature", f, feat['options'])
    link_functions_to_features(f, feat['functions'])
    f.save()

# TODO - add exception handling
def add_new_function(func):
    print(func['name'])
    print(func['globals'])
    print(func['import'])
    print(func['declaration'])
    print(func['definition'])
    f = Function.objects.get_or_create(name=func['name'])[0]
    add_update_table(Function_Global_Variables, "function", f, func['globals'])
    add_update_table(Function_Import, "function", f, func['import'])
    add_update_table(Header, "function", f, func['declaration'])
    add_update_table(Definition, "function", f, func['definition'])

# func - string
# deletes a function by name (string)
def delete_function(func):
    f = Function.objects.get_or_create(name=func)[0]
    f.delete()

# table - models.table name
# rName - name of the record to be deleted (string)
# deletes a feature by name (string)
def delete_record_by_name(table, rName):
    r = table.objects.get_or_create(name=rName)[0]
    r.delete()

def add_update_table(table, fKeyType, fKey, g_lines):
    line_counter = 0
    empty = True
    g = None
    if fKeyType == "function":
        g = table.objects.all().filter(function=fKey)
    elif fKeyType == "feature":
        g = table.objects.all().filter(feature=fKey)
    # if we found records, we need to remove excess rows
    if g:
        # we will need to remember if there were records in the table
        empty = False
        end = len(g)
        while end > len(g_lines):
            end -= 1
            g[end].delete()

    for l in g_lines:
        # only if we found recods in the table with len(g) succeed
        if not empty:
            if line_counter < len(g):
                g[line_counter].text = l
                g[line_counter].save()
                line_counter += 1
                continue

        if table is Toolkit_Option:
            table.objects.get_or_create(feature=fKey, \
                text=l, number=line_counter)[0].save()
        elif fKeyType == "function":
            table.objects.get_or_create(function=fKey, \
                text=l, line_number=line_counter)[0].save()
        elif fKeyType == "feature":
            table.objects.get_or_create(feature=fKey, \
                text=l, line_number=line_counter)[0].save()
        line_counter += 1

# adds records to the aggregate table linking functions to features
#TODO this needs to be tested against the server database since my local copy
# is not working correctly
def link_functions_to_features(feat, functions):
    rec_counter = 0
    empty = True
    records = Feature_Function.objects.all().filter(feature=feat)
    # if we found records, we need to remove excess rows
    if records:
        # we will need to remember if there were records in the table
        empty = False
        end = len(records)
        while end > len(functions):
            end -= 1
            records[end].delete()

    func = None
    for f in functions:
        # only if we found recods in the table with len(records) succeed
        func = Function.objects.get(name=f)
        if not empty:
            if rec_counter < len(records):
                if func:
                    records[rec_counter].function = func
                    records[rec_counter].save()
                else:
                    continue
                rec_counter += 1
                continue

        Feature_Function.objects.get_or_create(feature=feat, \
            function=func)[0].save()
        rec_counter += 1


# name - String name of the function or feautre
# fKeyType - the foreign key type, either 'function' or 'feature' (string)
# returns a QuerySet of models.'Table' object tied to the given function
# returns None if nothing was found or error encountered
def get_records_from_table(table, fKeyType, name):
    records = None
    try:
        # use function_id to query for the import objects
        if fKeyType == "function":
            # get the function id (foreign key)
            # based on the function name passed in
            fKey = Function.objects.get(name=name)
            records = table.objects.all().filter(function=fKey)
        elif fKeyType == 'feature':
            # get the feature id (foreign key)
            # based on the feature name passed in
            fKey = Feature.objects.get(name=name)
            records = table.objects.all().filter(feature=fKey)
    except ObjectDoesNotExist as e:
        print("bad news")
        # print(fKeyType +  " '" + name + "' does not exist")
    return records

def get_feature_function_names(feat):
    func_names = []
    funcs = get_records_from_table(Feature_Function, 'feature', feat)
    for f in funcs:
        print(f.function)
        func_names.append(Function.objects.get(function=f.function))
    return func_names

def get_python_feature_obj(feat_name):
    try:
        feat = Feature.objects.get(name=feat_name)
    except Feature.DoesNotExist:
        return None
    body = get_records_from_table(Body, 'feature', feat)
    parsed_b = make_array(body)
    parsed_b = replaceBadChars(parsed_b)
    global_vars = get_records_from_table(Feature_Global_Variables, 'feature', feat_name)
    imports = get_records_from_table(Feature_Import, 'feature', feat_name)
    try:
        feat_type = Feature_Type.objects.get(feature=feat).name
    except Feature_Type.DoesNotExist:
        feat_type = ""
    options = get_records_from_table(Toolkit_Option, 'feature', feat)
    parsed_opts = make_array(options)
    funcs = get_records_from_table(Feature_Function, 'feature', feat)
    func_names = []
    for f in funcs:
        func_names.append(f.function.name)
    featObj = {}
    featObj['name'] = feat_name
    featObj['globals'] = replaceBadChars(make_array(global_vars))
    featObj['imports'] = replaceBadChars(make_array(imports))
    featObj['priority'] = feat.priority
    featObj['type'] = feat_type
    featObj['body'] = parsed_b
    featObj['functions'] = func_names
    featObj['prompt'] = feat.toolkit_prompt
    featObj['options'] = parsed_opts

    return featObj

def get_python_function_obj(func_name):
    global_vars = get_records_from_table(Function_Global_Variables, 'function', func_name)
    imports = get_records_from_table(Function_Import, 'function', func_name)
    header = get_records_from_table(Header, 'function', func_name)
    definition = get_records_from_table(Definition, 'function', func_name)

    funcObj = {}
    funcObj['name'] = func_name
    funcObj['globals'] = replaceBadChars(make_array(global_vars))
    funcObj['imports'] = replaceBadChars(make_array(imports))
    funcObj['header'] = replaceBadChars(make_array(header))
    func_def = make_array(definition)
    funcObj['definition'] = replaceBadChars(func_def)

    return funcObj

def make_array(qSet):
    array = []
    if(qSet != None):
        for s in qSet:
            array.append(s.text)
    return array

def replaceBadChars(code_list):
    i = 0
    while i < len(code_list):
        code_list[i] = re.sub(r'\{\{1\}\}', '\"', code_list[i])
        code_list[i] = re.sub(r'\{\{2\}\}', '\'', code_list[i])
        code_list[i] = re.sub(r'\{\{3\}\}', '<', code_list[i])
        code_list[i] = re.sub(r'\{\{4\}\}', '>', code_list[i])
        i += 1
    return code_list
        # var data = data.replace(/"/g, '\{\{1\}\}');
    #   return data.replace(/'/g, '\{\{2\}\}');
