example/ida/depgraph.py
from __future__ import print_function
from builtins import map
from builtins import range
import os
import tempfile
from future.utils import viewitems, viewvalues
import idautils
import idc
import ida_funcs
import ida_kernwin
from miasm.core.bin_stream_ida import bin_stream_ida
from miasm.core.asmblock import *
from miasm.expression import expression as m2_expr
from miasm.core.locationdb import LocationDB
from miasm.expression.simplifications import expr_simp
from miasm.analysis.depgraph import DependencyGraph
from miasm.ir.ir import AssignBlock, IRBlock
from utils import guess_machine
class depGraphSettingsForm(ida_kernwin.Form):
def __init__(self, lifter_model_call, ircfg, mn):
self.lifter_model_call = lifter_model_call
self.ircfg = ircfg
self.mn = mn
self.stk_args = {'ARG%d' % i:i for i in range(10)}
self.stk_unalias_force = False
self.address = idc.get_screen_ea()
cur_block = None
for loc_key in ircfg.getby_offset(self.address):
block = ircfg.get_block(loc_key)
offset = self.ircfg.loc_db.get_location_offset(block.loc_key)
if offset is not None:
# Only one block non-generated
assert cur_block is None
cur_block = block
assert cur_block is not None
line_nb = None
for line_nb, assignblk in enumerate(cur_block):
if assignblk.instr.offset == self.address:
break
assert line_nb is not None
cur_loc_key = str(cur_block.loc_key)
loc_keys = sorted(map(str, ircfg.blocks))
regs = sorted(lifter_model_call.arch.regs.all_regs_ids_byname)
regs += list(self.stk_args)
reg_default = regs[0]
for i in range(10):
opnd = idc.print_operand(self.address, i).upper()
if opnd in regs:
reg_default = opnd
break
ida_kernwin.Form.__init__(self,
r"""BUTTON YES* Launch
BUTTON CANCEL NONE
Dependency Graph Settings
Track the element:
<Before the line:{rBeforeLine}>
<After the line:{rAfterLine}>
<At the end of the basic block:{rEndBlock}>{cMode}>
<Target basic block:{cbBBL}>
<Register to track:{cbReg}>
<Line number:{iLineNb}>
INFO: To track stack argument number n, use "ARGn"
Method to use:
<Follow Memory:{rNoMem}>
<Follow Call:{rNoCall}>
<Implicit dependencies:{rImplicit}>
<Unalias stack:{rUnaliasStack}>{cMethod}>
<Highlight color:{cColor}>
""", {
'cbReg': ida_kernwin.Form.DropdownListControl(
items=regs,
readonly=False,
selval=reg_default),
'cMode': ida_kernwin.Form.RadGroupControl(
("rBeforeLine", "rAfterLine", "rEndBlock")),
'cMethod': ida_kernwin.Form.ChkGroupControl(
("rNoMem", "rNoCall", "rImplicit", "rUnaliasStack")),
'iLineNb': ida_kernwin.Form.NumericInput(
tp=ida_kernwin.Form.FT_RAWHEX,
value=line_nb),
'cbBBL': ida_kernwin.Form.DropdownListControl(
items=loc_keys,
readonly=False,
selval=cur_loc_key),
'cColor': ida_kernwin.Form.ColorInput(value=0xc0c020),
})
self.Compile()
@property
def loc_key(self):
value = self.cbBBL.value
for real_loc_key in self.ircfg.blocks:
if str(real_loc_key) == value:
return real_loc_key
raise ValueError("Bad loc_key")
@property
def line_nb(self):
value = self.iLineNb.value
mode = self.cMode.value
if mode == 0:
return value
elif mode == 1:
return value + 1
else:
return len(self.ircfg.blocks[self.loc_key])
@property
def elements(self):
value = self.cbReg.value
if value in self.stk_args:
line = self.ircfg.blocks[self.loc_key][self.line_nb].instr
arg_num = self.stk_args[value]
stk_high = m2_expr.ExprInt(idc.get_spd(line.offset), lifter.sp.size)
stk_off = m2_expr.ExprInt(self.lifter_model_call.sp.size // 8 * arg_num, lifter.sp.size)
element = m2_expr.ExprMem(self.mn.regs.regs_init[lifter.sp] + stk_high + stk_off, self.lifter_model_call.sp.size)
element = expr_simp(element)
# Force stack unaliasing
self.stk_unalias_force = True
elif value:
element = self.lifter_model_call.arch.regs.all_regs_ids_byname.get(value, None)
else:
raise ValueError("Unknown element '%s'!" % value)
return set([element])
@property
def depgraph(self):
value = self.cMethod.value
return DependencyGraph(self.ircfg,
implicit=value & 4,
follow_mem=value & 1,
follow_call=value & 2)
@property
def unalias_stack(self):
return self.cMethod.value & 8 or self.stk_unalias_force
@property
def color(self):
return self.cColor.value
def clean_lines():
"Remove previous comments"
global comments
for offset in comments:
idc.set_color(offset, idc.CIC_ITEM, 0xffffff)
idc.set_cmt(offset, "", 0)
comments = {}
def treat_element():
"Display an element"
global graphs, comments, sol_nb, settings, addr, lifter, ircfg
try:
graph = next(graphs)
except StopIteration:
comments = {}
print("Done: %d solutions" % (sol_nb))
return
sol_nb += 1
print("Get graph number %02d" % sol_nb)
filename = os.path.join(tempfile.gettempdir(), "solution_0x%08x_%02d.dot" % (addr, sol_nb))
print("Dump the graph to %s" % filename)
open(filename, "w").write(graph.graph.dot())
for node in graph.relevant_nodes:
try:
offset = ircfg.blocks[node.loc_key][node.line_nb].instr.offset
except IndexError:
print("Unable to highlight %s" % node)
continue
comments[offset] = comments.get(offset, []) + [node.element]
idc.set_color(offset, idc.CIC_ITEM, settings.color)
if graph.has_loop:
print('Graph has dependency loop: symbolic execution is inexact')
else:
print("Possible value: %s" % next(iter(viewvalues(graph.emul(lifter)))))
for offset, elements in viewitems(comments):
idc.set_cmt(offset, ", ".join(map(str, elements)), 0)
def next_element():
"Display the next element"
clean_lines()
treat_element()
def launch_depgraph():
global graphs, comments, sol_nb, settings, addr, lifter, ircfg
# Get the current function
addr = idc.get_screen_ea()
func = ida_funcs.get_func(addr)
# Init
machine = guess_machine(addr=func.start_ea)
mn, dis_engine, lifter_model_call = machine.mn, machine.dis_engine, machine.lifter_model_call
bs = bin_stream_ida()
loc_db = LocationDB()
mdis = dis_engine(bs, loc_db=loc_db, dont_dis_nulstart_bloc=True)
lifter = lifter_model_call(loc_db)
# Populate symbols with ida names
for ad, name in idautils.Names():
if name is None:
continue
loc_db.add_location(name, ad)
asmcfg = mdis.dis_multiblock(func.start_ea)
# Generate IR
ircfg = lifter.new_ircfg_from_asmcfg(asmcfg)
# Get settings
settings = depGraphSettingsForm(lifter, ircfg, mn)
settings.Execute()
loc_key, elements, line_nb = settings.loc_key, settings.elements, settings.line_nb
# Simplify assignments
for irb in list(viewvalues(ircfg.blocks)):
irs = []
offset = loc_db.get_location_offset(irb.loc_key)
fix_stack = offset is not None and settings.unalias_stack
for assignblk in irb:
if fix_stack:
stk_high = m2_expr.ExprInt(idc.get_spd(assignblk.instr.offset), lifter.sp.size)
fix_dct = {lifter.sp: mn.regs.regs_init[lifter.sp] + stk_high}
new_assignblk = {}
for dst, src in viewitems(assignblk):
if fix_stack:
src = src.replace_expr(fix_dct)
if dst != lifter.sp:
dst = dst.replace_expr(fix_dct)
dst, src = expr_simp(dst), expr_simp(src)
new_assignblk[dst] = src
irs.append(AssignBlock(new_assignblk, instr=assignblk.instr))
ircfg.blocks[irb.loc_key] = IRBlock(irb.loc_db, irb.loc_key, irs)
# Get dependency graphs
dg = settings.depgraph
graphs = dg.get(loc_key, elements, line_nb,
set([loc_db.get_offset_location(func.start_ea)]))
# Display the result
comments = {}
sol_nb = 0
# Register and launch
ida_kernwin.add_hotkey("Shift-N", next_element)
treat_element()
if __name__ == "__main__":
launch_depgraph()