def main():
info("Showing number of tainted hit per image.")
# filename = idaapi.askfile_c(0, "pintool.log", "Trace file to load.")
filename = """/Users/anon/workspace/instrumentation/CodeCoverage/trace.log"""
if filename is None:
info("Aborting ...")
return
# Get loaded binary name
image_name = idc.GetInputFile().lower()
info("IDB binary name '%s'" % image_name)
# Get the image base
image_base = idaapi.get_imagebase()
info("IDB binary base 0x%.16x" % image_base)
# Load the trace file.
trace = TraceReader(filename)
trace.parse(match_events="HL")
# The IDB matches one and only one loaded image. Find it.
traced_image = None
for image in trace.getLoadedImages():
if image_name.lower() in os.path.basename(image.name).lower():
traced_image = image
break
if not traced_image:
info("Error, could not find the traced image '%s' in the list of loaded images." % image_name)
return set()
# Collect all the hits that belong to the ida database.
hits = set(filter(lambda x: traced_image.contains(x.address), trace.getBasicBlockHits()))
hits = set(map(lambda x: image_base + traced_image.get_offset(x.address), hits))
reached_functions = set()
for hit in hits:
f = idaapi.get_func(hit)
if f:
if f.startEA in reached_functions:
continue
info("Reached -> %s" % (GetFunctionName(f.startEA)))
reached_functions.add(f.startEA)
if idaapi.askyn_c(1, "Do you want to mark all the FUNCTIONS reached?") == 1:
FUNCTION_COLOR = rgb_to_bgr(0xBCF5D1)
for function in reached_functions:
SetFunctionColor(function, FUNCTION_COLOR)
if idaapi.askyn_c(0, "Do you want to mark all the BASIC BLOCKS reached?") == 1:
BBLOCK_COLOR = rgb_to_bgr(0xf2ddda)
for hit in hits:
SetBasicBlockColor(hit, BBLOCK_COLOR)
def OnCommand(self, n, cmd_id):
if cmd_id == self.cmd_show_reasons:
match = self.items[n]
reasons = match[len(match)-1]
msg = "\n".join(reasons)
info(msg)
elif cmd_id == self.cmd_import_all:
if askyn_c(0, "HIDECANCEL\nDo you really want to import all matched functions as well as struct, union, enum and typedef definitions?") == 1:
import_items = []
for item in self.items:
src_id, src_name, bin_ea = int(item[1]), item[2], int(item[3], 16)
import_items.append([src_id, src_name, bin_ea])
self.importer.import_items(import_items)
elif cmd_id == self.cmd_import_selected:
if len(self.selected_items) == 1 or askyn_c(1, "HIDECANCEL\nDo you really want to import the selected functions?") == 1:
import_items = []
for index in self.selected_items:
item = self.items[index]
src_id, src_name, bin_ea = int(item[1]), item[2], int(item[3], 16)
import_items.append([src_id, src_name, bin_ea])
import_definitions = askyn_c(0, "HIDECANCEL\nDo you also want to import all struct, union, enum and typedef definitions?") == 1
self.importer.import_items(import_items, import_definitions = import_definitions)
elif cmd_id == self.cmd_diff_c:
html_diff = CHtmlDiff()
item = self.items[n]
src_id = long(item[1])
cur = self.differ.db.cursor()
sql = "select source from src.functions where id = ?"
cur.execute(sql, (src_id,))
row = cur.fetchone()
cur.close()
if not row:
Warning("Cannot find the source function.")
return False
ea = long(item[3], 16)
proto = self.differ.decompile_and_get(ea)
if not proto:
Warning("Cannot decompile function 0x%08x" % ea)
return False
buf1 = indent_source(row[0])
buf2 = proto
buf2 += "\n".join(self.differ.pseudo[ea])
new_buf = indent_source(buf2)
src = html_diff.make_file(new_buf.split("\n"), buf1.split("\n"))
title = "Diff pseudo-source %s - %s" % (item[2], item[4])
cdiffer = CHtmlViewer()
cdiffer.Show(src, title)
def ask_to_checkout_modified_files(self):
modified_objects = ""
checkout_head = False
for modified_object in self.repo.get_modified_objects():
# check if modified file is original idb
original_idb = get_original_idb_name(idc.GetIdbPath())
if original_idb == modified_object:
# original idb modification detected, create a backup
try:
shutil.copy(
original_idb,
"%s_bkp_%s" % (original_idb, time.ctime().replace(
" ", "_").replace(":", "_")))
except Exception as e:
traceback.print_exc()
raise e
checkout_head = True
else:
modified_objects += "%s\n" % modified_object
if len(modified_objects) > 0:
message = "%s\nhas been modified, this is not normal, do you want to checkout these files ? "
message += "(Rebasing will be disabled if you answer no)"
if idaapi.askyn_c(True, message % modified_objects):
self.repo.checkout_head()
else:
self.repo_auto_sync = False
return
if checkout_head:
# checkout silently
self.repo.checkout_head()
def retrieve_reset(self, *args):
if not idaapi.askyn_c(
False,
"All your local changes will be lost !\nDo you really want to proceed ?"
):
return
# disable all yaco hooks
self.YaCoUI.unhook()
# create a backup of current idb
yatools.copy_idb_to_local_file(
"_bkp_%s" % time.ctime().replace(" ", "_").replace(":", "_"))
# delete all modified objects
self.repo_manager.repo.checkout_head()
# get reset
self.repo_manager.fetch_origin()
self.repo_manager.rebase_from_origin()
original_idb_name = yatools.get_original_idb_name(idc.GetIdbPath())
# remove current idb
os.remove(idc.GetIdbPath())
# recreate local idb
shutil.copy(original_idb_name,
yatools.get_local_idb_name(original_idb_name))
# local should not be overwritten, so we have to close IDA brutally !
idaapi.cvar.database_flags |= idaapi.DBFL_KILL
idc.Warning("Force pull complete, you can restart IDA")
idc.Exit(0)
def different_versions(self):
ret = False
db = sqlite3.connect(self.db_filename)
cur = db.cursor()
sql = "select value, status from version"
try:
cur.execute(sql)
row = cur.fetchone()
if row:
version = row[0]
status = row[1]
if version != VERSION_VALUE:
msg = "HIDECANCEL\nDatabase version (%s) is different to current version (%s).\n"
msg += "Do you want to re-create the database?"
msg += "\n\nNOTE: Selecting 'NO' will try to use the non updated database."
ret = askyn_c(0, msg % (version, VERSION_VALUE)) == 1
elif status != "done":
ret = True
else:
ret = False
except:
print("Error checking version: %s" % str(sys.exc_info()[1]))
ret = True
cur.close()
return ret
def OnDeleteLine(self, n):
ans = idaapi.askyn_c(
1,
"HIDECANCEL\nAre you sure you want to delete function [%s] @ [%s]?"
% (self.items[n][3], self.items[n][4]))
if ans == 1:
asms = Assembler.LoadSavedAssemblers()
item = int(self.items[n][2], 16)
if asms != None and len(asms.keys()) > 0:
for asm in asms.itervalues():
if asm.functions.has_key(item):
print "Removed [%08x]!" % item
del asm.functions[item]
asm.SaveState()
opty_ea = int(self.items[n][4], 16)
print "set_name[%08x]" % opty_ea
idc.MakeComm(opty_ea, "")
idaapi.set_name(opty_ea, "")
idc.DelFunction(opty_ea)
comment = idc.Comment(item)
comment = re.sub(r"(?i)OPTY@\[[\d+abcdef]+\];\s*", "", comment)
idc.MakeComm(item, comment)
self.populate_items()
return n
def OnEditLine(self, n):
#this is "Delete Segment"
ans = idaapi.askyn_c(
1,
"HIDECANCEL\nAre you sure you want to delete segment and all optimized data from disk?"
)
if ans == 1:
opty_dir = idc.GetIdbPath()
opty_dir = opty_dir[:opty_dir.rfind(
os.sep)] + os.sep + "optimice_%s" % idc.GetInputFile()
print opty_dir
if not os.path.isdir(opty_dir):
print ">GUI_FunctionManager:OnEditLine - Error [%s] not a directory!" % opty_dir
return 0
shutil.rmtree(opty_dir)
print ">GUI_FunctionManager: Optimice directory deleted: [%s]" % opty_dir
idc.SegDelete(int(self.items[n][0], 16), 0)
print ">GUI_FunctionManager: Optimice segment deleted: [%s]" % self.items[
n][0]
self.populate_items()
return 0
def main():
try:
res = idaapi.askyn_c(1, "Do you want to export (YES) or import (NO)?")
exporter = CFunctionsMatcher()
if res == 1:
exporter.export()
elif res == 0:
exporter.doImport()
except:
print "Error:", sys.exc_info()[1]
traceback.print_exc(file=sys.stdout)
def idb_to_mega(filename):
"""Convert IDB to mega file."""
mode = "w"
# File already exists
if os.path.exists(filename):
ask = idaapi.askyn_c(False, "Append result to existing file?")
if ask == -1:
return
elif ask == 0:
if 1 != idaapi.askyn_c(False, "Overwrite existing file?"):
return
elif ask == 1:
mode = "a+"
# Process functions
with open(filename, mode) as f:
for i in range(idaapi.get_func_qty()):
line = process_func(i)
if line:
line += "\n"
f.write(line)
def clear_bap_comments(self):
"""Ask user for confirmation and then clear (BAP ..) comments."""
if idaapi.askyn_c(ASKBTN_YES,
"Delete all `BAP: ..` comments?") != ASKBTN_YES:
return
for ea in ida.addresses(): # TODO: store actually commented addresses
comm = idaapi.get_cmt(ea, 0)
if comm and comm.startswith('BAP:'):
idaapi.set_cmt(ea, '', 0)
def clear_bap_comments(self):
"""Ask user for confirmation and then clear (BAP ..) comments."""
if idaapi.askyn_c(ASKBTN_YES,
"Delete all `BAP: ..` comments?") != ASKBTN_YES:
return
for ea in ida.addresses(): # TODO: store actually commented addresses
comm = idaapi.get_cmt(ea, 0)
if comm and comm.startswith('BAP:'):
idaapi.set_cmt(ea, '', 0)
def run(self):
"run BAP instance"
if len(BapIda.instances) > 0:
answer = idaapi.askyn_c(
idaapi.ASKBTN_YES,
"Previous instances of BAP didn't finish yet.\
Do you really want to start a new one?".
format(len(BapIda.instances)))
if answer == idaapi.ASKBTN_YES:
self._do_run()
else:
self._do_run()
idc.Message("BAP> total number of running instances: {0}\n".
format(len(BapIda.instances)))
def run(self):
"run BAP instance"
if len(BapIda.instances) > 0:
answer = idaapi.askyn_c(
idaapi.ASKBTN_YES,
"Previous instances of BAP didn't finish yet.\
Do you really want to start a new one?".format(
len(BapIda.instances)))
if answer == idaapi.ASKBTN_YES:
self._do_run()
else:
self._do_run()
idc.Message("BAP> total number of running instances: {0}\n".format(
len(BapIda.instances)))
def retrieve_reset(self, *args):
if not idaapi.askyn_c(
False,
"All your local changes will be lost !\nDo you really want to proceed ?"
):
return
# disable all yaco hooks
self.YaCoUI.unhook()
self.repo_manager.discard_and_pull_idb()
# current idb should not be overwritten, so we have to close IDA brutally !
idaapi.set_database_flag(idaapi.DBFL_KILL)
idc.Warning("Force pull complete, you can restart IDA")
idc.Exit(0)
def different_versions(self):
ret = False
db = sqlite3.connect(self.db_filename)
cur = db.cursor()
sql = "select value from version"
cur.execute(sql)
row = cur.fetchone()
if row:
version = row[0]
if version != VERSION_VALUE:
msg = "HIDECANCEL\nDatabase version (%s) is different to current version (%s).\n"
msg += "Do you want to re-create the database?"
msg += "\n\nNOTE: Selecting 'NO' will try to use the non updated database."
if askyn_c(0, msg % (version, VERSION_VALUE)) == 1:
ret = True
cur.close()
return ret
def assemble(self, ea, asm, save_state=True, opt_fix=True, opt_nop=True):
""" Assemble into memory """
#Fixup the assemble
if opt_fix:
regvars = self._getregvars(ea)
parts_arr = [self.PART_RE.split(i) for i in asm]
asm = []
for parts in parts_arr:
asm.append(self._fixup(parts, regvars))
#Assemble to a string
success, data = idautils.Assemble(ea, asm)
if not success:
return success, data
blob = ''.join(data)
if len(blob) > instr_size(ea):
if idaapi.askyn_c(
0,
"The assembled instruction is bigger than the current instruction. This will clobber following instructions. Continue?"
) != 1:
return
#Pad the blob with nops
if opt_nop:
nsuccess, nop_instr = idautils.Assemble(ea, 'nop')
if not nsuccess:
return nsuccess, nop_instr
i = ea
while i < ea + len(blob):
i += instr_size(i)
#Only pad if we trashed the next instruction
sz_diff = (i - (ea + len(blob))) / len(nop_instr)
blob += nop_instr * sz_diff
#Write out the data
old = read_data(ea, len(blob))
if save_state:
self._pushundo([(ea, old)])
self.redo_buffer = []
write_data(ea, blob)
return success, old
def assemble(self, ea, asm, save_state=True, opt_fix=True, opt_nop=True):
""" Assemble into memory """
#Fixup the assemble
if opt_fix:
regvars = self._getregvars(ea)
parts_arr = [self.PART_RE.split(i) for i in asm]
asm = []
for parts in parts_arr:
asm.append(self._fixup(parts, regvars))
#Assemble to a string
success, data = idautils.Assemble(ea, asm)
if not success:
return success, data
blob = ''.join(data)
if len(blob) > instr_size(ea):
if idaapi.askyn_c(0, "The assembled instruction is bigger than the current instruction. This will clobber following instructions. Continue?") != 1:
return
#Pad the blob with nops
if opt_nop:
nsuccess, nop_instr = idautils.Assemble(ea, 'nop')
if not nsuccess:
return nsuccess, nop_instr
i = ea
while i < ea + len(blob):
i += instr_size(i)
#Only pad if we trashed the next instruction
sz_diff = (i - (ea + len(blob))) / len(nop_instr)
blob += nop_instr * sz_diff
#Write out the data
old = read_data(ea, len(blob))
if save_state:
self._pushundo(
[(ea, old)]
)
self.redo_buffer = []
write_data(ea, blob)
return success, old
def OnEditLine(self, n):
#this is "Delete Segment"
ans = idaapi.askyn_c(1, "HIDECANCEL\nAre you sure you want to delete segment and all optimized data from disk?")
if ans == 1:
opty_dir = idc.GetIdbPath()
opty_dir = opty_dir[:opty_dir.rfind(os.sep)] + os.sep + "optimice_%s" % idc.GetInputFile()
print opty_dir
if not os.path.isdir(opty_dir):
print ">GUI_FunctionManager:OnEditLine - Error [%s] not a directory!" % opty_dir
return 0
shutil.rmtree(opty_dir)
print ">GUI_FunctionManager: Optimice directory deleted: [%s]" % opty_dir
idc.SegDelete(int(self.items[n][0], 16), 0)
print ">GUI_FunctionManager: Optimice segment deleted: [%s]" % self.items[n][0]
self.populate_items()
return 0
def reload_gui_info(self, from_arena_cb=False):
if self.heap is None:
return
try:
if not misc.is_process_suspended():
answer = idaapi.askyn_c(
idaapi.ASKBTN_YES,
"HIDECANCEL\nThe process must be suspended to reload the info.\n\
Do you want to suspend it?")
if answer == idaapi.ASKBTN_NO:
return
if not idaapi.suspend_process():
warning("Unable to suspend the process")
return
idaapi.refresh_debugger_memory()
if not self.heap.get_heap_base():
self.show_warning("Heap not initialized")
return
if not config.libc_base:
self.show_warning("Unable to resolve glibc base address.")
return
self.hide_warning()
self.arenas_widget.setVisible(True)
if not from_arena_cb:
self.populate_arenas()
self.arena_widget.populate_table()
self.tcache_widget.populate_table()
self.bins_widget.populate_tables()
except Exception as e:
self.show_warning(e.message)
idaapi.warning(traceback.format_exc())
def OnDeleteLine(self, n):
ans = idaapi.askyn_c(1, "HIDECANCEL\nAre you sure you want to delete function [%s] @ [%s]?" % (self.items[n][3], self.items[n][4]) )
if ans == 1:
asms = Assembler.LoadSavedAssemblers()
item = int(self.items[n][2], 16)
if asms != None and len(asms.keys()) > 0:
for asm in asms.itervalues():
if asm.functions.has_key(item):
print "Removed [%08x]!" % item
del asm.functions[item]
asm.SaveState()
opty_ea = int(self.items[n][4], 16)
print "set_name[%08x]" % opty_ea
idc.MakeComm(opty_ea, "")
idaapi.set_name(opty_ea, "")
idc.DelFunction(opty_ea)
comment = idc.Comment(item)
comment = re.sub(r"(?i)OPTY@\[[\d+abcdef]+\];\s*", "", comment)
idc.MakeComm(item, comment)
self.populate_items()
return n
def NasmAssemble(self, function_ea, write_ea):
dir = self.opty_dir
nasm = "C:\\Program Files\\nasm\\nasm.exe"
arg1 = "f_%08x.asm" % function_ea
arg2 = "-o f_%08x.o" % function_ea
arg3 = "-l f_%08x.lst" % function_ea
arg4 = "-Ox"
orig_dir = os.getcwd()
os.chdir(dir)
idc.Batch(0)
while 1:
try:
p = subprocess.Popen([nasm, arg1, arg2, arg3, arg4],
stdout=subprocess.PIPE,
stderr=subprocess.PIPE,
stdin=subprocess.PIPE)
o, e = p.communicate()
if o != "": print o
if e != "": print e
fop = open("f_%08x.o" % function_ea, "rb")
ans = idaapi.askyn_c(
0,
"HIDECANCEL\nDo you want to manually edit function before writing to IDA?"
)
if ans == 1:
os.startfile(arg1, "open")
idaapi.warning(
"Press OK button when you're done editing file.")
fop.close()
continue
else:
idc.Batch(1)
break
except:
error_msg = '\n'.join(e.split("\n")[:15])
os.startfile(arg1, "open")
ans = idaapi.askyn_c(
1, """HIDECANCEL\nNASM failed to assemble [f_%08x.o] file.
You can manually edit and NASM this file [f_%08x.asm] and click Yes when you're done.
File is located in directory where your IDB is located.
If you want to skip this function press No.
Nasm output:
%s""" % (function_ea, function_ea, error_msg))
if ans == 1:
continue
else:
os.chdir(orig_dir)
idc.Batch(1)
return None
os.chdir(orig_dir)
print ">>>Writing function [%08x] @ [%08x]" % (function_ea, write_ea)
data = fop.read()
data_len = len(data)
for offset in xrange(0, data_len):
idc.PatchByte(write_ea + offset, ord(data[offset]))
fop.close()
idc.MakeCode(write_ea)
fp = open("%s\\f_%08x.lst" % (dir, function_ea), "r")
asm_lst = fp.read()
base_addr = re.search(r"ORG ([\dABCDEF]+)H", asm_lst).group(1)
base_addr = int(base_addr, 16)
for jt in self.jmp_table_refs:
m = re.search(r"\s*\d+\s+([\dABCDEF]{8}).*?%s" % re.escape(jt),
asm_lst, re.IGNORECASE)
if m != None:
jt_ea = int(m.group(1), 16)
jt_str = re.search(r"SJ_.{8}", jt, re.IGNORECASE).group()
for m in re.findall(
r"(?i)\n\s*\d+\s+[\dABCDEF]{8}\s+.*?\s+%s" %
re.escape(jt_str), asm_lst):
r = re.search(r"\d+\s([\dABCDEF]{8})", m.strip(),
re.IGNORECASE).group(1)
#print "AddCodeXref(0x%08x, 0x%08x, idc.XREF_USER)" % (jt_ea+base_addr, idc.Dword(int(r, 16)+base_addr))
idc.AddCodeXref(jt_ea + base_addr,
idc.Dword(int(r, 16) + base_addr),
idc.XREF_USER)
else:
raise MiscError
for line in asm_lst.split("\n"):
comment = re.search(r"###(.*?)###", line)
if comment != None:
data = re.search(r"\s*\d+\s([\dABCDEF]+)\s([\dABCDEF\(\)]+)",
line)
if data != None:
offset = int(data.group(1), 16)
idc.MakeComm(write_ea + offset, comment.group(1))
fp.close()
return write_ea + data_len + 10
def confirm(msg):
return idaapi.askyn_c(idaapi.ASKBTN_YES, msg) == idaapi.ASKBTN_YES
def main():
# XXX: Change this to the corresponding directory.
input_dir = "/Users/anon/images/ttf/libFontParser"
if os.path.isdir(input_dir):
traces = filter(
lambda x: x.endswith(".trace"),
map(lambda x: os.path.join(input_dir, x), os.listdir(input_dir)))
else:
traces = [input_dir]
# Get loaded binary name
image_name = idc.GetInputFile().lower()
info("IDB binary name '%s'" % image_name)
# Get the image base
image_base = idaapi.get_imagebase()
info("IDB binary base 0x%.16x" % image_base)
# Gather tuples of
coverage_info = []
for filename in traces:
debug("Loading code coverage from '%s'." % filename)
# Get all the hits on this .idb file.
hits = get_image_hits(filename, image_name, image_base)
if not len(hits):
debug("No hits could be loaded from image")
continue
# Save the coverage information.
coverage_info.append(CoverageInformation(filename, hits))
if not len(coverage_info):
info("No coverage information was present for image '%s'" % image_name)
sys.exit()
all_hits = set()
shared_hits = set.intersection(*[x.hits for x in coverage_info])
reached_functions = set()
for element in coverage_info:
all_hits.update(element.hits)
for hit in element.hits:
f = idaapi.get_func(hit)
if f:
reached_functions.add(f.startEA)
info("Covered %d basic blocks in total using %d files" %
(len(all_hits), len(coverage_info)))
info(" Number of shared basic locks %d" % (len(shared_hits)))
info(" Number of reached functions %d" % (len(reached_functions)))
if idaapi.askyn_c(1,
"Do you want to mark all the FUNCTIONS reached?") == 1:
FUNCTION_COLOR = 0xBCF5D1
for function in reached_functions:
info("Reached -> %s" % GetFunctionName(function))
SetFunctionColor(function, FUNCTION_COLOR)
if idaapi.askyn_c(
0, "Do you want to mark all the BASIC BLOCKS reached?") == 1:
BBLOCK_COLOR_1 = 0xA3A9E3
BBLOCK_COLOR_2 = 0xA3D1E3
for hit in all_hits:
SetBasicBlockColor(hit, BBLOCK_COLOR_1)
for hit in shared_hits:
SetBasicBlockColor(hit, BBLOCK_COLOR_2)
return
trace_to_new = {}
for filename in traces:
info("Loading code coverage from '%s'." % filename)
# Get all the hits on this .idb file.
hits = get_image_hits(filename, image_name, image_base)
reached_functions = set()
for e in hits:
f = idaapi.get_func(e)
if not f:
continue
reached_functions.add(f.startEA)
# Get the elements that are introduced by this new trace.
diff_hits = hits - global_hits
diff_functions = reached_functions - global_reached_functions
trace_to_new[filename] = diff_functions
global_hits.update(hits)
global_reached_functions.update(reached_functions)
info("Image '%s' got %d hits (global) and %d function hits (global)." %
(image_name, len(global_hits), len(global_reached_functions)))
for trace_name, introduced_functions in trace_to_new.iteritems():
# Get the original file name.
file_name = trace_name.replace(".trace", "")
# We remove the files that did not introduce any new functions.
if not len(introduced_functions):
assert os.path.exists(file_name)
assert os.path.exists(trace_name)
debug("Removing input file '%s'", os.path.basename(file_name))
debug("Removing trace file '%s'", os.path.basename(trace_name))
os.remove(file_name)
os.remove(trace_name)
continue
fileName, fileExtension = os.path.splitext(file_name)
fileDir = os.path.dirname(file_name)
hash_ = hashlib.sha224(file(file_name).read()).hexdigest()
new_file_name = os.path.join(fileDir, hash_ + fileExtension)
new_trace_name = new_file_name + ".trace"
os.rename(file_name, new_file_name)
os.rename(trace_name, new_trace_name)
info("Trace '%s' introduced functions:" % new_trace_name)
for func in introduced_functions:
info(" %s" % GetFunctionName(func))
if idaapi.askyn_c(1,
"Do you want to mark all the FUNCTIONS reached?") == 1:
FUNCTION_COLOR = 0xBCF5D1
for function in global_reached_functions:
SetFunctionColor(function, FUNCTION_COLOR)
if idaapi.askyn_c(
0, "Do you want to mark all the BASIC BLOCKS reached?") == 1:
BBLOCK_COLOR = 0xf2ddda
for hit in global_hits:
SetBasicBlockColor(hit, BBLOCK_COLOR)
def NasmAssemble(self, function_ea, write_ea):
dir = self.opty_dir
nasm = "C:\\Program Files\\nasm\\nasm.exe"
arg1 = "f_%08x.asm" % function_ea
arg2 = "-o f_%08x.o" % function_ea
arg3 = "-l f_%08x.lst" % function_ea
arg4 = "-Ox"
orig_dir = os.getcwd()
os.chdir(dir)
idc.Batch(0)
while 1:
try:
p = subprocess.Popen([nasm, arg1, arg2, arg3, arg4], stdout=subprocess.PIPE, stderr=subprocess.PIPE, stdin=subprocess.PIPE)
o, e = p.communicate()
if o != "": print o
if e != "": print e
fop = open("f_%08x.o" % function_ea, "rb")
ans = idaapi.askyn_c(0, "HIDECANCEL\nDo you want to manually edit function before writing to IDA?")
if ans == 1:
os.startfile(arg1, "open")
idaapi.warning("Press OK button when you're done editing file.")
fop.close()
continue
else:
idc.Batch(1)
break
except:
error_msg = '\n'.join(e.split("\n")[:15])
os.startfile(arg1, "open")
ans = idaapi.askyn_c(1, """HIDECANCEL\nNASM failed to assemble [f_%08x.o] file.
You can manually edit and NASM this file [f_%08x.asm] and click Yes when you're done.
File is located in directory where your IDB is located.
If you want to skip this function press No.
Nasm output:
%s""" % (function_ea, function_ea, error_msg))
if ans == 1:
continue
else:
os.chdir(orig_dir)
idc.Batch(1)
return None
os.chdir(orig_dir)
print ">>>Writing function [%08x] @ [%08x]" % (function_ea, write_ea)
data = fop.read()
data_len = len(data)
for offset in xrange(0, data_len):
idc.PatchByte(write_ea+offset, ord(data[offset]))
fop.close()
idc.MakeCode(write_ea)
fp = open("%s\\f_%08x.lst" % (dir, function_ea), "r")
asm_lst = fp.read()
base_addr = re.search(r"ORG ([\dABCDEF]+)H", asm_lst).group(1)
base_addr = int(base_addr, 16)
for jt in self.jmp_table_refs:
m = re.search(r"\s*\d+\s+([\dABCDEF]{8}).*?%s" % re.escape(jt), asm_lst, re.IGNORECASE)
if m != None:
jt_ea = int(m.group(1), 16)
jt_str = re.search(r"SJ_.{8}", jt, re.IGNORECASE).group()
for m in re.findall(r"(?i)\n\s*\d+\s+[\dABCDEF]{8}\s+.*?\s+%s" % re.escape(jt_str), asm_lst):
r = re.search(r"\d+\s([\dABCDEF]{8})", m.strip(), re.IGNORECASE).group(1)
#print "AddCodeXref(0x%08x, 0x%08x, idc.XREF_USER)" % (jt_ea+base_addr, idc.Dword(int(r, 16)+base_addr))
idc.AddCodeXref(jt_ea+base_addr, idc.Dword(int(r, 16)+base_addr), idc.XREF_USER)
else:
raise MiscError
for line in asm_lst.split("\n"):
comment = re.search(r"###(.*?)###", line)
if comment != None:
data = re.search(r"\s*\d+\s([\dABCDEF]+)\s([\dABCDEF\(\)]+)", line)
if data != None:
offset = int(data.group(1), 16)
idc.MakeComm(write_ea+offset, comment.group(1))
fp.close()
return write_ea + data_len + 10
def Assemble(self, function, nasm=True):
'''
Main Assemble Function
'''
if self.functions.has_key(function.start_ea):
ans = idaapi.askyn_c(1, "HIDECANCEL\nFunction @ [%08x] was already optimized, are you sure you want to overwrite old one?")
if ans == 0:
return
idc.Batch(1)
self.jmp_table = ""
self.jmp_table_refs = []
self.jmp_deps = {}
self.ref_instr = {}
self.bb_head_ea = {}
function_ea = self.free_ea
function_end = None
if debug:
print ">Assemble(%08x) - Starting" % function.start_ea
if nasm:
if self.opty_dir == None:
self.opty_dir = idc.GetIdbPath()
self.opty_dir = self.opty_dir[:self.opty_dir.rfind(os.sep)] + os.sep + "optimice_%s" % idc.GetInputFile()
if not os.path.isdir(self.opty_dir):
os.mkdir(self.opty_dir)
self.nasm = True
self.nasmfw = open(self.opty_dir + os.sep + 'f_%08x.asm' % function.start_ea, "w+")
self.NasmWriteToFile("[BITS 32]\n")
self.NasmWriteToFile("\torg %09xh\n" % function_ea)
else:
self.nasm = False
for bb_ea in function.DFSFalseTraverseBlocks():
if self.nasm:
self.NasmWriteToFile("\nL%08x:\n" % bb_ea)
self.bb_head_ea[bb_ea] = True
else:
self.bb_head_ea[bb_ea] = self.free_ea
for instr in function.GetBBInstructions(bb_ea):
if instr == None:
continue
mnem = self.BuildAsmString(instr, function)
if debug:
print ">Assemble - Assembling @ %08x [%s]" % (instr.GetOriginEA(), mnem)
if not self.nasm:
self.AsmAndWrite(mnem, instr, function=function)
if not instr.IsCFI():
ref_from = list(function.GetRefsFrom(instr.GetOriginEA()))[0][0]
if self.bb_head_ea.has_key(ref_from):
if self.nasm:
self.NasmWriteToFile("\tjmp L%08x ; \n" % ref_from)
else:
self.AsmAndWrite("jmp %09xh" % self.bb_head_ea[ref_from])
elif instr.GetMnem() == "call":
for ref,path in function.GetRefsFrom(instr.GetOriginEA()):
if path == False:
if self.bb_head_ea.has_key(ref):
if self.nasm:
self.NasmWriteToFile("\tjmp L%08x ; ###FAKE 2 JMP###\n" % ref)
else:
self.AsmAndWrite("jmp %09xh" % self.bb_head_ea[ref])
elif instr.IsJcc():
ref_from = list(function.GetRefsFrom(instr.GetOriginEA()))
for ref,path in ref_from:
if path == False:
break
if path == True:
raise MiscError
elif self.bb_head_ea.has_key(ref):
print "Check this out @ [%08x] ref:[%08x]" % (instr.GetOriginEA(), ref)
if self.nasm:
self.NasmWriteToFile("\tjmp L%08x ; ###FAKE 2 JMP###\n" % ref)
else:
self.AsmAndWrite("jmp %09xh" % self.bb_head_ea[ref])
ret = None
if self.nasm:
self.NasmWriteToFile(self.jmp_table)
self.nasmfw.close()
self.nasmfw = None
ret = self.NasmAssemble(function.start_ea, self.free_ea)
if ret != None:
self.free_ea = ret
else:
ret = self.AsmAndWriteBranches(function)
self.free_ea = ret
self.ref_instr = {}
idc.Batch(0)
if ret != None:
idc.MakeFunction(function_ea)
comment = idc.Comment(function_ea)
if comment != None:
comment = 'Origin@[%08x]; %s' % (function.start_ea, comment)
else:
comment = 'Origin@[%08x];' % function.start_ea
idc.MakeComm(function_ea, comment)
comment = idc.Comment(function.start_ea)
if comment != None:
comment = 'OPTY@[%08x]; %s' % (function_ea, comment)
else:
comment = 'OPTY@[%08x];' % function_ea
idc.MakeComm(function.start_ea, comment)
idaapi.set_name(function_ea, "om_%s" % idc.Name(function.start_ea))
function_end = self.free_ea
self.functions[function.start_ea] = (function_ea, function_end)
idaapi.refresh_idaview_anyway()
def run(self, arg):
if not idaapi.autoIsOk():
if idaapi.askyn_c(
ASKBTN_CANCEL, "HIDECANCEL\n",
"The autoanalysis has not finished yet.\n",
"The result might be incomplete. Do you want to continue?"
) < ASKBTN_NO:
return
form_title = "ETM trace"
form = idaapi.find_tform(form_title)
if form != None:
print "ETM trace window already open. Switching to it."
idaapi.switchto_tform(form, True)
return
trace_file_name = idaapi.askfile_c(0, "",
"Select a trace to display...")
if len(trace_file_name) < 1:
return
image_name = idaapi.get_root_filename()
f = open(trace_file_name, "r")
#trace format: filename[0] id[1] type[2] description[3] src_addr[4] src_func_offset[5] src_image[6] =>[7] dst_addr[8] dst_func_offset[9] dst_image[10]
start_branch = f.readline().split()
if not start_branch:
return
while len(start_branch) != 11:
start_branch = f.readline().split()
if not start_branch:
return
self.c = EtmTraceChoose2(form_title, modal=False)
self.c.callgraph.append("start")
while True:
next_branch = f.readline().split()
if not next_branch:
break
start_branch[10] = start_branch[10].replace("(",
"").replace(")", "")
start_branch[6] = start_branch[6].replace("(", "").replace(")", "")
if start_branch[10].split("/")[-1] != image_name and start_branch[
6].split("/")[-1] != image_name:
start_branch = next_branch
continue
if start_branch[10].split("/")[-1] != image_name:
#to external lib
self.c.add_jump_to_external(start_branch[1], start_branch[8],
start_branch[10])
start_branch = next_branch
continue
if start_branch[6].split("/")[-1] != image_name:
#from external lib
self.c.add_jump_from_external(start_branch[1], start_branch[4],
start_branch[6])
self.c.add_instruction_range(start_branch[1],
[start_branch[8], next_branch[4]])
start_branch = next_branch
self.c.show()
def confirm(msg):
return idaapi.askyn_c(idaapi.ASKBTN_YES, msg) == idaapi.ASKBTN_YES
def confirm(msg):
from idaapi import askyn_c, ASKBTN_CANCEL, ASKBTN_YES
return askyn_c(ASKBTN_CANCEL, msg) == ASKBTN_YES
def repo_init(self, ask_for_remote=True):
# create git
try:
self.repo = ya.GitRepo(".")
self.repo.init()
self.ensure_git_globals()
# add current IDB to repo
self.repo.add_file(self.idb_filename)
# create an initial commit with IDB
self.repo.commit("Initial commit")
if IDA_RUNNING and IDA_IS_INTERACTIVE:
# add a remote to git repo
if ask_for_remote:
url = idaapi.askstr(0, "ssh://gitolite@repo/",
"Specify a remote origin :")
else:
url = None
if url not in [None, ""]:
self.repo.create_remote("origin", url)
if not url.startswith("ssh://"):
if not os.path.exists(url):
if idaapi.askyn_c(
True,
"The target directory doesn't exist, do you want to create it ?"
) == 1:
os.mkdir(url)
temp_repo = ya.GitRepo(url)
temp_repo.init_bare()
# copy idb to local idb
copy_idb_to_local_file()
# push master to remote
self.push_origin_master()
return
except Exception as exc:
logger.error("an error occured during repo_init, error :%s" % exc)
traceback.print_exc()
if idaapi.askyn_c(True,
"could not initialised repo, try again ?") == 1:
while True:
try:
url = idaapi.askstr(0, "ssh://gitolite@repo/",
"Specify a remote origin :")
if url is not None:
url = url.strip()
while url.endswith("/"):
url = url[:-1]
self.repo.remove_remote("origin")
self.repo.create_remote("origin", url)
# push master to remote
self.push_origin_master()
return
except Exception as exc:
logger.error(
"an error occured during repo_init, error :%s" %
exc)
traceback.print_exc()
if idaapi.askyn_c(
True, "could not initialised repo, try again ?"
) != 1:
raise exc
raise exc
def Assemble(self, function, nasm=True):
'''
Main Assemble Function
'''
if self.functions.has_key(function.start_ea):
ans = idaapi.askyn_c(
1,
"HIDECANCEL\nFunction @ [%08x] was already optimized, are you sure you want to overwrite old one?"
)
if ans == 0:
return
idc.Batch(1)
self.jmp_table = ""
self.jmp_table_refs = []
self.jmp_deps = {}
self.ref_instr = {}
self.bb_head_ea = {}
function_ea = self.free_ea
function_end = None
if debug:
print ">Assemble(%08x) - Starting" % function.start_ea
if nasm:
if self.opty_dir == None:
self.opty_dir = idc.GetIdbPath()
self.opty_dir = self.opty_dir[:self.opty_dir.rfind(
os.sep)] + os.sep + "optimice_%s" % idc.GetInputFile()
if not os.path.isdir(self.opty_dir):
os.mkdir(self.opty_dir)
self.nasm = True
self.nasmfw = open(
self.opty_dir + os.sep + 'f_%08x.asm' % function.start_ea,
"w+")
self.NasmWriteToFile("[BITS 32]\n")
self.NasmWriteToFile("\torg %09xh\n" % function_ea)
else:
self.nasm = False
for bb_ea in function.DFSFalseTraverseBlocks():
if self.nasm:
self.NasmWriteToFile("\nL%08x:\n" % bb_ea)
self.bb_head_ea[bb_ea] = True
else:
self.bb_head_ea[bb_ea] = self.free_ea
for instr in function.GetBBInstructions(bb_ea):
if instr == None:
continue
mnem = self.BuildAsmString(instr, function)
if debug:
print ">Assemble - Assembling @ %08x [%s]" % (
instr.GetOriginEA(), mnem)
if not self.nasm:
self.AsmAndWrite(mnem, instr, function=function)
if not instr.IsCFI():
ref_from = list(function.GetRefsFrom(
instr.GetOriginEA()))[0][0]
if self.bb_head_ea.has_key(ref_from):
if self.nasm:
self.NasmWriteToFile("\tjmp L%08x ; \n" % ref_from)
else:
self.AsmAndWrite("jmp %09xh" %
self.bb_head_ea[ref_from])
elif instr.GetMnem() == "call":
for ref, path in function.GetRefsFrom(instr.GetOriginEA()):
if path == False:
if self.bb_head_ea.has_key(ref):
if self.nasm:
self.NasmWriteToFile(
"\tjmp L%08x ; ###FAKE 2 JMP###\n" % ref)
else:
self.AsmAndWrite("jmp %09xh" %
self.bb_head_ea[ref])
elif instr.IsJcc():
ref_from = list(function.GetRefsFrom(instr.GetOriginEA()))
for ref, path in ref_from:
if path == False:
break
if path == True:
raise MiscError
elif self.bb_head_ea.has_key(ref):
print "Check this out @ [%08x] ref:[%08x]" % (
instr.GetOriginEA(), ref)
if self.nasm:
self.NasmWriteToFile(
"\tjmp L%08x ; ###FAKE 2 JMP###\n" % ref)
else:
self.AsmAndWrite("jmp %09xh" % self.bb_head_ea[ref])
ret = None
if self.nasm:
self.NasmWriteToFile(self.jmp_table)
self.nasmfw.close()
self.nasmfw = None
ret = self.NasmAssemble(function.start_ea, self.free_ea)
if ret != None:
self.free_ea = ret
else:
ret = self.AsmAndWriteBranches(function)
self.free_ea = ret
self.ref_instr = {}
idc.Batch(0)
if ret != None:
idc.MakeFunction(function_ea)
comment = idc.Comment(function_ea)
if comment != None:
comment = 'Origin@[%08x]; %s' % (function.start_ea, comment)
else:
comment = 'Origin@[%08x];' % function.start_ea
idc.MakeComm(function_ea, comment)
comment = idc.Comment(function.start_ea)
if comment != None:
comment = 'OPTY@[%08x]; %s' % (function_ea, comment)
else:
comment = 'OPTY@[%08x];' % function_ea
idc.MakeComm(function.start_ea, comment)
idaapi.set_name(function_ea, "om_%s" % idc.Name(function.start_ea))
function_end = self.free_ea
self.functions[function.start_ea] = (function_ea, function_end)
idaapi.refresh_idaview_anyway()
def confirm(msg):
from idaapi import askyn_c, ASKBTN_CANCEL, ASKBTN_YES
return askyn_c(ASKBTN_CANCEL, msg) == ASKBTN_YES