"""Determine if an instruction is a call or an

unconditional jump

@param op: a distorm3 Op object

"""

"""Disassemble a block of data and yield possible

calls to imported functions. We're looking for

instructions such as these:

x86:

CALL DWORD [0x1000400]

JMP DWORD [0x1000400]

We also capture indirect call where the API address is moved into

a register and the regsiter is called:

MOV ECX [0x9400c]

CALL ECX Register

"""

# Call scan has two options based on the limit settings

# if there are no limits then only potential IAT pointers

# to inside the scanned data will be kept. This is useful

# if you are scanning a data blob and are not able to resolve

# outside the blob.

#

# If the limits are set then any potential IAT pointers to

# inside the limits are kept. This is useful if you are

# scanning the code segment of a PE file but your IAT may

# be located in another segment in the PE file. Set the

# limits to be the start and end of the mapped PE.

#

if start_limit == None:

start_limit = data_vr_address

if end_limit ==None:

end_limit = data_vr_address + len(data)

iat_ptrs=[]

reg_redirect = {"EAX":0x0, "EBX":0x0, "ECX":0x0, "EDX":0x0}

mode = distorm3.Decode32Bits

for op in distorm3.DecomposeGenerator(data_vr_address, data, mode):

if not op.valid:

continue

iat_loc = None

if (_call_or_unc_jmp(op) and op.operands[0].type == 'AbsoluteMemoryAddress'):

iat_loc = (op.operands[0].disp) & 0xffffffff

if op.mnemonic == "MOV" and op.operands[0].type == 'Register' and op.operands[1].type == 'AbsoluteMemory':

#print "MOV %s %s %s" % (op.operands[0], op.operands[1], op.operands[1].type)

reg_redirect[str(op.operands[0])] = op.address + op.operands[1].disp

if op.mnemonic == "MOV" and op.operands[0].type == 'Register' and op.operands[1].type == 'AbsoluteMemoryAddress':

#print "MOV %s %s %s" % (op.operands[0], op.operands[1], op.operands[1].type)

reg_redirect[str(op.operands[0])] =op.operands[1].disp

if op.mnemonic == "CALL" and op.operands[0].type == 'Register':

#print "CALL %s %s" % (op.operands[0], op.operands[0].type)

iat_loc = reg_redirect[str(op.operands[0])]

if (not iat_loc or (iat_loc < start_limit) or (iat_loc > end_limit)):

continue

# resolve iat_loc to API

#print iat_loc

if iat_loc not in iat_ptrs:

iat_ptrs.append(iat_loc)

"""Determine if an instruction is able to reference an IAT pointer

@param op: a distorm3 Op object

"""

"""Use winappdbg to resolve IAT pointers to their respective module and function names

@param pid: process ID to connect to

@param iat_ptrs: list of pointer addresses to be resolved

"""

######################################################################

#

# Attach to process and start using winappdbg

#

######################################################################

# Request debug privileges.

System.request_debug_privileges()

# Attach to process

process = Process(pid)

# Lookup the process modules.

process.scan_modules()

# imp_table[ <funct_pointer> ] = [ <module_name>, <function_name> ]

imp_table = {}

for iat_ptr in iat_ptrs:

# For each iat pointer get the function name as a label populated by winappdbg

label = process.get_label_at_address(process.peek_dword(iat_ptr))

module,function,offset = Process.split_label_strict(label)

# Only add functions that have valid labels

if function != None:

imp_table[iat_ptr] = [module, function]

assert len(imp_table) != 0, "Unable to find imports in code!"

######################################################################

#

# Because we may have missed some IAT pointers with our scanner we

# are going to attempt to locate the full mapped IAT directory in the

# section then enumerate ever pointer in the directory. And use that

# list instead.

#

######################################################################

imp_table_new={}

for iat_ptr in range(min(imp_table.keys()), max(imp_table.keys())+4, 4):

# Resolve the requested label address.

label = process.get_label_at_address(process.peek_dword(iat_ptr))

module,function,offset = Process.split_label_strict(label)

if function != None:

imp_table_new[iat_ptr] = [module, function]

"""Rebuild the import address table for the pe_data that was passed.

@param pid: process ID for winappdbg to attach to and dump IAT offsets

@param pe_data: full PE file read in as a binary string

@param base_address: base address of PE (this override the base addres set in the pe_data)

@param oep: original entry point of the PE (this override the base addres set in the pe_data)

@param newimpdir: name for new section that will contain imports

@param newiat: name for new section that will contain new IAT

@param loadfrommem: pe data is mapped or unmapped (default mapped)

"""

pf = pe_init.PE(loadfrommem=loadfrommem, pestr=pe_data)

pf.NThdr.ImageBase = base_address

# get offset to oep

rva_oep = oep - base_address

pf.Opthdr.AddressOfEntryPoint = rva_oep

# clear the existing import table

# there are two different versions of elfesteem one that uses an extra .l[] object

# and one that does not. If we get the wrong one, just catch the error and continue.

try:

pf.DirImport.impdesc.l=[]

except:

pf.DirImport.impdesc =[]

######################################################################

pdata = None

data_vr_addr = None

data_rva = None

# locate section that contains OEP

for tmp_sec in pf.SHList:

tmp_start = tmp_sec.addr

tmp_end = tmp_start + tmp_sec.size

if (rva_oep >= tmp_start) and (rva_oep <= tmp_end):

try:

pdata = pf._rva.get(tmp_start,tmp_end)

except AttributeError as e:

raise AttributeError("You are using the wrong version of elfesteem, don't use pip instead install from https://github.com/serpilliere/elfesteem")

data_vr_addr = base_address + tmp_start

data_rva = tmp_start

break

# make sure we have found the correct section

# WARNING! if the pdata section is not the same as the one loaded in memory the IAT extraction will

# fail because the addresses and offsets will be wrong.

assert pdata != None, "Unable to locate rva_oep: 0x%x in sections: %s" % (rva_oep, pf.SHList)

# find all call/jmp to possible IAT function pointers

# iat_ptrs is a list of all the addresses that are potential IAT pointers

iat_ptrs = call_scan(data_vr_addr, pdata, start_limit=base_address, end_limit=base_address+len(pe_data))

assert len(iat_ptrs) != 0, "Unable to find IAT pointer candidates in code!"

imp_table = reslove_iat_pointers(pid, iat_ptrs)

# Create a table with module names as the keys and all the functions assigned accordingly

#[module]:[func1, func2, ...]

mod_table = {}

for iat_ptr in imp_table.keys():

# TODO: it is possbile some module names may end with somethin other than .dll

tmp_mod = imp_table[iat_ptr][0]+".dll"

tmp_fn = imp_table[iat_ptr][1]

if tmp_mod in mod_table.keys():

f_arr = mod_table[tmp_mod]

# Only add function name if it isn't already assigned

if tmp_fn not in f_arr:

f_arr.append(tmp_fn)

mod_table[tmp_mod] = f_arr

else:

mod_table[tmp_mod] = [tmp_fn]

newiat_rawsize = (( (len(imp_table.keys()) * 4 ) / 0x1000) + 1) * 0x1000

# Create new section to hold new IAT

s_newiat = pf.SHList.add_section(name=newiat, rawsize=newiat_rawsize)

######################################################################

#

# elfesteem.PE has a special format for describing the IAT directory

# We are converting the mod_table into that format...

#

######################################################################

new_dll=[({"name": mod_table.keys()[0],"firstthunk": s_newiat.addr}, mod_table[mod_table.keys()[0]])]

for mod in mod_table.keys()[1:]:

tmp_entry = ({"name": mod,"firstthunk": None}, mod_table[mod])

new_dll.append(tmp_entry)

######################################################################

#

# Add the new imports table directory to PE file

#

######################################################################

pf.DirImport.add_dlldesc(new_dll)

newimpdir_rawsize = ((len(pf.DirImport) / 0x1000) + 1) * 0x1000

s_newimpdir = pf.SHList.add_section(name=newimpdir, rawsize=newimpdir_rawsize)

pf.SHList.align_sections(0x1000, 0x1000)

pf.DirImport.set_rva(s_newimpdir.addr)

######################################################################

#

# Create a mapping from the old IAT pointers to the new ones

# iat_map[ <old_pointer> ] = <new_pointer>

#

# TODO: handle import by ordinal

#

######################################################################

iat_map ={}

for iat_ptr in imp_table.keys():

if imp_table[iat_ptr][1] == None:

continue

tmp_mod = imp_table[iat_ptr][0]+".dll"

tmp_fn = imp_table[iat_ptr][1]

try:

iat_map[iat_ptr] = pf.DirImport.get_funcvirt(tmp_mod, tmp_fn)

except:

continue

# If there is some error building the module map stop!

assert iat_map != {}, "The iat_map is empty, we are unable to find the new IAT pointers for the old pointers in imp_table."

######################################################################

#

# Patch the code section in the PE and replace all references to the

# old IAT pointers with references to the new pointers.

#

######################################################################

# Make a copy of the code data to work on

odata = pdata

mode = distorm3.Decode32Bits

for op in distorm3.DecomposeGenerator(data_vr_addr, pdata, mode):

if not op.valid:

continue

iat_loc = None

if _iat_candidate(op):

for operand in op.operands:

test_operand = operand.disp & 0xffffffff

if test_operand in iat_map.keys():

#print "Fixing IAT pointer for: %s" % op

#op_ptr = op.address - base_address

#TODO: is this right??

op_ptr = op.address - data_vr_addr

op_size = op.size

orig_op = odata[op_ptr:op_ptr+op_size]

orig_operand = struct.pack('<I',test_operand)

new_operand = struct.pack('<I',iat_map[test_operand])

new_op = orig_op.replace(orig_operand, new_operand)

odata = odata[:op_ptr] + new_op + odata[op_ptr+op_size:]

#stop testing operands

break

# Copy the patched section back to the PE

pf._rva.set(data_rva,odata)

# Disable rebase, since addresses are absolute any rebase will make this explode

pf.NThdr.dllcharacteristics = 0x0

"""Get a memory map of process

@param process: winappdbg.process object

@return list: [ {'BaseAddress': <>, 'RegionSize' <>, 'State': <>, 'Protect': <>, 'Type': <> 'Owner': <> }, ... ]

"""

# Get the process memory map

memoryMap = process.get_memory_map()

# For each memory block in the map...

#mem_map = [... , {Address, Size, State, Access, Type, Owner}]

mem_map_arr = []

for mbi in memoryMap:

mem_page = {}

# Address and size of memory block.

mem_page['BaseAddress'] = mbi.BaseAddress

mem_page['RegionSize'] = mbi.RegionSize

# State (free or allocated).

if mbi.State == MEM_RESERVE:

mem_page['State'] = "Reserved"

elif mbi.State == MEM_COMMIT:

mem_page['State'] = "Commited"

elif mbi.State == MEM_FREE:

mem_page['State'] = "Free"

else:

mem_page['State'] = "Unknown"

# Page protection bits (R/W/X/G).

if mbi.State != MEM_COMMIT:

mem_page['Protect'] = ""

else:

if mbi.Protect & PAGE_NOACCESS:

mem_page['Protect'] = "--- "

elif mbi.Protect & PAGE_READONLY:

mem_page['Protect'] = "R-- "

elif mbi.Protect & PAGE_READWRITE:

mem_page['Protect'] = "RW- "

elif mbi.Protect & PAGE_WRITECOPY:

mem_page['Protect'] = "RC- "

elif mbi.Protect & PAGE_EXECUTE:

mem_page['Protect'] = "--X "

elif mbi.Protect & PAGE_EXECUTE_READ:

mem_page['Protect'] = "R-X "

elif mbi.Protect & PAGE_EXECUTE_READWRITE:

mem_page['Protect'] = "RWX "

elif mbi.Protect & PAGE_EXECUTE_WRITECOPY:

mem_page['Protect'] = "RCX "

else:

mem_page['Protect'] = "??? "

if mbi.Protect & PAGE_GUARD:

mem_page['Protect'] += "G"

else:

mem_page['Protect'] += "-"

if mbi.Protect & PAGE_NOCACHE:

mem_page['Protect'] += "N"

else:

mem_page['Protect'] += "-"

if mbi.Protect & PAGE_WRITECOMBINE:

mem_page['Protect'] += "W"

else:

mem_page['Protect'] += "-"

# Type (file mapping, executable image, or private memory).

if mbi.Type == MEM_IMAGE:

mem_page['Type'] = "Image"

elif mbi.Type == MEM_MAPPED:

mem_page['Type'] = "Mapped"

elif mbi.Type == MEM_PRIVATE:

mem_page['Type'] = "Private"

elif mbi.Type == 0:

mem_page['Type'] = "Free"

else:

mem_page['Type'] = "Unknown"

# Get the page owner

hProcess = process.get_handle( PROCESS_VM_READ | PROCESS_QUERY_INFORMATION )

mem_page['Owner'] = ''

if mbi.Type in (MEM_IMAGE, MEM_MAPPED):

try:

fileName = GetMappedFileName(hProcess, mbi.BaseAddress)

file_path = winappdbg.PathOperations.native_to_win32_pathname(fileName)

mem_page['Owner'] = p_os.path.basename(file_path)

except WindowsError, e:

mem_page['Owner'] = "???"

#add page info to map

mem_map_arr.append(mem_page)

'''Dump pe-based packer process and rebuild with new original entry point.

This function requires the original PE file in order to use the header and

header corrumption anti-dumping techniques.

@param pid: process ID

@param oep: original entry point

@param orig_pe: binary string containing unmapped original PE file

@param newimpdir: name for new section that will contain imports

@param newiat: name for new section that will contain new IAT

'''

System.request_debug_privileges()

process = Process( pid )

try:

process.suspend()

except WindowsError as e:

pass

file_path = process.get_filename()

file_name = p_os.path.basename(file_path)

#######################################################################

#

# REBUILD THE DUMPED PE

#

# I'm sure there is a better way to do this because all we are really

# doing is dumping the PE mapped sections. Suggestions welcome!

#

# The crazy way we do this is to get a memory map of the whole process

# then find the pages that are owned by the file that spawned the process.

#

#######################################################################

mem_map = get_mem_map(process)

temp_data_arr = {}

for page in mem_map:

if file_name.upper() in page["Owner"].upper():

dump_data = process.peek(page["BaseAddress"],page["RegionSize"])

temp_data_arr[page["BaseAddress"]] = dump_data

# we need to work with the dump as one contiguous data block in "mapped" format.

ordered_mem = temp_data_arr.keys()

ordered_mem.sort()

block_data = temp_data_arr[ordered_mem[0]]

for addr_ptr in range(1,len(ordered_mem)):

padding_len = ordered_mem[addr_ptr] - (ordered_mem[0] + len(block_data))

#print "Padding: %d" % padding_len

# These should be contiguous pages so there should be no need for padding!

block_data += temp_data_arr[ordered_mem[addr_ptr]] + '\x00'*padding_len

# The lowest mapped section is the base address

base_address = ordered_mem[0]

# Elfesteem has a small issue with the way it loads mapped PE files

# instead of using the virtual size for segments it uses the raw size

# this messes up unpacker dumps so we will fix it manually.

pf = pe_init.PE(loadfrommem=False, pestr=orig_pe)

new_sections = []

for tmp_section in pf.SHList:

new_sections.append({"name": tmp_section.name ,"offset": tmp_section.addr ,"size": tmp_section.size ,"addr": tmp_section.addr ,"flags": tmp_section.flags ,"rawsize": tmp_section.size})

# Remove existing sections

pf.SHList.shlist=[]

for tmp_section in new_sections:

pf.SHList.add_section(name=tmp_section["name"],

data=block_data[tmp_section["offset"]:tmp_section["offset"] + tmp_section["rawsize"]],

size=tmp_section["size"],

addr=tmp_section["addr"],

offset=tmp_section["offset"],

rawsize=tmp_section["rawsize"])

pf.NThdr.ImageBase = base_address

pf.Opthdr.AddressOfEntryPoint = oep

# Disable rebase, since addresses are absolute any rebase will make this explode

pf.NThdr.dllcharacteristics = 0x0

# Null out the imports they will be wrong anyway and may cause issues when importing into elfesteem

try:

pf.DirImport.impdesc.l=[]

except:

pf.DirImport.impdesc =[]

#######################################################################

#

# At this point pf contains a fully reconstructed PE but with a

# broken IAT. Fix the IAT!

#

#######################################################################

'''Dump process and rebuild with new original entry point.

@param pid: process ID

@param oep: original entry point

@param newimpdir: name for new section that will contain imports

@param newiat: name for new section that will contain new IAT

'''

System.request_debug_privileges()

process = Process( pid )

try:

process.suspend()

except WindowsError as e:

pass

file_path = process.get_filename()

file_name = p_os.path.basename(file_path)

#######################################################################

#

# REBUILD THE DUMPED PE

#

# I'm sure there is a better way to do this because all we are really

# doing is dumping the PE mapped sections. Suggestions welcome!

#

# The crazy way we do this is to get a memory map of the whole process

# then find the pages that are owned by the file that spawned the process.

#

#######################################################################

mem_map = get_mem_map(process)

temp_data_arr = {}

for page in mem_map:

if file_name.upper() in page["Owner"].upper():

dump_data = process.peek(page["BaseAddress"],page["RegionSize"])

temp_data_arr[page["BaseAddress"]] = dump_data

# we need to work with the dump as one contiguous data block in "mapped" format.

ordered_mem = temp_data_arr.keys()

ordered_mem.sort()

block_data = temp_data_arr[ordered_mem[0]]

for addr_ptr in range(1,len(ordered_mem)):

padding_len = ordered_mem[addr_ptr] - (ordered_mem[0] + len(block_data))

#print "Padding: %d" % padding_len

# These should be contiguous pages so there should be no need for padding!

block_data += temp_data_arr[ordered_mem[addr_ptr]] + '\x00'*padding_len

# The lowest mapped section is the base address

base_address = ordered_mem[0]

# Elfesteem has a small issue with the way it loads mapped PE files

# instead of using the virtual size for segments it uses the raw size

# this messes up unpacker dumps so we will fix it manually.

pf = pe_init.PE(loadfrommem=True, pestr=block_data)

new_sections = []

for tmp_section in pf.SHList:

new_sections.append({"name": tmp_section.name ,"offset": tmp_section.addr ,"size": tmp_section.size ,"addr": tmp_section.addr ,"flags": tmp_section.flags ,"rawsize": tmp_section.size})

# Remove existing sections

pf.SHList.shlist=[]

for tmp_section in new_sections:

pf.SHList.add_section(name=tmp_section["name"],

data=block_data[tmp_section["offset"]:tmp_section["offset"] + tmp_section["rawsize"]],

size=tmp_section["size"],

addr=tmp_section["addr"],

offset=tmp_section["offset"],

rawsize=tmp_section["rawsize"])

pf.NThdr.ImageBase = base_address

pf.Opthdr.AddressOfEntryPoint = oep

# Disable rebase, since addresses are absolute any rebase will make this explode

pf.NThdr.dllcharacteristics = 0x0

#######################################################################

#

# At this point pf contains a fully reconstructed PE but with a

# broken IAT. Fix the IAT!

#

#######################################################################

parser = argparse.ArgumentParser(description="Simple example of PyIATRebuild library in use!")

subparsers = parser.add_subparsers(help='', dest='subparser_name')

# create the parser for the load command

parser_rebuild = subparsers.add_parser('rebuild', help='Load dumped PE from file, attach to process, and rebuild IAT.')

parser_rebuild.add_argument("infile", help="The file to fix IAT.")

parser_rebuild.add_argument("outfile", help="The file to write results.")

parser_rebuild.add_argument('--pid',dest="in_pid",type=int,default=None,required=True,help="Specify process ID to export IAT from.")

parser_rebuild.add_argument('--base_address',dest="in_base_address",type=int,default=None,required=True,help="Specify base address the process is loaded at (will overwrite PE).")

parser_rebuild.add_argument('--oep',dest="in_oep",type=int,default=None,required=True,help="Specify original entry point for process, virtual address not RVA (will overwrite PE).")

# create the parser for the results command

parser_dump = subparsers.add_parser('dump', help='Attach to process, dump, and rebuild IAT.')

parser_dump.add_argument("outfile", help="The file to write results.")

parser_dump.add_argument('--pid',dest="in_pid",type=int,default=None,required=True,help="Specify process ID to export IAT from.")

parser_dump.add_argument('--oep',dest="in_oep",type=int,default=None,required=True,help="Specify original entry point for process, virtual address not RVA (will overwrite PE).")

args = parser.parse_args()

if args.subparser_name == "rebuild":

with open(args.infile,"rb") as fp:

pe_data= fp.read()

new_pe_data = rebuild_iat(args.in_pid, pe_data, args.in_base_address, args.in_oep)

open(args.outfile, 'wb').write(new_pe_data)

elif args.subparser_name == "dump":

new_pe_data = dump_and_rebuild(args.in_pid, args.in_oep)