mirror of
https://github.com/cea-sec/miasm
synced 2026-08-17 10:26:03 -04:00
276 lines
8.7 KiB
Python
276 lines
8.7 KiB
Python
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()
|