220 lines
7.8 KiB
Python
Executable File
220 lines
7.8 KiB
Python
Executable File
from immlib import *
|
|
|
|
class BasicBlock:
|
|
|
|
def __init__(self, start_addr=None, start_op=None):
|
|
self.start_addr = start_addr
|
|
self.end_addr = None
|
|
|
|
self.start_op = start_op
|
|
self.end_op = None
|
|
|
|
# A list of instruction objects in the current Basic Block in
|
|
# address order
|
|
self.ops = [start_op]
|
|
|
|
def __iter__(self):
|
|
for op in self.ops:
|
|
yield op
|
|
|
|
def __str__(self):
|
|
return ["\n".join([op.address for op in self.ops])]
|
|
|
|
def log(self, imm):
|
|
imm.log("BB @ %s" % hex(self.start_addr))
|
|
for op in self.ops:
|
|
imm.log(" %s" % hex(op.address), op.address)
|
|
|
|
class BasicBlockGraph:
|
|
|
|
def __init__(self):
|
|
# Dictionary of addresses to basic blocks
|
|
self.basic_blocks = {}
|
|
# The outgoing edges for a basic block identified by its
|
|
# starting address. A dictionary of addresses to a list of
|
|
# addresses of basic blocks that the address may branch to
|
|
self.bb_out_edges = {}
|
|
# The incoming edges for a basic block identified by its
|
|
# starting address. As above.
|
|
self.bb_in_edges = {}
|
|
|
|
def __iter__(self):
|
|
for addr, bb in self.basic_blocks.items():
|
|
yield bb
|
|
|
|
def log(self, imm):
|
|
imm.log("** BASIC BLOCKS **")
|
|
for bb_addr in self.basic_blocks:
|
|
bb = self.basic_blocks[bb_addr]
|
|
imm.log("BB start %s, end %s" % (hex(bb.start_addr),
|
|
hex(bb.end_addr)),
|
|
bb.start_addr)
|
|
|
|
imm.log("** OUT EDGES **")
|
|
for bb_addr in self.bb_out_edges:
|
|
out = []
|
|
[out.append(hex(addr)) \
|
|
for addr in self.bb_out_edges[bb_addr]]
|
|
imm.log("%s -> %s" % (hex(bb_addr), ','.join(out)),
|
|
bb_addr)
|
|
|
|
imm.log("** IN EDGES **")
|
|
for bb_addr in self.bb_in_edges:
|
|
out = []
|
|
[out.append(hex(addr)) \
|
|
for addr in self.bb_in_edges[bb_addr]]
|
|
imm.log("%s -> %s" % (hex(bb_addr), ','.join(out)),
|
|
bb_addr)
|
|
|
|
class CodeStructureAnalyzer:
|
|
|
|
def __init__(self, imm):
|
|
self.imm = imm
|
|
|
|
def buildIncomingGraph(self, basic_blocks, out_graph):
|
|
"""
|
|
This function takes a dictionary of basic blocks and a
|
|
dictionary describing outgoing edges from the basic blocks
|
|
and constructs a dictionary of the incoming edges to each
|
|
basic block.
|
|
"""
|
|
|
|
bb_in_edges = {}
|
|
|
|
for out_bb_addr in out_graph:
|
|
for in_bb_addr in out_graph[out_bb_addr]:
|
|
if in_bb_addr in bb_in_edges:
|
|
bb_in_edges[in_bb_addr].append(out_bb_addr)
|
|
else:
|
|
bb_in_edges[in_bb_addr] = [out_bb_addr]
|
|
|
|
return bb_in_edges
|
|
|
|
def getBasicBlockGraph(self, start_addr):
|
|
work_list = []
|
|
bb_out_edges = {}
|
|
bb_in_edges = {}
|
|
basic_blocks = {}
|
|
|
|
next_addr = start_addr
|
|
next_op = self.imm.disasm(next_addr, DISASM_FILE)
|
|
next_op_type = next_op.getOpType()[0] & DEC_TYPEMASK
|
|
|
|
curr_bb = BasicBlock(start_addr=next_addr, start_op=next_op)
|
|
basic_blocks[next_addr] = curr_bb
|
|
|
|
while True:
|
|
if curr_bb.end_addr is not None:
|
|
try:
|
|
curr_bb = work_list.pop()
|
|
prev_addr = curr_bb.start_addr
|
|
|
|
prev_op = curr_bb.start_op
|
|
next_addr = prev_addr
|
|
except IndexError:
|
|
break
|
|
else:
|
|
prev_addr = next_addr
|
|
prev_op = next_op
|
|
|
|
next_addr = prev_addr + prev_op.getOpSize()
|
|
|
|
next_op = self.imm.disasm(next_addr, DISASM_FILE)
|
|
next_op_type = next_op.getOpType()[0] & DEC_TYPEMASK
|
|
decode_info = Decode(next_addr)
|
|
|
|
if decode_info.isJmpDestination(next_addr) or \
|
|
decode_info.isCallDestination(next_addr):
|
|
|
|
if next_addr == curr_bb.start_addr:
|
|
continue
|
|
|
|
curr_bb.end_addr = prev_addr
|
|
curr_bb.end_op = prev_op
|
|
|
|
if curr_bb.start_addr in bb_out_edges:
|
|
bb_out_edges[curr_bb.start_addr].append(next_addr)
|
|
else:
|
|
bb_out_edges[curr_bb.start_addr] = [next_addr]
|
|
|
|
if not next_addr in basic_blocks:
|
|
new_bb = BasicBlock(start_addr=next_addr,
|
|
start_op=next_op)
|
|
basic_blocks[next_addr] = new_bb
|
|
work_list.append(new_bb)
|
|
continue
|
|
|
|
if next_op.isJmp():
|
|
curr_bb.ops.append(next_op)
|
|
curr_bb.end_addr = next_addr
|
|
curr_bb.end_op = next_op
|
|
|
|
jmp_target = next_op.getJmpAddr()
|
|
|
|
if curr_bb.start_addr in bb_out_edges:
|
|
bb_out_edges[curr_bb.start_addr].append(jmp_target)
|
|
else:
|
|
bb_out_edges[curr_bb.start_addr] = [jmp_target]
|
|
|
|
if not jmp_target in basic_blocks:
|
|
jmp_op = self.imm.disasm(jmp_target, DISASM_FILE)
|
|
|
|
new_bb = BasicBlock(start_addr=jmp_target,
|
|
start_op=jmp_op)
|
|
basic_blocks[jmp_target] = new_bb
|
|
|
|
work_list.append(new_bb)
|
|
continue
|
|
elif next_op.isConditionalJmp():
|
|
curr_bb.ops.append(next_op)
|
|
curr_bb.end_addr = next_addr
|
|
curr_bb.end_op = next_op
|
|
|
|
jmp_target = next_op.getJmpAddr()
|
|
if curr_bb.end_addr in bb_out_edges:
|
|
bb_out_edges[curr_bb.start_addr].append(jmp_target)
|
|
else:
|
|
bb_out_edges[curr_bb.start_addr] = [jmp_target]
|
|
|
|
if not jmp_target in basic_blocks:
|
|
jmp_op = self.imm.disasm(jmp_target, DISASM_FILE)
|
|
|
|
new_bb = BasicBlock(start_addr=jmp_target,
|
|
start_op=jmp_op)
|
|
basic_blocks[jmp_target] = new_bb
|
|
|
|
work_list.append(new_bb)
|
|
|
|
fall_thru_addr = next_addr + next_op.getSize()
|
|
if curr_bb.start_addr in bb_out_edges:
|
|
bb_out_edges[curr_bb.start_addr].append(fall_thru_addr)
|
|
else:
|
|
bb_out_edges[curr_bb.start_addr] = [fall_thru_addr]
|
|
|
|
if not fall_thru_addr in basic_blocks:
|
|
fall_thru_op = self.imm.disasm(fall_thru_addr, DISASM_FILE)
|
|
|
|
new_bb = BasicBlock(start_addr=fall_thru_addr,
|
|
start_op=fall_thru_op)
|
|
basic_blocks[fall_thru_addr] = new_bb
|
|
|
|
work_list.append(new_bb)
|
|
continue
|
|
elif next_op.isRet():
|
|
curr_bb.ops.append(next_op)
|
|
curr_bb.end_addr = next_addr
|
|
curr_bb.end_op = next_op
|
|
continue
|
|
else:
|
|
curr_bb.ops.append(next_op)
|
|
|
|
bb_in_edges = self.buildIncomingGraph(basic_blocks,
|
|
bb_out_edges)
|
|
|
|
bb_graph = BasicBlockGraph()
|
|
bb_graph.basic_blocks = basic_blocks
|
|
bb_graph.bb_out_edges = bb_out_edges
|
|
bb_graph.bb_in_edges = bb_in_edges
|
|
|
|
return bb_graph
|