예제 #1
0
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)
예제 #2
0
  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)
예제 #3
0
 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()
예제 #4
0
    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)
예제 #5
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
예제 #6
0
    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
예제 #7
0
    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
예제 #8
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)
예제 #9
0
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)
예제 #10
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 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)
예제 #12
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)))
예제 #13
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)))
예제 #14
0
파일: YaCo.py 프로젝트: silverf0x/YaCo
    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)
예제 #15
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
예제 #16
0
파일: Fentanyl.py 프로젝트: zha0/Fentanyl
    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
예제 #17
0
    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
예제 #18
0
 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
예제 #19
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())
예제 #20
0
 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
예제 #21
0
    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
예제 #22
0
def confirm(msg):
    return idaapi.askyn_c(idaapi.ASKBTN_YES, msg) == idaapi.ASKBTN_YES
예제 #23
0
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)
예제 #24
0
    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
예제 #25
0
    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()
예제 #26
0
    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()
예제 #27
0
def confirm(msg):
    return idaapi.askyn_c(idaapi.ASKBTN_YES, msg) == idaapi.ASKBTN_YES
예제 #28
0
 def confirm(msg):
     from idaapi import askyn_c, ASKBTN_CANCEL, ASKBTN_YES
     return askyn_c(ASKBTN_CANCEL, msg) == ASKBTN_YES
예제 #29
0
    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
예제 #30
0
    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()
예제 #31
0
 def confirm(msg):
     from idaapi import askyn_c, ASKBTN_CANCEL, ASKBTN_YES
     return askyn_c(ASKBTN_CANCEL, msg) == ASKBTN_YES