def main(): button_result = QMessageBox.warning( None, "Warning", "This script will attempt to undo ida_taint2.py changes. It is not perfect and will unpredictably change comments if you've made changes to comments since running ida_taint2.py. Do you want to continue?", QMessageBox.Yes | QMessageBox.No, QMessageBox.No) if button_result == QMessageBox.No: return snapshot = ida_loader.snapshot_t() snapshot.desc = "Before undo_ida_taint2.py @ %s" % ( datetime.datetime.now()) ida_kernwin.take_database_snapshot(snapshot) for segea in Segments(): for funcea in Functions(segea, SegEnd(segea)): function_name = GetFunctionName(funcea) if function_name.startswith("TAINTED_"): MakeName(funcea, function_name.replace("TAINTED_", "")) SetColor(funcea, CIC_FUNC, UNDO_COLOR) for (startea, endea) in Chunks(funcea): for head in Heads(startea, endea): comment = str(Comment(head)) print(comment) if "taint labels" in comment: SetColor(head, CIC_ITEM, UNDO_COLOR) MakeComm(head, label_regex.sub("", comment))
def main(): filename, _ = QFileDialog.getOpenFileName(None, "Open file", ".", "CSV Files(*.csv)") if filename == "": return processes = set() input_file = open(filename, "r") reader = csv.reader(input_file) next(reader, None) for row in reader: processes.add((row[0], int(row[1]))) input_file.close() selected_pid = ProcessSelectDialog.selectProcess(processes) if not selected_pid: return snapshot = ida_loader.snapshot_t() snapshot.desc = "Before ida_taint2.py @ %s" % (datetime.datetime.now()) ida_kernwin.take_database_snapshot(snapshot) input_file = open(filename, "r") reader = csv.reader(input_file) labels_for_pc = {} # skip header next(reader, None) for row in reader: pid = int(row[1]) pc = int(row[2], 16) label = int(row[3]) if pid != selected_pid: continue fn = idaapi.get_func(pc) if not fn: continue fn_start = fn.startEA fn_name = GetFunctionName(fn_start) if "TAINTED" not in fn_name: MakeName(fn_start, "TAINTED_" + fn_name) SetColor(pc, CIC_FUNC, FUNC_COLOR) SetColor(pc, CIC_ITEM, INST_COLOR) if pc not in labels_for_pc: labels_for_pc[pc] = set() labels_for_pc[pc].add(label) input_file.close() for pc, labels in labels_for_pc.iteritems(): comment = Comment(pc) if not comment: comment = "" label_portion = "taint labels = {}".format(list(labels)) if comment == "": comment = label_portion else: comment += ", " + label_portion MakeComm(pc, comment)
def main(): button_result = QMessageBox.warning(None, "Warning", "This script will attempt to undo ida_taint2.py changes. It is not perfect and will unpredictably change comments if you've made changes to comments since running ida_taint2.py. Do you want to continue?", QMessageBox.Yes | QMessageBox.No, QMessageBox.No) if button_result == QMessageBox.No: return snapshot = ida_loader.snapshot_t() snapshot.desc = "Before undo_ida_taint2.py @ %s" % (datetime.datetime.now()) ida_kernwin.take_database_snapshot(snapshot) for segea in Segments(): for funcea in Functions(segea, SegEnd(segea)): function_name = GetFunctionName(funcea) if function_name.startswith("TAINTED_"): MakeName(funcea, function_name.replace("TAINTED_", "")) SetColor(funcea, CIC_FUNC, UNDO_COLOR) for (startea, endea) in Chunks(funcea): for head in Heads(startea, endea): comment = str(Comment(head)) print(comment) if "taint labels" in comment: SetColor(head, CIC_ITEM, UNDO_COLOR) MakeComm(head, label_regex.sub("", comment))
def main(): filename, _ = QFileDialog.getOpenFileName(None, "Open file", ".", "CSV Files(*.csv)") if filename == "": return processes = set() input_file = open(filename, "r") reader = csv.reader(input_file) # what mode was used to produce this output? fmt_row = next(reader, None) mode = fmt_row[0] # where to find the pertinent data depends upon the mode that produced it if ("process" == mode): id_index = 1 # process ID and thread ID are in decimal id_radix = 10 tid_index = 2 pc_index = 4 size_index = 5 prefix = "process_" title = "Select Process" headers = ("Process Name", "PID") has_tid = True else: id_index = 0 # ASID is in hex id_radix = 16 tid_index = -1 pc_index = 2 size_index = 3 prefix = "ASID_" title = "Select Address Space ID" headers = ("Address Space ID", "ASID") has_tid = False binary_name = get_root_filename() match_len = 0 matched_process = None # skip column headers next(reader, None) for row in reader: processes.add((prefix + row[0], int(row[id_index], id_radix))) # the process names from PANDA may be truncated, so the longest one # that is a substring (or equal to) the binary's name is probably the # one we want if (("process" == mode) and (len(row[0]) > match_len) and binary_name.startswith(row[0])): match_len = len(row[0]) matched_process = prefix + row[0] input_file.close() selections = ProcessSelectDialog.selectProcess( processes, matched_process, title, headers, has_tid) if not selections: return if not selections['selected_id']: return snapshot = ida_loader.snapshot_t() snapshot.desc = "Before coverage.py @ %s" % (datetime.datetime.now()) ida_kernwin.take_database_snapshot(snapshot) # the next bit may take a while, if the input file is large ida_kernwin.show_wait_box("HIDECANCEL\nProcessing file " + filename + "...") colored_fn = False info_for_pcs = {} input_file = open(filename, "r") reader = csv.reader(input_file) # skip mode and column headers next(reader, None) next(reader, None) seq_num = 0 for row in reader: cur_id = int(row[id_index], id_radix) pc = int(row[pc_index], 16) size = int(row[size_index]) if (has_tid): cur_tid = int(row[tid_index], id_radix) if cur_id != selections['selected_id']: continue # get the function containing pc fn = idaapi.get_func(pc) if not fn: continue colored_fn = True seq_num = seq_num + 1 color_blocks(fn, pc, size) if (selections['add_tids'] or selections['add_seqs']): if (pc not in info_for_pcs): info_for_pcs[pc] = set() # want to keep the thread IDs (if have them) with the matching # sequence numbers if (selections['add_tids'] and selections['add_seqs']): info_pair = "(" + str(seq_num) + ", " + str(cur_tid) + ")" info_for_pcs[pc].add(info_pair) elif (selections['add_tids']): info_for_pcs[pc].add(cur_tid) else: info_for_pcs[pc].add(seq_num) input_file.close() if (not colored_fn): print("WARNING: did not find any selected functions") else: add_comments(info_for_pcs, selections) ida_kernwin.hide_wait_box()
def _file_downloaded(self, database, progress, reply): """Called when the file has been downloaded.""" progress.close() # Get the absolute path of the file app_path = QCoreApplication.applicationFilePath() app_name = QFileInfo(app_path).fileName() file_ext = "i64" if "64" in app_name else "idb" file_name = "%s_%s.%s" % (database.project, database.name, file_ext) file_path = self._plugin.user_resource("files", file_name) # Write the file to disk with open(file_path, "wb") as output_file: output_file.write(reply.content) self._plugin.logger.info("Saved file %s" % file_name) # Save the old database database = ida_loader.get_path(ida_loader.PATH_TYPE_IDB) if database: ida_loader.save_database(database, ida_loader.DBFL_TEMP) # This is a very ugly hack used to open a database into IDA. We don't # have any function for this in the SDK, so I sorta hijacked the # snapshot functionality in this effect. # Get the library to call functions not present in the bindings idaname = "ida64" if "64" in app_name else "ida" if sys.platform == "win32": dllname, dlltype = idaname + ".dll", ctypes.windll elif sys.platform == "linux2": dllname, dlltype = "lib" + idaname + ".so", ctypes.cdll elif sys.platform == "darwin": dllname, dlltype = "lib" + idaname + ".dylib", ctypes.cdll dllpath = ida_diskio.idadir(None) if not os.path.exists(os.path.join(dllpath, dllname)): dllpath = dllpath.replace("ida64", "ida") dll = dlltype[os.path.join(dllpath, dllname)] # Close the old database using the term_database library function old_path = ida_loader.get_path(ida_loader.PATH_TYPE_IDB) if old_path: dll.term_database() # Open the new database using the init_database library function # This call only won't be enough because the user interface won't # be initialized, this is why the snapshot functionality is used for args = [app_name, file_path] argc = len(args) argv = (ctypes.POINTER(ctypes.c_char) * (argc + 1))() for i, arg in enumerate(args): arg = arg.encode("utf-8") argv[i] = ctypes.create_string_buffer(arg) v = ctypes.c_int(0) av = ctypes.addressof(v) pv = ctypes.cast(av, ctypes.POINTER(ctypes.c_int)) dll.init_database(argc, argv, pv) # Create a temporary copy of the new database because we cannot use # the snapshot functionality to restore the currently opened database file_ext = ".i64" if "64" in app_name else ".idb" tmp_file, tmp_path = tempfile.mkstemp(suffix=file_ext) shutil.copyfile(file_path, tmp_path) # This hook is used to delete the temporary database when all done class UIHooks(ida_kernwin.UI_Hooks): def database_inited(self, is_new_database, idc_script): self.unhook() os.close(tmp_file) if os.path.exists(tmp_path): os.remove(tmp_path) hooks = UIHooks() hooks.hook() # Call the restore_database_snapshot library function # This will initialize the user interface, completing the process s = ida_loader.snapshot_t() s.filename = tmp_path # Use the temporary database ida_kernwin.restore_database_snapshot(s, None, None)
def _database_downloaded(self, branch, progress, reply): """ Called when the file has been downloaded. :param branch: the branch :param progress: the progress dialog :param reply: the reply from the server """ # Close the progress dialog progress.close() # Get the absolute path of the file appPath = QCoreApplication.applicationFilePath() appName = QFileInfo(appPath).fileName() fileExt = 'i64' if '64' in appName else 'idb' fileName = '%s_%s.%s' % (branch.repo, branch.name, fileExt) filePath = local_resource('files', fileName) # Write the packet content to disk with open(filePath, 'wb') as outputFile: outputFile.write(reply.content) logger.info("Saved file %s" % fileName) # Save the old database database = ida_loader.get_path(ida_loader.PATH_TYPE_IDB) if database: ida_loader.save_database(database, ida_loader.DBFL_TEMP) # Get the dynamic library idaname = 'ida64' if '64' in appName else 'ida' if sys.platform == 'win32': dllname, dlltype = idaname + '.dll', ctypes.windll elif sys.platform == 'linux2': dllname, dlltype = 'lib' + idaname + '.so', ctypes.cdll elif sys.platform == 'darwin': dllname, dlltype = 'lib' + idaname + '.dylib', ctypes.cdll dllpath = ida_diskio.idadir(None) if not os.path.exists(os.path.join(dllpath, dllname)): dllpath = dllpath.replace('ida64', 'ida') dll = dlltype[os.path.join(dllpath, dllname)] # Close the old database oldPath = ida_loader.get_path(ida_loader.PATH_TYPE_IDB) if oldPath: dll.term_database() # Open the new database LP_c_char = ctypes.POINTER(ctypes.c_char) args = [appName, filePath] argc = len(args) argv = (LP_c_char * (argc + 1))() for i, arg in enumerate(args): arg = arg.encode('utf-8') argv[i] = ctypes.create_string_buffer(arg) LP_c_int = ctypes.POINTER(ctypes.c_int) v = ctypes.c_int(0) av = ctypes.addressof(v) pv = ctypes.cast(av, LP_c_int) dll.init_database(argc, argv, pv) # Create a copy of the new database fileExt = '.i64' if '64' in appName else '.idb' tmpFile, tmpPath = tempfile.mkstemp(suffix=fileExt) shutil.copyfile(filePath, tmpPath) class UIHooks(ida_kernwin.UI_Hooks): def database_inited(self, is_new_database, idc_script): self.unhook() # Remove the tmp database os.close(tmpFile) if os.path.exists(tmpPath): os.remove(tmpPath) hooks = UIHooks() hooks.hook() # Open the tmp database s = ida_loader.snapshot_t() s.filename = tmpPath ida_kernwin.restore_database_snapshot(s, None, None)
def main(): filename, _ = QFileDialog.getOpenFileName(None, "Open file", ".", "CSV Files(*.csv)") if filename == "": return processes = set() input_file = open(filename, "r") reader = csv.reader(input_file) next(reader, None) for row in reader: processes.add((row[0], int(row[1]))) input_file.close() selected_pid = ProcessSelectDialog.selectProcess(processes) # N.B.: 0 is a valid process ID if (None == selected_pid): return semantic_labels = read_semantic_labels(filename + ".semantic_labels") snapshot = ida_loader.snapshot_t() snapshot.desc = "Before ida_taint2.py @ %s" % (datetime.datetime.now()) ida_kernwin.take_database_snapshot(snapshot) input_file = open(filename, "r") reader = csv.reader(input_file) labels_for_pc = {} # skip header next(reader, None) for row in reader: pid = int(row[1]) pc = int(row[2], 16) label = int(row[3]) try: label = semantic_labels[label] except KeyError: pass if pid != selected_pid: continue fn = ida_funcs.get_func(pc) if not fn: continue fn_start = fn.start_ea fn_name = ida_funcs.get_func_name(fn_start) if "TAINTED" not in fn_name: ida_name.set_name(fn_start, "TAINTED_" + fn_name, ida_name.SN_CHECK) fn.color = FUNC_COLOR ida_nalt.set_item_color(pc, INST_COLOR) if pc not in labels_for_pc: labels_for_pc[pc] = set() labels_for_pc[pc].add(label) input_file.close() for pc, labels in labels_for_pc.items(): comment = ida_bytes.get_cmt(pc, 0) if not comment: comment = "" label_portion = "taint labels = {}".format(list(labels)) if comment == "": comment = label_portion else: comment += ", " + label_portion ida_bytes.set_cmt(pc, comment, 0)