# TODO:
# * modify exports using lief
# * zero out rich header (if it exists) --> requires updating OptionalHeader's checksum ("Rich Header" only in Microsoft-produced executables)
# * tinker with resources: https://lief.quarkslab.com/doc/tutorials/07_pe_resource.html

import lief  # pip install https://github.com/lief-project/LIEF/releases/download/0.7.0/linux_lief-0.7.0_py3.6.tar.gz
import json
import os
import sys
import array
import struct  # byte manipulations
import random
import tempfile
import subprocess
import functools
import signal
import multiprocessing
import pickle
import pefile
import hashlib

module_path = os.path.split(os.path.abspath(sys.modules[__name__].__file__))[0]


#COMMON_SECTION_NAMES = open(os.path.join(module_path, 'section_names.txt'), 'r').read().rstrip().split('\n')
#COMMON_IMPORTS = open(os.path.join(module_path, 'imports.txt'), 'r').read().rstrip().split('\n')
#COMMON_IMPORTS = json.load(
   #open(os.path.join(module_path, 'small_dll_imports.json'), 'r'))

COMMON_SECTION_NAMES = pickle.load(open(os.path.join(module_path, '../../../RL_Features/adversarial_sections_set.pk'), "rb"))
COMMON_IMPORTS = pickle.load(open(os.path.join(module_path, '../../../RL_Features/adversarial_imports_set.pk'), "rb"))
#COMMON_SECTION_NAMES = pickle.load(open(os.path.join(module_path, '../../../RL_Features/adversarial_sections_set1.pk'), "rb"))
#COMMON_IMPORTS = pickle.load(open(os.path.join(module_path, '../../../RL_Features/importsFinal.pk'), "rb"))
section_content = "manipulation_content/section-content.txt"

class MalwareManipulator(object):
    def __init__(self, bytez):
        self.bytez = bytez
        self.min_append_log2 = 5
        self.max_append_log2 = 8

    def __random_length(self):
        return 2**random.randint(self.min_append_log2, self.max_append_log2)

    def __binary_to_bytez(self, binary, dos_stub=False, imports=False, overlay=False, relocations=False, resources=False, tls=False):
        # write the file back as bytez
        builder = lief.PE.Builder(binary)
        builder.build_dos_stub(dos_stub) # rebuild DOS stub

        builder.build_imports(imports) # rebuild IAT in another section
        builder.patch_imports(imports) # patch original import table with trampolines to new import table

        builder.build_overlay(overlay) # rebuild overlay
        builder.build_relocations(relocations) # rebuild relocation table in another section
        builder.build_resources(resources) # rebuild resources in another section
        builder.build_tls(tls) # rebuilt TLS object in another section

        builder.build() # perform the build process

        # return bytestring
        return array.array('B', builder.get_build()).tobytes()

    def overlay_append(self, seed=None):
        random.seed(seed)
        L = self.__random_length()
        # choose the upper bound for a uniform distribution in [0,upper]
        upper = random.randrange(256)
        # upper chooses the upper bound on uniform distribution:
        # upper=0 would append with all 0s
        # upper=126 would append with "printable ascii"
        # upper=255 would append with any character
        return self.bytez + bytes([random.randint(0, upper) for _ in range(L)])

    def imports_append(self, seed=None):
        #COMMON_IMPORTS_NAMES = ['ADVAPI32.DLL', 'SHLWAPI.DLL', 'KERNEL32.DLL','USER32.DLL']
        importsFile = open("imports.txt" , 'w')

        libname = random.choice(list(COMMON_IMPORTS))

        while(len(list(COMMON_IMPORTS[libname])) < 20 ):
            libname = random.choice(list(COMMON_IMPORTS))

        importsFile.write(libname + '\n')
        for fun in (list(COMMON_IMPORTS[libname])):
            importsFile.write(fun + '\n')

        importsFile.close()

        #print('adding import library : ' + libname)
        with open("modified.exe", 'wb') as file1:
            file1.write(self.bytez)

        cmd = "portable-executable/project-add-imports/bin/Debug/project-append-import modified.exe imports.txt modified.exe"
        os.system(cmd)

        with open("modified.exe", "rb") as binfile:
            self.bytez = binfile.read()

        return self.bytez


    def section_add(self, seed=None): 
  
        section = ""
        while(len(section.strip()) < 1):   
            section = random.choice(COMMON_SECTION_NAMES)
        with open("modified.exe", 'wb') as file1:
            file1.write(self.bytez)

        cmd = "portable-executable/project-add-sections/bin/Debug/project-append-section modified.exe " + section + " " + section_content + " modified.exe"
        os.system(cmd)

        #print('adding section : ' + section)

        with open("modified.exe", "rb") as binfile:
            self.bytez = binfile.read()

        return self.bytez



    def section_append(self, seed=None):
        # append to a section (changes size and entropy)
        random.seed(seed)
        binary1 = lief.parse(self.bytez, name="")
        with open("modified.exe", 'wb') as file1:
            file1.write(bytez)

        binary = lief.PE.Binary("modified.exe", lief.PE.PE_TYPE.PE32)

        if(len(binary1.sections) == 0):
            #print("no sections in parse")
            self.bytez = self.__binary_to_bytez(binary)
            return self.bytez

        if(len(binary.sections) == 0):
            #print("no sections in binary")
            self.bytez = self.__binary_to_bytez(binary)
            return self.bytez


        targeted_section = random.choice(binary.sections)
        L = self.__random_length()
        available_size = targeted_section.size - len(targeted_section.content)
        if L > available_size:
            L = available_size

        upper = random.randrange(256)
        targeted_section.content = targeted_section.content + [random.randint(0, upper) for _ in range(L)]

        self.bytez = self.__binary_to_bytez(binary)
        return self.bytez

    def section_rename(self, seed=None):
        # rename a random section
        #random.seed(seed)
        binary = lief.parse(self.bytez, name="")
        for i in range(0, 10):   
            targeted_section = random.choice(binary.sections)
            targeted_section.name = random.choice(COMMON_SECTION_NAMES)[:7] # current version of lief not allowing 8 chars?

        self.bytez = self.__binary_to_bytez(binary)

        return self.bytez

    def upx_pack(self, seed=None):
        # tested with UPX 3.91
        random.seed(seed)
        tmpfilename = os.path.join(
            tempfile._get_default_tempdir(), next(tempfile._get_candidate_names()))

        # dump bytez to a temporary file
        with open(tmpfilename, 'wb') as outfile:
            outfile.write(self.bytez)

        options = ['--force', '--overlay=copy']
        compression_level = random.randint(1, 9)
        options += ['-{}'.format(compression_level)]
        # --exact
        # compression levels -1 to -9
        # --overlay=copy [default]

        # optional things:
        # --compress-exports=0/1
        # --compress-icons=0/1/2/3
        # --compress-resources=0/1
        # --strip-relocs=0/1
        options += ['--compress-exports={}'.format(random.randint(0, 1))]
        options += ['--compress-icons={}'.format(random.randint(0, 3))]
        options += ['--compress-resources={}'.format(random.randint(0, 1))]
        options += ['--strip-relocs={}'.format(random.randint(0, 1))]

        with open(os.devnull, 'w') as DEVNULL:
            retcode = subprocess.call(
                ['upx'] + options + [tmpfilename, '-o', tmpfilename + '_packed'], stdout=DEVNULL, stderr=DEVNULL)

        os.unlink(tmpfilename)

        if retcode == 0:  # successfully packed

            with open(tmpfilename + '_packed', 'rb') as infile:
                self.bytez = infile.read()

            os.unlink(tmpfilename + '_packed')

        return self.bytez

    def upx_unpack(self, seed=None):
        # dump bytez to a temporary file
        tmpfilename = os.path.join(
            tempfile._get_default_tempdir(), next(tempfile._get_candidate_names()))

        with open(tmpfilename, 'wb') as outfile:
            outfile.write(self.bytez)

        with open(os.devnull, 'w') as DEVNULL:
            retcode = subprocess.call(
                ['upx', tmpfilename, '-d', '-o', tmpfilename + '_unpacked'], stdout=DEVNULL, stderr=DEVNULL)

        os.unlink(tmpfilename)

        if retcode == 0:  # sucessfully unpacked
            with open(tmpfilename + '_unpacked', 'rb') as result:
                self.bytez = result.read()

            os.unlink(tmpfilename + '_unpacked')

        return self.bytez

    def remove_signature(self, seed=None):
        binary = lief.parse(self.bytez, name="")

        if binary.has_signature:
            for i, e in enumerate(binary.data_directories):
                if e.type == lief.PE.DATA_DIRECTORY.CERTIFICATE_TABLE:
                    break
            if e.type == lief.PE.DATA_DIRECTORY.CERTIFICATE_TABLE:
                # remove signature from certificate table
                e.rva = 0
                e.size = 0
                self.bytez = self.__binary_to_bytez(binary)
                return self.bytez
        # if no signature found, self.bytez is unmodified
        return self.bytez

    def add_signature(self, seed=None):
        # with open("modified.exe", 'wb') as file1:
        #     file1.write(self.bytez)

        # cmd = "osslsigncode sign -certs \"cert.pem\" -key \"key.pem\" -pass \"PaSSw0Rd1!\" -n \"MyApp 1.2\" -i \"https://www.mysite.com\" -t \"http://timestamp.comodoca.com/authenticode\" -in \"modified.exe\" -out \"modified.exe\""
        # os.system(cmd)

        # with open("modified.exe", "rb") as binfile:
        #     self.bytez = binfile.read()

        return self.bytez

        


    def edit_tls(self, seed):
        with open("modified.exe", 'wb') as file1:
            file1.write(self.bytez)

        cmd = 'portable-executable/test-other/bin/edit-tls/test-other modified.exe' 
        os.system(cmd)

        with open("modified.exe", "rb") as binfile:
            self.bytez = binfile.read()

        return self.bytez

    def pe_stripper(self, seed):
        with open("modified.exe", 'wb') as file1:
            file1.write(self.bytez)

        cmd = 'portable-executable/test-other/bin/pe-stripper/test-other modified.exe' 
        os.system(cmd)

        with open("modified.exe", "rb") as binfile:
            self.bytez = binfile.read()

        return self.bytez

    def load_config_dir(self, seed):
        with open("modified.exe", 'wb') as file1:
            file1.write(self.bytez)

        cmd = 'portable-executable/test-other/bin/load-config-dir/test-other modified.exe'
        os.system(cmd)

        with open("modified.exe", "rb") as binfile:
            self.bytez = binfile.read()

        return self.bytez


    def remove_debug(self, seed=None):
        binary = lief.parse(self.bytez, name="")

        if binary.has_debug:
            for i, e in enumerate(binary.data_directories):
                if e.type == lief.PE.DATA_DIRECTORY.DEBUG:
                    break
            if e.type == lief.PE.DATA_DIRECTORY.DEBUG:
                # remove signature from certificate table
                e.rva = 0
                e.size = 0
                self.bytez = self.__binary_to_bytez(binary)
                return self.bytez
        # if no signature found, self.bytez is unmodified
        return self.bytez

    def break_optional_header_checksum(self, seed=None):
        binary = lief.parse(self.bytez, name="")
        binary.optional_header.checksum = 0
        self.bytez = self.__binary_to_bytez(binary)
        return self.bytez

    


##############################
def identity(bytez, seed=None):
    return bytez


######################
# explicitly list so that these may be used externally
# action table for azure rl
ACTION_TABLE = {
    # 'do_nothing': identity,
    'overlay_append': 'overlay_append',
    'section_rename' : 'section_rename',
    'section_add' : 'section_add',
    'imports_append' : 'imports_append'

}


# ACTION_TABLE = {
#     # 'do_nothing': identity,
#     'overlay_append': 'overlay_append',
#     'section_rename' : 'section_rename',
#     'add_signature' : 'add_signature',
#     'edit_tls' : 'edit_tls',
#     'load_config_dir' : 'load_config_dir',
#     'section_add' : 'section_add',
#     'imports_append' : 'imports_append',
#     'remove_signature': 'remove_signature',
#     'remove_debug': 'remove_debug'

# }


def calculate_hash(bytez):
    m = hashlib.sha256()
    m.update( bytez )
    #print("new hash: {}".format(m.hexdigest()))

def modify_without_breaking(bytez, actions=[], seed=None):
    for action in actions:

        _action = ACTION_TABLE[action]

        # we run manipulation in a child process to shelter
        # our malware model from rare parsing errors in LIEF that
        # may segfault or timeout
        def helper(_action,shared_list):
            # TODO: LIEF is chatty. redirect stdout and stderr to /dev/null

            # for this process, change segfault of the child process
            # to a RuntimeEror
            def sig_handler(signum, frame):
                raise RuntimeError
            signal.signal(signal.SIGSEGV, sig_handler)

            bytez = array.array('B', shared_list[:]).tobytes()
            # TODO: LIEF is chatty. redirect output to /dev/null
            if type(_action) is str:
                _action = MalwareManipulator(bytez).__getattribute__(_action)
            else:
                _action = functools.partial( _action, bytez )

            # redirect standard out only in this queue
            try:
                shared_list[:] = _action(seed) 
            except (RuntimeError,UnicodeDecodeError,TypeError,lief.not_found) as e:
                # some exceptions that have yet to be handled by public release of LIEF
                print("==== exception in child process ===")
                print(e)
                # shared_bytez remains unchanged                


        # communicate with the subprocess through a shared list
        # can't use multiprocessing.Array since the subprocess may need to
        # change the size
        manager = multiprocessing.Manager()
        shared_list = manager.list() 
        shared_list[:] = bytez # copy bytez to shared array
        # define process
        p = multiprocessing.Process( target=helper, args=(_action,shared_list) ) 
        p.start() # start the process
        try:
            p.join(5) # allow this to take up to 5 seconds...
        except multiprocessing.TimeoutError: # ..then become petulant
            print('==== timeouterror ')
            p.terminate()

        bytez = array.array('B', shared_list[:]).tobytes() # copy result from child process

    import hashlib
    m = hashlib.sha256()
    m.update( bytez )
    #print("new hash: {}".format(m.hexdigest()))
    return bytez



if __name__ == "__main__":


    filename = sys.argv[1]

    with open(filename, "rb") as binfile:
        bytez = binfile.read()

    #binary = lief.PE.parse(bytez , name="")
    calculate_hash(bytez)
    manip = MalwareManipulator(bytez) 


    print('pe_stripper')
    bytez = manip.pe_stripper(bytez)
    with open("modified_pe_stripper.exe", 'wb') as file1:
        file1.write(bytez)  

    # print('overlay_append')
    # bytez = manip.overlay_append(bytez)
    # calculate_hash(bytez)

    # with open("modified_overlay_append.exe", 'wb') as file1:
    #     file1.write(bytez)    

    # print('remove_signature')
    # bytez = manip.remove_signature(bytez)
    # calculate_hash(bytez)

    # with open("modified_remove_sig.exe", 'wb') as file1:
    #     file1.write(bytez)    

    # print('add_signature')
    # bytez = manip.add_signature(bytez)
    # calculate_hash(bytez)

    # with open("modified_add_sig.exe", 'wb') as file1:
    #     file1.write(bytez)

    # print('edit_tls')
    # bytez = manip.edit_tls(bytez)
    # calculate_hash(bytez)

    # with open("modified_edit_tls.exe", 'wb') as file1:
    #     file1.write(bytez)

    # print('load_config_dir')
    # bytez = manip.load_config_dir(bytez)
    # calculate_hash(bytez)

    # with open("modified_load_config_dir.exe", 'wb') as file1:
    #     file1.write(bytez)

    # print('import append')
    # bytez = manip.imports_append(bytez)
    # calculate_hash(bytez)


    # with open("modified_import_append.exe", 'wb') as file1:
    #     file1.write(bytez)

    # print('section append')
    # bytez = manip.section_add(bytez)
    # calculate_hash(bytez)

    # with open("modified_section_append.exe", 'wb') as file1:
    #     file1.write(bytez)


    # print('remove_debug')
    # bytez = manip.remove_debug(bytez)
    # calculate_hash(bytez)

    # with open("modified_remove_debug.exe", 'wb') as file1:
    #     file1.write(bytez)    

    # print('section_rename')
    # bytez = manip.section_rename()
    # calculate_hash(bytez)

    # with open("modified_section_rename.exe", 'wb') as file1:
    #     file1.write(bytez)



    







