# 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

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 = json.load(
#   open(os.path.join(module_path, 'small_dll_imports.json'), 'r'))

COMMON_SECTION_NAMES = pickle.load(open(os.path.join(module_path, 'adversarial_sections_set.pk'), "rb"))
COMMON_IMPORTS = pickle.load(open(os.path.join(module_path, 'adversarial_imports_set.pk'), "rb"))


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']
        binary = lief.PE.parse(self.bytez, name="")
        common_imports = []
        for im in binary.imports:
          if im.name.upper() in list(COMMON_IMPORTS.keys()):
            common_imports.append(im.name.upper())

        if len(common_imports) == 0:
            common_imports = list(COMMON_IMPORTS.keys())

        
        binary = lief.PE.parse(self.bytez, name="")
        libname = random.choice(list(common_imports))
        funcname = random.choice(list(COMMON_IMPORTS[libname]))
        lowerlibname = libname.lower()
        lib = None

        lib = binary.add_library(lowerlibname)
        names = set([e.name for e in lib.entries])
        if funcname in names or 'ord' in funcname:
          funcname = random.choice(list(COMMON_IMPORTS[libname]))

        lib.add_entry(funcname)
        print("lib name : " + lowerlibname + " , func name : " + funcname)
        self.bytez = self.__binary_to_bytez(binary,imports=True)
        return self.bytez


    def section_add(self, seed=None):
        binary = lief.PE.parse(self.bytez, name="")
        sec_name = random.choice(COMMON_SECTION_NAMES)[:7]
        new_section = lief.PE.Section(sec_name)
        new_section.content = [0xCC] * 0x1000

        sec = [
               lief.PE.SECTION_TYPES.BSS,
               lief.PE.SECTION_TYPES.IDATA,
               lief.PE.SECTION_TYPES.RELOCATION,
               lief.PE.SECTION_TYPES.RESOURCE,
               lief.PE.SECTION_TYPES.TEXT,
               lief.PE.SECTION_TYPES.UNKNOWN
                           ]

        new_sec = random.choice(sec)
        binary.add_section(new_section, new_sec)                  
        print("section name : " + sec_name +", ssection type : " + str(new_sec))

        if(new_sec == lief.PE.SECTION_TYPES.RELOCATION):
        	self.bytez = self.__binary_to_bytez(binary, relocations=True)

        elif(new_sec == lief.PE.SECTION_TYPES.RESOURCE):
        	self.bytez = self.__binary_to_bytez(binary, resources=True)

        else:
        	self.bytez = self.__binary_to_bytez(binary)

        return self.bytez


    def section_append(self, seed=None):
        # append to a section (changes size and entropy)
        random.seed(seed)
        binary = lief.PE.parse(self.bytez, name="")
        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.PE.parse(self.bytez, name="")
        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 identity(bytez, seed=None):
    return bytez


######################
# explicitly list so that these may be used externally
ACTION_TABLE = {
    # 'do_nothing': identity,
    'overlay_append': 'overlay_append',
    'imports_append': 'imports_append',
    'section_add': 'section_add',
    #'section_append': 'section_append'
    'section_rename' : 'section_rename'
}


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
