formatting

This commit is contained in:
Dominik Maier 2021-07-15 12:38:15 +02:00
parent 7fb2d90c49
commit c279750689

View File

@ -29,6 +29,7 @@ from unicornafl.mips_const import *
# If Capstone libraries are availible (only check once)
try:
from capstone import *
CAPSTONE_EXISTS = 1
except:
CAPSTONE_EXISTS = 0
@ -49,6 +50,7 @@ ALIGN_PAGE_UP = lambda x: (x + UNICORN_PAGE_SIZE - 1) & ~(UNICORN_PAGE_SIZE-1)
# ---------------------------------------
# ---- Unicorn-based heap implementation
class UnicornSimpleHeap(object):
"""Use this class to provide a simple heap implementation. This should
be used if malloc/free calls break things during emulation. This heap also
@ -59,14 +61,22 @@ class UnicornSimpleHeap(object):
# Helper data-container used to track chunks
class HeapChunk(object):
def __init__(self, actual_addr, total_size, data_size):
self.total_size = total_size # Total size of the chunk (including padding and guard page)
self.total_size = (
total_size # Total size of the chunk (including padding and guard page)
)
self.actual_addr = actual_addr # Actual start address of the chunk
self.data_size = data_size # Size requested by the caller of actual malloc call
self.data_addr = actual_addr + UNICORN_PAGE_SIZE # Address where data actually starts
self.data_size = (
data_size # Size requested by the caller of actual malloc call
)
self.data_addr = (
actual_addr + UNICORN_PAGE_SIZE
) # Address where data actually starts
# Returns true if the specified buffer is completely within the chunk, else false
def is_buffer_in_chunk(self, addr, size):
if addr >= self.data_addr and ((addr + size) <= (self.data_addr + self.data_size)):
if addr >= self.data_addr and (
(addr + size) <= (self.data_addr + self.data_size)
):
return True
else:
return False
@ -98,7 +108,11 @@ class UnicornSimpleHeap(object):
self._uc.mem_map(addr, total_chunk_size, UC_PROT_READ | UC_PROT_WRITE)
chunk = self.HeapChunk(addr, total_chunk_size, size)
if self._debug_print:
print("Allocating 0x{0:x}-byte chunk @ 0x{1:016x}".format(chunk.data_size, chunk.data_addr))
print(
"Allocating 0x{0:x}-byte chunk @ 0x{1:016x}".format(
chunk.data_size, chunk.data_addr
)
)
break
except UcError as e:
continue
@ -115,14 +129,21 @@ class UnicornSimpleHeap(object):
def realloc(self, ptr, new_size):
# Wrapper around malloc(new_size) / memcpy(new, old, old_size) / free(old)
if self._debug_print:
print("Reallocating chunk @ 0x{0:016x} to be 0x{1:x} bytes".format(ptr, new_size))
print(
"Reallocating chunk @ 0x{0:016x} to be 0x{1:x} bytes".format(
ptr, new_size
)
)
old_chunk = None
for chunk in self._chunks:
if chunk.data_addr == ptr:
old_chunk = chunk
new_chunk_addr = self.malloc(new_size)
if old_chunk != None:
self._uc.mem_write(new_chunk_addr, str(self._uc.mem_read(old_chunk.data_addr, old_chunk.data_size)))
self._uc.mem_write(
new_chunk_addr,
str(self._uc.mem_read(old_chunk.data_addr, old_chunk.data_size)),
)
self.free(old_chunk.data_addr)
return new_chunk_addr
@ -130,7 +151,11 @@ class UnicornSimpleHeap(object):
for chunk in self._chunks:
if chunk.is_buffer_in_chunk(addr, 1):
if self._debug_print:
print("Freeing 0x{0:x}-byte chunk @ 0x{0:016x}".format(chunk.req_size, chunk.data_addr))
print(
"Freeing 0x{0:x}-byte chunk @ 0x{0:016x}".format(
chunk.req_size, chunk.data_addr
)
)
self._uc.mem_unmap(chunk.actual_addr, chunk.total_size)
self._chunks.remove(chunk)
return True
@ -139,19 +164,27 @@ class UnicornSimpleHeap(object):
# Implements basic guard-page functionality
def __check_mem_access(self, uc, access, address, size, value, user_data):
for chunk in self._chunks:
if address >= chunk.actual_addr and ((address + size) <= (chunk.actual_addr + chunk.total_size)):
if address >= chunk.actual_addr and (
(address + size) <= (chunk.actual_addr + chunk.total_size)
):
if chunk.is_buffer_in_chunk(address, size) == False:
if self._debug_print:
print("Heap over/underflow attempting to {0} 0x{1:x} bytes @ {2:016x}".format( \
"write" if access == UC_MEM_WRITE else "read", size, address))
print(
"Heap over/underflow attempting to {0} 0x{1:x} bytes @ {2:016x}".format(
"write" if access == UC_MEM_WRITE else "read",
size,
address,
)
)
# Force a memory-based crash
uc.force_crash(UcError(UC_ERR_READ_PROT))
# ---------------------------
# ---- Loading function
class AflUnicornEngine(Uc):
class AflUnicornEngine(Uc):
def __init__(self, context_directory, enable_trace=False, debug_print=False):
"""
Initializes an AflUnicornEngine instance, which extends standard the UnicornEngine
@ -166,51 +199,56 @@ class AflUnicornEngine(Uc):
# Make sure the index file exists and load it
index_file_path = os.path.join(context_directory, INDEX_FILE_NAME)
if not os.path.isfile(index_file_path):
raise Exception("Index file not found. Expected it to be at {}".format(index_file_path))
raise Exception(
"Index file not found. Expected it to be at {}".format(index_file_path)
)
# Load the process context from the index file
if debug_print:
print("Loading process context index from {}".format(index_file_path))
index_file = open(index_file_path, 'r')
index_file = open(index_file_path, "r")
context = json.load(index_file)
index_file.close()
# Check the context to make sure we have the basic essential components
if 'arch' not in context:
if "arch" not in context:
raise Exception("Couldn't find architecture information in index file")
if 'regs' not in context:
if "regs" not in context:
raise Exception("Couldn't find register information in index file")
if 'segments' not in context:
if "segments" not in context:
raise Exception("Couldn't find segment/memory information in index file")
# Set the UnicornEngine instance's architecture and mode
self._arch_str = context['arch']['arch']
self._arch_str = context["arch"]["arch"]
arch, mode = self.__get_arch_and_mode(self._arch_str)
Uc.__init__(self, arch, mode)
# Load the registers
regs = context['regs']
regs = context["regs"]
reg_map = self.__get_register_map(self._arch_str)
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']
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)
# 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'])
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)
self.__map_segments(context["segments"], context_directory, debug_print)
if enable_trace:
self.hook_add(UC_HOOK_BLOCK, self.__trace_block)
self.hook_add(UC_HOOK_CODE, self.__trace_instruction)
self.hook_add(UC_HOOK_MEM_WRITE | UC_HOOK_MEM_READ, self.__trace_mem_access)
self.hook_add(UC_HOOK_MEM_WRITE_UNMAPPED | UC_HOOK_MEM_READ_INVALID, self.__trace_mem_invalid_access)
self.hook_add(
UC_HOOK_MEM_WRITE_UNMAPPED | UC_HOOK_MEM_READ_INVALID,
self.__trace_mem_invalid_access,
)
if debug_print:
print("Done loading context.")
@ -229,9 +267,15 @@ class AflUnicornEngine(Uc):
You can pass the exception received from Uc.emu_start
"""
mem_errors = [
UC_ERR_READ_UNMAPPED, UC_ERR_READ_PROT, UC_ERR_READ_UNALIGNED,
UC_ERR_WRITE_UNMAPPED, UC_ERR_WRITE_PROT, UC_ERR_WRITE_UNALIGNED,
UC_ERR_FETCH_UNMAPPED, UC_ERR_FETCH_PROT, UC_ERR_FETCH_UNALIGNED,
UC_ERR_READ_UNMAPPED,
UC_ERR_READ_PROT,
UC_ERR_READ_UNALIGNED,
UC_ERR_WRITE_UNMAPPED,
UC_ERR_WRITE_PROT,
UC_ERR_WRITE_UNALIGNED,
UC_ERR_FETCH_UNMAPPED,
UC_ERR_FETCH_PROT,
UC_ERR_FETCH_UNALIGNED,
]
if uc_error.errno in mem_errors:
# Memory error - throw SIGSEGV
@ -245,13 +289,18 @@ class AflUnicornEngine(Uc):
def dump_regs(self):
""" Dumps the contents of all the registers to STDOUT """
for reg in sorted(self.__get_register_map(self._arch_str).items(), key=lambda reg: reg[0]):
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]):
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?")
@ -307,7 +356,11 @@ class AflUnicornEngine(Uc):
reg_write_retry = False
except Exception as e:
if debug_print:
print("ERROR writing register: {}, value: {} -- {}".format(register, value, repr(e)))
print(
"ERROR writing register: {}, value: {} -- {}".format(
register, value, repr(e)
)
)
if reg_write_retry:
if debug_print:
@ -316,7 +369,11 @@ class AflUnicornEngine(Uc):
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)))
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
@ -330,24 +387,30 @@ class AflUnicornEngine(Uc):
if mem_start_aligned != mem_start or mem_end_aligned != mem_end:
print("Aligning segment to page boundary:")
print(" name: {}".format(name))
print(" start: {0:016x} -> {1:016x}".format(mem_start, mem_start_aligned))
print(
" start: {0:016x} -> {1:016x}".format(mem_start, mem_start_aligned)
)
print(" end: {0:016x} -> {1:016x}".format(mem_end, mem_end_aligned))
print("Mapping segment from {0:016x} - {1:016x} with perm={2}: {3}".format(mem_start_aligned, mem_end_aligned, perms, name))
if(mem_start_aligned < mem_end_aligned):
print(
"Mapping segment from {0:016x} - {1:016x} with perm={2}: {3}".format(
mem_start_aligned, mem_end_aligned, perms, name
)
)
if mem_start_aligned < mem_end_aligned:
self.mem_map(mem_start_aligned, mem_end_aligned - mem_start_aligned, perms)
def __map_segments(self, segment_list, context_directory, debug_print=False):
for segment in segment_list:
# Get the segment information from the index
name = segment['name']
seg_start = segment['start']
seg_end = segment['end']
perms = \
(UC_PROT_READ if segment['permissions']['r'] == True else 0) | \
(UC_PROT_WRITE if segment['permissions']['w'] == True else 0) | \
(UC_PROT_EXEC if segment['permissions']['x'] == True else 0)
name = segment["name"]
seg_start = segment["start"]
seg_end = segment["end"]
perms = (
(UC_PROT_READ if segment["permissions"]["r"] == True else 0)
| (UC_PROT_WRITE if segment["permissions"]["w"] == True else 0)
| (UC_PROT_EXEC if segment["permissions"]["x"] == True else 0)
)
if debug_print:
print("Handling segment {}".format(name))
@ -376,48 +439,86 @@ class AflUnicornEngine(Uc):
# Map memory into the address space if it is of an acceptable size.
if (seg_end - seg_start) > MAX_ALLOWABLE_SEG_SIZE:
if debug_print:
print("Skipping segment (LARGER THAN {0}) from {1:016x} - {2:016x} with perm={3}: {4}".format(MAX_ALLOWABLE_SEG_SIZE, seg_start, seg_end, perms, name))
print(
"Skipping segment (LARGER THAN {0}) from {1:016x} - {2:016x} with perm={3}: {4}".format(
MAX_ALLOWABLE_SEG_SIZE, seg_start, seg_end, perms, name
)
)
continue
elif not found: # Make sure it's not already mapped
if overlap_start: # Partial overlap (start)
self.__map_segment(name, tmp, seg_end - tmp, perms, debug_print)
elif overlap_end: # Patrial overlap (end)
self.__map_segment(name, seg_start, tmp - seg_start, perms, debug_print)
self.__map_segment(
name, seg_start, tmp - seg_start, perms, debug_print
)
else: # Not found
self.__map_segment(name, seg_start, seg_end - seg_start, perms, debug_print)
self.__map_segment(
name, seg_start, seg_end - seg_start, perms, debug_print
)
else:
if debug_print:
print("Segment {} already mapped. Moving on.".format(name))
# Load the content (if available)
if 'content_file' in segment and len(segment['content_file']) > 0:
content_file_path = os.path.join(context_directory, segment['content_file'])
if "content_file" in segment and len(segment["content_file"]) > 0:
content_file_path = os.path.join(
context_directory, segment["content_file"]
)
if not os.path.isfile(content_file_path):
raise Exception("Unable to find segment content file. Expected it to be at {}".format(content_file_path))
raise Exception(
"Unable to find segment content file. Expected it to be at {}".format(
content_file_path
)
)
# if debug_print:
# print("Loading content for segment {} from {}".format(name, segment['content_file']))
content_file = open(content_file_path, 'rb')
content_file = open(content_file_path, "rb")
compressed_content = content_file.read()
content_file.close()
self.mem_write(seg_start, zlib.decompress(compressed_content))
else:
if debug_print:
print("No content found for segment {0} @ {1:016x}".format(name, seg_start))
self.mem_write(seg_start, b'\x00' * (seg_end - seg_start))
print(
"No content found for segment {0} @ {1:016x}".format(
name, seg_start
)
)
self.mem_write(seg_start, b"\x00" * (seg_end - seg_start))
def __get_arch_and_mode(self, arch_str):
arch_map = {
"x64": [UC_X86_REG_RIP, UC_ARCH_X86, UC_MODE_64],
"x86": [UC_X86_REG_EIP, UC_ARCH_X86, UC_MODE_32],
"arm64be" : [ UC_ARM64_REG_PC, UC_ARCH_ARM64, UC_MODE_ARM | UC_MODE_BIG_ENDIAN ],
"arm64le" : [ UC_ARM64_REG_PC, UC_ARCH_ARM64, UC_MODE_ARM | UC_MODE_LITTLE_ENDIAN ],
"arm64be": [
UC_ARM64_REG_PC,
UC_ARCH_ARM64,
UC_MODE_ARM | UC_MODE_BIG_ENDIAN,
],
"arm64le": [
UC_ARM64_REG_PC,
UC_ARCH_ARM64,
UC_MODE_ARM | UC_MODE_LITTLE_ENDIAN,
],
"armbe": [UC_ARM_REG_PC, UC_ARCH_ARM, UC_MODE_ARM | UC_MODE_BIG_ENDIAN],
"armle": [UC_ARM_REG_PC, UC_ARCH_ARM, UC_MODE_ARM | UC_MODE_LITTLE_ENDIAN],
"armbethumb": [ UC_ARM_REG_PC, UC_ARCH_ARM, UC_MODE_THUMB | UC_MODE_BIG_ENDIAN ],
"armlethumb": [ UC_ARM_REG_PC, UC_ARCH_ARM, UC_MODE_THUMB | UC_MODE_LITTLE_ENDIAN ],
"armbethumb": [
UC_ARM_REG_PC,
UC_ARCH_ARM,
UC_MODE_THUMB | UC_MODE_BIG_ENDIAN,
],
"armlethumb": [
UC_ARM_REG_PC,
UC_ARCH_ARM,
UC_MODE_THUMB | UC_MODE_LITTLE_ENDIAN,
],
"mips": [UC_MIPS_REG_PC, UC_ARCH_MIPS, UC_MODE_MIPS32 | UC_MODE_BIG_ENDIAN],
"mipsel" : [ UC_MIPS_REG_PC, UC_ARCH_MIPS, UC_MODE_MIPS32 | UC_MODE_LITTLE_ENDIAN ],
"mipsel": [
UC_MIPS_REG_PC,
UC_ARCH_MIPS,
UC_MODE_MIPS32 | UC_MODE_LITTLE_ENDIAN,
],
}
return (arch_map[arch_str][1], arch_map[arch_str][2])
@ -487,7 +588,7 @@ class AflUnicornEngine(Uc):
"pc": UC_ARM_REG_PC,
"sp": UC_ARM_REG_SP,
"lr": UC_ARM_REG_LR,
"cpsr": UC_ARM_REG_CPSR
"cpsr": UC_ARM_REG_CPSR,
},
"arm64": {
"x0": UC_ARM64_REG_X0,
@ -562,8 +663,8 @@ class AflUnicornEngine(Uc):
"fp": UC_MIPS_REG_FP,
"ra": UC_MIPS_REG_RA,
"hi": UC_MIPS_REG_HI,
"lo": UC_MIPS_REG_LO
}
"lo": UC_MIPS_REG_LO,
},
}
return registers[arch]
@ -610,18 +711,17 @@ class AflUnicornEngine(Uc):
"d29": UC_ARM_REG_D29,
"d30": UC_ARM_REG_D30,
"d31": UC_ARM_REG_D31,
"fpscr": UC_ARM_REG_FPSCR
"fpscr": UC_ARM_REG_FPSCR,
}
}
return registers[arch];
return registers[arch]
# ---------------------------
# Callbacks for tracing
# 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):
if CAPSTONE_EXISTS == 1:
# If Capstone is installed then we'll dump disassembly, otherwise just dump the binary.
@ -651,11 +751,23 @@ class AflUnicornEngine(Uc):
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))
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))
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))
@ -664,15 +776,27 @@ class AflUnicornEngine(Uc):
def __trace_mem_access(self, uc, access, address, size, value, user_data):
if access == UC_MEM_WRITE:
print(" >>> Write: addr=0x{0:016x} size={1} data=0x{2:016x}".format(address, size, value))
print(
" >>> Write: addr=0x{0:016x} size={1} data=0x{2:016x}".format(
address, size, value
)
)
else:
print(" >>> Read: addr=0x{0:016x} size={1}".format(address, size))
def __trace_mem_invalid_access(self, uc, access, address, size, value, user_data):
if access == UC_MEM_WRITE_UNMAPPED:
print(" >>> INVALID Write: addr=0x{0:016x} size={1} data=0x{2:016x}".format(address, size, value))
print(
" >>> INVALID Write: addr=0x{0:016x} size={1} data=0x{2:016x}".format(
address, size, value
)
)
else:
print(" >>> INVALID Read: addr=0x{0:016x} size={1}".format(address, size))
print(
" >>> INVALID Read: addr=0x{0:016x} size={1}".format(
address, size
)
)
def bit_size_arch(self):
arch = self.get_arch()