updated helper_scripts from battelle/afl-unicorn

This commit is contained in:
Dominik Maier 2020-12-04 15:40:38 +01:00
parent f0e81b2301
commit 330f33a435
3 changed files with 239 additions and 96 deletions

View File

@ -44,6 +44,7 @@ MAX_SEG_SIZE = 128 * 1024 * 1024
# Name of the index file
INDEX_FILE_NAME = "_index.json"
#----------------------
#---- Helper Functions
@ -59,14 +60,14 @@ def map_arch():
return "arm64be"
elif 'armeb' in arch:
# check for THUMB mode
cpsr = get_register('cpsr')
cpsr = get_register('$cpsr')
if (cpsr & (1 << 5)):
return "armbethumb"
else:
return "armbe"
elif 'arm' in arch:
# check for THUMB mode
cpsr = get_register('cpsr')
cpsr = get_register('$cpsr')
if (cpsr & (1 << 5)):
return "armlethumb"
else:
@ -88,12 +89,8 @@ def dump_regs():
reg_state = {}
for reg in current_arch.all_registers:
reg_val = get_register(reg)
# current dumper script looks for register values to be hex strings
# reg_str = "0x{:08x}".format(reg_val)
# if "64" in get_arch():
# reg_str = "0x{:016x}".format(reg_val)
# reg_state[reg.strip().strip('$')] = reg_str
reg_state[reg.strip().strip('$')] = reg_val
return reg_state
@ -146,6 +143,21 @@ def dump_process_memory(output_dir):
return final_segment_list
#---------------------------------------------
#---- ARM Extention (dump floating point regs)
def dump_float(rge=32):
reg_convert = ""
if map_arch() == "armbe" or map_arch() == "armle" or map_arch() == "armbethumb" or map_arch() == "armbethumb":
reg_state = {}
for reg_num in range(32):
value = gdb.selected_frame().read_register("d" + str(reg_num))
reg_state["d" + str(reg_num)] = int(str(value["u64"]), 16)
value = gdb.selected_frame().read_register("fpscr")
reg_state["fpscr"] = int(str(value), 16)
return reg_state
#----------
#---- Main
@ -173,6 +185,7 @@ def main():
context = {
"arch": dump_arch_info(),
"regs": dump_regs(),
"regs_extended": dump_float(),
"segments": dump_process_memory(output_path),
}
@ -187,4 +200,3 @@ def main():
if __name__ == "__main__":
main()

View File

@ -26,6 +26,13 @@ from unicorn.arm64_const import *
from unicorn.x86_const import *
from unicorn.mips_const import *
# If Capstone libraries are availible (only check once)
try:
from capstone import *
CAPSTONE_EXISTS = 1
except:
CAPSTONE_EXISTS = 0
# Name of the index file
INDEX_FILE_NAME = "_index.json"
@ -86,7 +93,7 @@ class UnicornSimpleHeap(object):
total_chunk_size = UNICORN_PAGE_SIZE + ALIGN_PAGE_UP(size) + UNICORN_PAGE_SIZE
# Gross but efficient way to find space for the chunk:
chunk = None
for addr in xrange(self.HEAP_MIN_ADDR, self.HEAP_MAX_ADDR, UNICORN_PAGE_SIZE):
for addr in range(self.HEAP_MIN_ADDR, self.HEAP_MAX_ADDR, UNICORN_PAGE_SIZE):
try:
self._uc.mem_map(addr, total_chunk_size, UC_PROT_READ | UC_PROT_WRITE)
chunk = self.HeapChunk(addr, total_chunk_size, size)
@ -184,29 +191,17 @@ class AflUnicornEngine(Uc):
# Load the registers
regs = context['regs']
reg_map = self.__get_register_map(self._arch_str)
for register, value in regs.iteritems():
if debug_print:
print("Reg {0} = {1}".format(register, value))
if not reg_map.has_key(register.lower()):
if debug_print:
print("Skipping Reg: {}".format(register))
else:
reg_write_retry = True
try:
self.reg_write(reg_map[register.lower()], value)
reg_write_retry = False
except Exception as e:
if debug_print:
print("ERROR writing register: {}, value: {} -- {}".format(register, value, repr(e)))
self.__load_registers(regs, reg_map, debug_print)
# If we have extra FLOATING POINT regs, load them in!
if 'regs_extended' in context:
if context['regs_extended']:
regs_extended = context['regs_extended']
reg_map = self.__get_registers_extended(self._arch_str)
self.__load_registers(regs_extended, reg_map, debug_print)
if reg_write_retry:
if debug_print:
print("Trying to parse value ({}) as hex string".format(value))
try:
self.reg_write(reg_map[register.lower()], int(value, 16))
except Exception as e:
if debug_print:
print("ERROR writing hex string register: {}, value: {} -- {}".format(register, value, repr(e)))
# For ARM, sometimes the stack pointer is erased ??? (I think I fixed this (issue with ordering of dumper.py, I'll keep the write anyways)
if self.__get_arch_and_mode(self.get_arch_str())[0] == UC_ARCH_ARM:
self.reg_write(UC_ARM_REG_SP, regs['sp'])
# Setup the memory map and load memory content
self.__map_segments(context['segments'], context_directory, debug_print)
@ -253,21 +248,76 @@ class AflUnicornEngine(Uc):
for reg in sorted(self.__get_register_map(self._arch_str).items(), key=lambda reg: reg[0]):
print(">>> {0:>4}: 0x{1:016x}".format(reg[0], self.reg_read(reg[1])))
def dump_regs_extended(self):
""" Dumps the contents of all the registers to STDOUT """
try:
for reg in sorted(self.__get_registers_extended(self._arch_str).items(), key=lambda reg: reg[0]):
print(">>> {0:>4}: 0x{1:016x}".format(reg[0], self.reg_read(reg[1])))
except Exception as e:
print("ERROR: Are extended registers loaded?")
# TODO: Make this dynamically get the stack pointer register and pointer width for the current architecture
"""
def dump_stack(self, window=10):
arch = self.get_arch()
mode = self.get_mode()
# Get stack pointers and bit sizes for given architecture
if arch == UC_ARCH_X86 and mode == UC_MODE_64:
stack_ptr_addr = self.reg_read(UC_X86_REG_RSP)
bit_size = 8
elif arch == UC_ARCH_X86 and mode == UC_MODE_32:
stack_ptr_addr = self.reg_read(UC_X86_REG_ESP)
bit_size = 4
elif arch == UC_ARCH_ARM64:
stack_ptr_addr = self.reg_read(UC_ARM64_REG_SP)
bit_size = 8
elif arch == UC_ARCH_ARM:
stack_ptr_addr = self.reg_read(UC_ARM_REG_SP)
bit_size = 4
elif arch == UC_ARCH_ARM and mode == UC_MODE_THUMB:
stack_ptr_addr = self.reg_read(UC_ARM_REG_SP)
bit_size = 4
elif arch == UC_ARCH_MIPS:
stack_ptr_addr = self.reg_read(UC_MIPS_REG_SP)
bit_size = 4
print("")
print(">>> Stack:")
stack_ptr_addr = self.reg_read(UC_X86_REG_RSP)
for i in xrange(-window, window + 1):
addr = stack_ptr_addr + (i*8)
print("{0}0x{1:016x}: 0x{2:016x}".format( \
'SP->' if i == 0 else ' ', addr, \
'SP->' if i == 0 else ' ', addr, \
struct.unpack('<Q', self.mem_read(addr, 8))[0]))
"""
#-----------------------------
#---- Loader Helper Functions
def __load_registers(self, regs, reg_map, debug_print):
for register, value in regs.items():
if debug_print:
print("Reg {0} = {1}".format(register, value))
if register.lower() not in reg_map:
if debug_print:
print("Skipping Reg: {}".format(register))
else:
reg_write_retry = True
try:
self.reg_write(reg_map[register.lower()], value)
reg_write_retry = False
except Exception as e:
if debug_print:
print("ERROR writing register: {}, value: {} -- {}".format(register, value, repr(e)))
if reg_write_retry:
if debug_print:
print("Trying to parse value ({}) as hex string".format(value))
try:
self.reg_write(reg_map[register.lower()], int(value, 16))
except Exception as e:
if debug_print:
print("ERROR writing hex string register: {}, value: {} -- {}".format(register, value, repr(e)))
def __map_segment(self, name, address, size, perms, debug_print=False):
# - size is unsigned and must be != 0
# - starting address must be aligned to 4KB
@ -354,7 +404,7 @@ class AflUnicornEngine(Uc):
else:
if debug_print:
print("No content found for segment {0} @ {1:016x}".format(name, seg_start))
self.mem_write(seg_start, '\x00' * (seg_end - seg_start))
self.mem_write(seg_start, b'\x00' * (seg_end - seg_start))
def __get_arch_and_mode(self, arch_str):
arch_map = {
@ -398,7 +448,6 @@ class AflUnicornEngine(Uc):
"r14": UC_X86_REG_R14,
"r15": UC_X86_REG_R15,
"rip": UC_X86_REG_RIP,
"rsp": UC_X86_REG_RSP,
"efl": UC_X86_REG_EFLAGS,
"cs": UC_X86_REG_CS,
"ds": UC_X86_REG_DS,
@ -415,7 +464,6 @@ class AflUnicornEngine(Uc):
"esi": UC_X86_REG_ESI,
"edi": UC_X86_REG_EDI,
"ebp": UC_X86_REG_EBP,
"esp": UC_X86_REG_ESP,
"eip": UC_X86_REG_EIP,
"esp": UC_X86_REG_ESP,
"efl": UC_X86_REG_EFLAGS,
@ -519,29 +567,97 @@ class AflUnicornEngine(Uc):
}
return registers[arch]
def __get_registers_extended(self, arch):
# Similar to __get_register_map, but for ARM floating point registers
if arch == "arm64le" or arch == "arm64be":
arch = "arm64"
elif arch == "armle" or arch == "armbe" or "thumb" in arch:
arch = "arm"
elif arch == "mipsel":
arch = "mips"
registers = {
"arm": {
"d0": UC_ARM_REG_D0,
"d1": UC_ARM_REG_D1,
"d2": UC_ARM_REG_D2,
"d3": UC_ARM_REG_D3,
"d4": UC_ARM_REG_D4,
"d5": UC_ARM_REG_D5,
"d6": UC_ARM_REG_D6,
"d7": UC_ARM_REG_D7,
"d8": UC_ARM_REG_D8,
"d9": UC_ARM_REG_D9,
"d10": UC_ARM_REG_D10,
"d11": UC_ARM_REG_D11,
"d12": UC_ARM_REG_D12,
"d13": UC_ARM_REG_D13,
"d14": UC_ARM_REG_D14,
"d15": UC_ARM_REG_D15,
"d16": UC_ARM_REG_D16,
"d17": UC_ARM_REG_D17,
"d18": UC_ARM_REG_D18,
"d19": UC_ARM_REG_D19,
"d20": UC_ARM_REG_D20,
"d21": UC_ARM_REG_D21,
"d22": UC_ARM_REG_D22,
"d23": UC_ARM_REG_D23,
"d24": UC_ARM_REG_D24,
"d25": UC_ARM_REG_D25,
"d26": UC_ARM_REG_D26,
"d27": UC_ARM_REG_D27,
"d28": UC_ARM_REG_D28,
"d29": UC_ARM_REG_D29,
"d30": UC_ARM_REG_D30,
"d31": UC_ARM_REG_D31,
"fpscr": UC_ARM_REG_FPSCR
}
}
return registers[arch];
#---------------------------
# Callbacks for tracing
# TODO: Make integer-printing fixed widths dependent on bitness of architecture
# (i.e. only show 4 bytes for 32-bit, 8 bytes for 64-bit)
# TODO: Figure out how best to determine the capstone mode and architecture here
"""
try:
# If Capstone is installed then we'll dump disassembly, otherwise just dump the binary.
from capstone import *
cs = Cs(CS_ARCH_MIPS, CS_MODE_MIPS32 + CS_MODE_BIG_ENDIAN)
def __trace_instruction(self, uc, address, size, user_data):
mem = uc.mem_read(address, size)
for (cs_address, cs_size, cs_mnemonic, cs_opstr) in cs.disasm_lite(bytes(mem), size):
print(" Instr: {:#016x}:\t{}\t{}".format(address, cs_mnemonic, cs_opstr))
except ImportError:
def __trace_instruction(self, uc, address, size, user_data):
print(" Instr: addr=0x{0:016x}, size=0x{1:016x}".format(address, size))
"""
# TODO: Extra mode for Capstone (i.e. Cs(cs_arch, cs_mode + cs_extra) not implemented
def __trace_instruction(self, uc, address, size, user_data):
print(" Instr: addr=0x{0:016x}, size=0x{1:016x}".format(address, size))
if CAPSTONE_EXISTS == 1:
# If Capstone is installed then we'll dump disassembly, otherwise just dump the binary.
arch = self.get_arch()
mode = self.get_mode()
bit_size = self.bit_size_arch()
# Map current arch to capstone labeling
if arch == UC_ARCH_X86 and mode == UC_MODE_64:
cs_arch = CS_ARCH_X86
cs_mode = CS_MODE_64
elif arch == UC_ARCH_X86 and mode == UC_MODE_32:
cs_arch = CS_ARCH_X86
cs_mode = CS_MODE_32
elif arch == UC_ARCH_ARM64:
cs_arch = CS_ARCH_ARM64
cs_mode = CS_MODE_ARM
elif arch == UC_ARCH_ARM and mode == UC_MODE_THUMB:
cs_arch = CS_ARCH_ARM
cs_mode = CS_MODE_THUMB
elif arch == UC_ARCH_ARM:
cs_arch = CS_ARCH_ARM
cs_mode = CS_MODE_ARM
elif arch == UC_ARCH_MIPS:
cs_arch = CS_ARCH_MIPS
cs_mode = CS_MODE_MIPS32 # No other MIPS supported in program
cs = Cs(cs_arch, cs_mode)
mem = uc.mem_read(address, size)
if bit_size == 4:
for (cs_address, cs_size, cs_mnemonic, cs_opstr) in cs.disasm_lite(bytes(mem), size):
print(" Instr: {:#08x}:\t{}\t{}".format(address, cs_mnemonic, cs_opstr))
else:
for (cs_address, cs_size, cs_mnemonic, cs_opstr) in cs.disasm_lite(bytes(mem), size):
print(" Instr: {:#16x}:\t{}\t{}".format(address, cs_mnemonic, cs_opstr))
else:
print(" Instr: addr=0x{0:016x}, size=0x{1:016x}".format(address, size))
def __trace_block(self, uc, address, size, user_data):
print("Basic Block: addr=0x{0:016x}, size=0x{1:016x}".format(address, size))
@ -558,3 +674,18 @@ class AflUnicornEngine(Uc):
else:
print(" >>> INVALID Read: addr=0x{0:016x} size={1}".format(address, size))
def bit_size_arch(self):
arch = self.get_arch()
mode = self.get_mode()
# Get bit sizes for given architecture
if arch == UC_ARCH_X86 and mode == UC_MODE_64:
bit_size = 8
elif arch == UC_ARCH_X86 and mode == UC_MODE_32:
bit_size = 4
elif arch == UC_ARCH_ARM64:
bit_size = 8
elif arch == UC_ARCH_ARM:
bit_size = 4
elif arch == UC_ARCH_MIPS:
bit_size = 4
return bit_size