Пример #1
0
class Lighthouse(idaapi.plugin_t):
    """
    The Lighthouse IDA Plugin.
    """

    flags = idaapi.PLUGIN_PROC | idaapi.PLUGIN_MOD | idaapi.PLUGIN_HIDE
    comment = "Code Coverage Explorer"
    help = ""
    wanted_name = "Lighthouse"
    wanted_hotkey = ""

    #--------------------------------------------------------------------------
    # IDA Plugin Overloads
    #--------------------------------------------------------------------------

    def init(self):
        """
        This is called by IDA when it is loading the plugin.
        """

        # attempt plugin initialization
        try:
            self._install_plugin()

        # failed to initialize or integrate the plugin, log and skip loading
        except Exception as e:
            logger.exception("Failed to initialize")
            return idaapi.PLUGIN_SKIP

        # plugin loaded successfully, print the Lighthouse banner
        self.print_banner()
        logger.info("Successfully initialized")

        # tell IDA to keep the plugin loaded (everything is okay)
        return idaapi.PLUGIN_KEEP

    def run(self, arg):
        """
        This is called by IDA when this file is loaded as a script.
        """
        idaapi.warning("The Lighthouse plugin cannot be run as a script.")

    def term(self):
        """
        This is called by IDA when it is unloading the plugin.
        """

        # attempt to cleanup and uninstall our plugin instance
        try:
            self._uninstall_plugin()

        # failed to cleanly remove the plugin, log failure
        except Exception as e:
            logger.exception("Failed to cleanly unload the plugin")

        logger.info("-"*75)
        logger.info("Plugin terminated")

    #--------------------------------------------------------------------------
    # Initialization
    #--------------------------------------------------------------------------

    def _install_plugin(self):
        """
        Initialize & integrate the plugin into IDA.
        """
        self._init()
        self._install_ui()

    def _init(self):
        """
        Initialize plugin members.
        """

        # plugin color palette
        self.palette = LighthousePalette()

        # the coverage engine
        self.director = CoverageDirector(self.palette)

        # the coverage painter
        self.painter = CoveragePainter(self.director, self.palette)

        # the coverage overview widget
        self._ui_coverage_overview = None

        # menu entry icons
        self._icon_id_file = idaapi.BADADDR
        self._icon_id_batch = idaapi.BADADDR
        self._icon_id_overview = idaapi.BADADDR

        # the directory to start the coverage file dialog in
        self._last_directory = idautils.GetIdbDir()

    def _install_ui(self):
        """
        Initialize & integrate all UI elements.
        """
        self._install_load_file()
        self._install_load_batch()
        self._install_open_coverage_overview()

    def print_banner(self):
        """
        Print the Lighthouse plugin banner.
        """

        # build the main banner title
        banner_params = (PLUGIN_VERSION, AUTHORS, DATE)
        banner_title  = "Lighthouse v%s - (c) %s - %s" % banner_params

        # print plugin banner
        lmsg("")
        lmsg("-"*75)
        lmsg("---[ %s" % banner_title)
        lmsg("-"*75)
        lmsg("")

    #--------------------------------------------------------------------------
    # Termination
    #--------------------------------------------------------------------------

    def _uninstall_plugin(self):
        """
        Cleanup & uninstall the plugin from IDA.
        """
        self._uninstall_ui()
        self._cleanup()

    def _uninstall_ui(self):
        """
        Cleanup & uninstall the plugin UI from IDA.
        """
        self._uninstall_open_coverage_overview()
        self._uninstall_load_batch()
        self._uninstall_load_file()

    def _cleanup(self):
        """
        IDB closing event, last chance to spin down threaded workers.
        """
        self.painter.terminate()
        self.director.terminate()

    #--------------------------------------------------------------------------
    # IDA Actions
    #--------------------------------------------------------------------------

    ACTION_LOAD_FILE         = "lighthouse:load_file"
    ACTION_LOAD_BATCH        = "lighthouse:load_batch"
    ACTION_COVERAGE_OVERVIEW = "lighthouse:coverage_overview"

    def _install_load_file(self):
        """
        Install the 'File->Load->Code coverage file...' menu entry.
        """

        # create a custom IDA icon
        icon_path = plugin_resource(os.path.join("icons", "load.png"))
        icon_data = str(open(icon_path, "rb").read())
        self._icon_id_file = idaapi.load_custom_icon(data=icon_data)

        # describe a custom IDA UI action
        action_desc = idaapi.action_desc_t(
            self.ACTION_LOAD_FILE,                     # The action name.
            "~C~ode coverage file...",                 # The action text.
            IDACtxEntry(self.interactive_load_file),   # The action handler.
            None,                                      # Optional: action shortcut
            "Load individual code coverage file(s)",   # Optional: tooltip
            self._icon_id_file                         # Optional: the action icon
        )

        # register the action with IDA
        result = idaapi.register_action(action_desc)
        if not result:
            RuntimeError("Failed to register load_file action with IDA")

        # attach the action to the File-> dropdown menu
        result = idaapi.attach_action_to_menu(
            "File/Load file/",       # Relative path of where to add the action
            self.ACTION_LOAD_FILE,   # The action ID (see above)
            idaapi.SETMENU_APP       # We want to append the action after ^
        )
        if not result:
            RuntimeError("Failed action attach load_file")

        logger.info("Installed the 'Code coverage file' menu entry")

    def _install_load_batch(self):
        """
        Install the 'File->Load->Code coverage batch...' menu entry.
        """

        # create a custom IDA icon
        icon_path = plugin_resource(os.path.join("icons", "batch.png"))
        icon_data = str(open(icon_path, "rb").read())
        self._icon_id_batch = idaapi.load_custom_icon(data=icon_data)

        # describe a custom IDA UI action
        action_desc = idaapi.action_desc_t(
            self.ACTION_LOAD_BATCH,                   # The action name.
            "~C~ode coverage batch...",               # The action text.
            IDACtxEntry(self.interactive_load_batch), # The action handler.
            None,                                     # Optional: action shortcut
            "Load and aggregate code coverage files", # Optional: tooltip
            self._icon_id_batch                       # Optional: the action icon
        )

        # register the action with IDA
        result = idaapi.register_action(action_desc)
        if not result:
            RuntimeError("Failed to register load_batch action with IDA")

        # attach the action to the File-> dropdown menu
        result = idaapi.attach_action_to_menu(
            "File/Load file/",          # Relative path of where to add the action
            self.ACTION_LOAD_BATCH,     # The action ID (see above)
            idaapi.SETMENU_APP          # We want to append the action after ^
        )
        if not result:
            RuntimeError("Failed action attach load_batch")

        logger.info("Installed the 'Code coverage batch' menu entry")

    def _install_open_coverage_overview(self):
        """
        Install the 'View->Open subviews->Coverage Overview' menu entry.
        """

        # create a custom IDA icon
        icon_path = plugin_resource(os.path.join("icons", "overview.png"))
        icon_data = str(open(icon_path, "rb").read())
        self._icon_id_overview = idaapi.load_custom_icon(data=icon_data)

        # describe a custom IDA UI action
        action_desc = idaapi.action_desc_t(
            self.ACTION_COVERAGE_OVERVIEW,            # The action name.
            "~C~overage Overview",                    # The action text.
            IDACtxEntry(self.open_coverage_overview), # The action handler.
            None,                                     # Optional: action shortcut
            "Open database code coverage overview",   # Optional: tooltip
            self._icon_id_overview                    # Optional: the action icon
        )

        # register the action with IDA
        result = idaapi.register_action(action_desc)
        if not result:
            RuntimeError("Failed to register open coverage overview action with IDA")

        # attach the action to the View-> dropdown menu
        result = idaapi.attach_action_to_menu(
            "View/Open subviews/Hex dump", # Relative path of where to add the action
            self.ACTION_COVERAGE_OVERVIEW,    # The action ID (see above)
            idaapi.SETMENU_INS             # We want to insert the action before ^
        )
        if not result:
            RuntimeError("Failed action attach to 'View/Open subviews' dropdown")

        logger.info("Installed the 'Coverage Overview' menu entry")

    def _uninstall_load_file(self):
        """
        Remove the 'File->Load file->Code coverage file...' menu entry.
        """

        # remove the entry from the File-> menu
        result = idaapi.detach_action_from_menu(
            "File/Load file/",
            self.ACTION_LOAD_FILE
        )
        if not result:
            return False

        # unregister the action
        result = idaapi.unregister_action(self.ACTION_LOAD_FILE)
        if not result:
            return False

        # delete the entry's icon
        idaapi.free_custom_icon(self._icon_id_file)
        self._icon_id_file = idaapi.BADADDR

        logger.info("Uninstalled the 'Code coverage file' menu entry")

    def _uninstall_load_batch(self):
        """
        Remove the 'File->Load file->Code coverage batch...' menu entry.
        """

        # remove the entry from the File-> menu
        result = idaapi.detach_action_from_menu(
            "File/Load file/",
            self.ACTION_LOAD_BATCH
        )
        if not result:
            return False

        # unregister the action
        result = idaapi.unregister_action(self.ACTION_LOAD_BATCH)
        if not result:
            return False

        # delete the entry's icon
        idaapi.free_custom_icon(self._icon_id_batch)
        self._icon_id_batch = idaapi.BADADDR

        logger.info("Uninstalled the 'Code coverage batch' menu entry")

    def _uninstall_open_coverage_overview(self):
        """
        Remove the 'View->Open subviews->Coverage Overview' menu entry.
        """

        # remove the entry from the View-> menu
        result = idaapi.detach_action_from_menu(
            "View/Open subviews/Hex dump",
            self.ACTION_COVERAGE_OVERVIEW
        )
        if not result:
            return False

        # unregister the action
        result = idaapi.unregister_action(self.ACTION_COVERAGE_OVERVIEW)
        if not result:
            return False

        # delete the entry's icon
        idaapi.free_custom_icon(self._icon_id_overview)
        self._icon_id_overview = idaapi.BADADDR

        logger.info("Uninstalled the 'Coverage Overview' menu entry")

    #--------------------------------------------------------------------------
    # UI - Actions
    #--------------------------------------------------------------------------

    def open_coverage_overview(self):
        """
        Open the 'Coverage Overview' dialog.
        """
        self.palette.refresh_colors()

        # the coverage overview is already open & visible, simply refresh it
        if self._ui_coverage_overview and self._ui_coverage_overview.isVisible():
            self._ui_coverage_overview.refresh()
            return

        # create a new coverage overview if there is not one visible
        self._ui_coverage_overview = CoverageOverview(self.director)
        self._ui_coverage_overview.show()

    def interactive_load_batch(self):
        """
        Interactive loading & aggregation of coverage files.
        """
        self.palette.refresh_colors()

        #
        # kick off an asynchronous metadata refresh. this collects underlying
        # database metadata while the user will be busy selecting coverage files.
        #

        future = self.director.metadata.refresh(progress_callback=metadata_progress)

        #
        # we will now prompt the user with an interactive file dialog so they
        # can select the coverage files they would like to load from disk.
        #

        loaded_files = self._select_and_load_coverage_files()

        # if no valid coveragee files were selected (and loaded), bail
        if not loaded_files:
            self.director.metadata.abort_refresh()
            return

        # prompt the user to name the new coverage aggregate
        default_name = "BATCH_%s" % self.director.peek_shorthand()
        ok, coverage_name = prompt_string(
            "Batch Name:",
            "Please enter a name for this coverage",
            default_name
        )

        # if user didn't enter a name for the batch, or hit cancel, we abort
        if not (ok and coverage_name):
            lmsg("Aborting batch load...")
            return

        #
        # to continue any further, we need the database metadata. hopefully
        # it has finished with its asynchronous collection, otherwise we will
        # block until it completes. the user will be shown a progress dialog.
        #

        idaapi.show_wait_box("Building database metadata...")
        await_future(future)

        # aggregate all the selected files into one new coverage set
        new_coverage = self._aggregate_batch(loaded_files)

        # inject the the aggregated coverage set
        idaapi.replace_wait_box("Mapping coverage...")
        self.director.create_coverage(coverage_name, new_coverage.data)

        # select the newly created batch coverage
        idaapi.replace_wait_box("Selecting coverage...")
        self.director.select_coverage(coverage_name)

        # all done, hide the IDA wait box
        idaapi.hide_wait_box()
        lmsg("Successfully loaded batch %s..." % coverage_name)

        # show the coverage overview
        self.open_coverage_overview()

    def _aggregate_batch(self, loaded_files):
        """
        Aggregate the given loaded_files data into a single coverage object.
        """
        idaapi.replace_wait_box("Aggregating coverage batch...")

        # create a new coverage set to manually aggregate data into
        coverage = DatabaseCoverage({}, self.palette)

        #
        # loop through the coverage data we have loaded from disk, and begin
        # the normalization process to translate / filter / flatten it for
        # insertion into the director (as a list of instruction addresses)
        #

        for i, data in enumerate(loaded_files, 1):

            # keep the user informed about our progress while loading coverage
            idaapi.replace_wait_box(
                "Aggregating batch data %u/%u" % (i, len(loaded_files))
            )

            # normalize coverage data to the open database
            try:
                addresses = self._normalize_coverage(data, self.director.metadata)

            # normalization failed, print & log it
            except Exception as e:
                lmsg("Failed to map coverage %s" % data.filepath)
                lmsg("- %s" % e)
                logger.exception("Error details:")
                continue

            # aggregate the addresses into the output coverage object
            coverage.add_addresses(addresses, False)

        # return the created coverage name
        return coverage

    def interactive_load_file(self):
        """
        Interactive loading of individual coverage files.
        """
        self.palette.refresh_colors()
        created_coverage = []

        #
        # kick off an asynchronous metadata refresh. this collects underlying
        # database metadata while the user will be busy selecting coverage files.
        #

        future = self.director.metadata.refresh(progress_callback=metadata_progress)

        #
        # we will now prompt the user with an interactive file dialog so they
        # can select the coverage files they would like to load from disk.
        #

        loaded_files = self._select_and_load_coverage_files()

        # if no valid coveragee files were selected (and loaded), bail
        if not loaded_files:
            self.director.metadata.abort_refresh()
            return

        #
        # to continue any further, we need the database metadata. hopefully
        # it has finished with its asynchronous collection, otherwise we will
        # block until it completes. the user will be shown a progress dialog.
        #

        idaapi.show_wait_box("Building database metadata...")
        await_future(future)

        #
        # stop the director's aggregate from updating. this is in the interest
        # of better performance when loading more than one new coverage set
        # into the director.
        #

        self.director.suspend_aggregation()

        #
        # loop through the coverage data we have loaded from disk, and begin
        # the normalization process to translate / filter / flatten its blocks
        # into a generic format the director can understand (a list of addresses)
        #

        for i, data in enumerate(loaded_files, 1):

            # keep the user informed about our progress while loading coverage
            idaapi.replace_wait_box(
                "Normalizing and mapping coverage %u/%u" % (i, len(loaded_files))
            )

            # normalize coverage data to the open database
            try:
                addresses = self._normalize_coverage(data, self.director.metadata)
            except Exception as e:
                lmsg("Failed to map coverage %s" % data.filepath)
                lmsg("- %s" % e)
                logger.exception("Error details:")
                continue

            #
            # ask the director to create and track a new coverage set from
            # the normalized coverage data we provide
            #

            coverage_name = os.path.basename(data.filepath)
            self.director.create_coverage(coverage_name, addresses)

            # save the coverage name to the list of succesful loads
            created_coverage.append(coverage_name)

        #
        # resume the director's aggregation capabilities, triggering an update
        # to recompute the aggregate with the newly loaded coverage
        #

        idaapi.replace_wait_box("Recomputing coverage aggregate...")
        self.director.resume_aggregation()

        # if nothing was mapped, then there's nothing else to do
        if not created_coverage:
            lmsg("No coverage files could be mapped...")
            idaapi.hide_wait_box()
            return

        #
        # select one (the first) of the newly loaded coverage file(s)
        #

        idaapi.replace_wait_box("Selecting coverage...")
        self.director.select_coverage(created_coverage[0])

        # all done, hide the IDA wait box
        idaapi.hide_wait_box()
        lmsg("Successfully loaded %u coverage file(s)..." % len(created_coverage))

        # show the coverage overview
        self.open_coverage_overview()

    def _select_and_load_coverage_files(self):
        """
        Interactive coverage file selection.
        """

        #
        # prompt the user with a QtFileDialog so that they can select any
        # number of coverage files to load at once.
        #
        # if no files are selected, we abort the coverage loading process.
        #

        filenames = self._select_coverage_files()
        if not filenames:
            return None

        # load the selected coverage files from disk and return them
        return self._load_coverage_files(filenames)

    def _select_coverage_files(self):
        """
        Open the 'Load Code Coverage' dialog and capture file selections.
        """

        # create & configure a Qt File Dialog for immediate use
        file_dialog = QtWidgets.QFileDialog(
            None,
            'Open code coverage file',
            self._last_directory,
            'All Files (*.*)'
        )
        file_dialog.setFileMode(QtWidgets.QFileDialog.ExistingFiles)

        # prompt the user with the file dialog, and await filename(s)
        filenames, _ = file_dialog.getOpenFileNames()

        #
        # remember the last directory we were in (parsed from a selected file)
        # for the next time the user comes to load coverage files
        #

        if filenames:
            self._last_directory = os.path.dirname(filenames[0]) + os.sep

        # log the captured (selected) filenames from the dialog
        logger.debug("Captured filenames from file dialog:")
        logger.debug('\n - ' + '\n - '.join(filenames))

        # return the captured filenames
        return filenames

    #--------------------------------------------------------------------------
    # Misc
    #--------------------------------------------------------------------------
    #
    #   NOTE / FUTURE / TODO
    #
    #    In my vision for Lighthouse, I always imagined that it would be
    #    able to dynamically detect and load coverage data from a variety of
    #    different coverage sources and formats (DR, PIN, an inst trace, etc)
    #
    #    The dream was that Lighthouse would have a folder of loaders to parse
    #    and normalize their data to the database / loaded executable so that
    #    they can be injected into the director for exploration.
    #
    #    I would still like to do this, but really haven't heard many people
    #    asking for additional coverage source support yet... so this feature
    #    keeps getting pushed back.
    #
    #    ...
    #
    #    In the mean time, we have a few random functions that are hardcoded
    #    here to load DrCov files and normalize them to the current databasae.
    #

    def _load_coverage_files(self, filenames):
        """
        Load multiple code coverage files from disk.
        """
        loaded_coverage = []

        #
        # loop through each of the given filenames and attempt to load/parse
        # their coverage data from disk
        #

        for filename in filenames:

            # attempt to load/parse a single coverage data file from disk
            try:
                coverage_data = self._load_coverage_file(filename)

            # catch all for parse errors / bad input / malformed files
            except Exception as e:
                lmsg("Failed to load coverage %s" % filename)
                logger.exception("Error details:")
                continue

            # save the loaded coverage data to the output list
            loaded_coverage.append(coverage_data)

        # return all the succesfully loaded coverage files
        return loaded_coverage

    def _load_coverage_file(self, filename):
        """
        Load a single code coverage file from disk.
        """
        return DrcovData(filename)

    def _normalize_coverage(self, coverage_data, metadata):
        """
        Normalize loaded DrCov data to the database metadata.
        """

        # extract the coverage relevant to this IDB (well, the root binary)
        root_filename   = idaapi.get_root_filename()
        coverage_blocks = coverage_data.filter_by_module(root_filename)

        # rebase the basic blocks
        base = idaapi.get_imagebase()
        rebased_blocks = rebase_blocks(base, coverage_blocks)

        # coalesce the blocks into larger contiguous blobs
        condensed_blocks = coalesce_blocks(rebased_blocks)

        # flatten the blobs into individual instructions or addresses
        return metadata.flatten_blocks(condensed_blocks)
Пример #2
0
class Lighthouse(object):
    __metaclass__ = abc.ABCMeta

    #--------------------------------------------------------------------------
    # Initialization
    #--------------------------------------------------------------------------

    def load(self):
        """
        Load the plugin, and integrate its UI into the disassembler.
        """
        self._init()
        self._install_ui()

        # plugin loaded successfully, print the plugin banner
        self.print_banner()
        logger.info("Successfully loaded plugin")

    def _init(self):
        """
        Initialize the core components of the plugin.
        """

        # the plugin's color palette
        self.palette = LighthousePalette()

        # the coverage engine
        self.director = CoverageDirector(self.palette)

        # the coverage painter
        self.painter = CoveragePainter(self.director, self.palette)

        # the coverage overview widget
        self._ui_coverage_overview = None

        # the directory to start the coverage file dialog in
        self._last_directory = None

    def print_banner(self):
        """
        Print the plugin banner.
        """

        # build the main banner title
        banner_params = (PLUGIN_VERSION, AUTHORS, DATE)
        banner_title = "Lighthouse v%s - (c) %s - %s" % banner_params

        # print plugin banner
        lmsg("")
        lmsg("-" * 75)
        lmsg("---[ %s" % banner_title)
        lmsg("-" * 75)
        lmsg("")

    #--------------------------------------------------------------------------
    # Termination
    #--------------------------------------------------------------------------

    def unload(self):
        """
        Unload the plugin, and remove any UI integrations.
        """
        self._uninstall_ui()
        self._cleanup()

        logger.info("-" * 75)
        logger.info("Plugin terminated")

    def _cleanup(self):
        """
        Spin down any lingering core components before plugin unload.
        """
        self.painter.terminate()
        self.director.terminate()

    #--------------------------------------------------------------------------
    # UI Integration (Internal)
    #--------------------------------------------------------------------------

    def _install_ui(self):
        """
        Initialize & integrate all plugin UI elements.
        """
        self._install_load_file()
        self._install_load_batch()
        self._install_open_coverage_overview()

    def _uninstall_ui(self):
        """
        Cleanup & remove all plugin UI integrations.
        """
        self._uninstall_open_coverage_overview()
        self._uninstall_load_batch()
        self._uninstall_load_file()

    @abc.abstractmethod
    def _install_load_file(self):
        """
        Install the 'File->Load->Code coverage file...' menu entry.
        """
        pass

    @abc.abstractmethod
    def _install_load_batch(self):
        """
        Install the 'File->Load->Code coverage batch...' menu entry.
        """
        pass

    @abc.abstractmethod
    def _install_open_coverage_overview(self):
        """
        Install the 'View->Open subviews->Coverage Overview' menu entry.
        """
        pass

    @abc.abstractmethod
    def _uninstall_load_file(self):
        """
        Remove the 'File->Load file->Code coverage file...' menu entry.
        """
        pass

    @abc.abstractmethod
    def _uninstall_load_batch(self):
        """
        Remove the 'File->Load file->Code coverage batch...' menu entry.
        """
        pass

    @abc.abstractmethod
    def _uninstall_open_coverage_overview(self):
        """
        Remove the 'View->Open subviews->Coverage Overview' menu entry.
        """
        pass

    #--------------------------------------------------------------------------
    # UI Actions (Public)
    #--------------------------------------------------------------------------

    def open_coverage_overview(self):
        """
        Open the dockable 'Coverage Overview' dialog.
        """
        self.palette.refresh_colors()

        # the coverage overview is already open & visible, simply refresh it
        if self._ui_coverage_overview and self._ui_coverage_overview.isVisible(
        ):
            self._ui_coverage_overview.refresh()
            return

        # create a new coverage overview if there is not one visible
        self._ui_coverage_overview = CoverageOverview(self)
        self._ui_coverage_overview.show()

    def interactive_load_batch(self):
        """
        Perform the user-interactive loading of a coverage batch.
        """
        self.palette.refresh_colors()

        #
        # kick off an asynchronous metadata refresh. this will run in the
        # background while the user is selecting which coverage files to load
        #

        future = self.director.refresh_metadata(
            progress_callback=metadata_progress)

        #
        # we will now prompt the user with an interactive file dialog so they
        # can select the coverage files they would like to load from disk
        #

        filenames = self._select_coverage_files()

        #
        # load the selected coverage files from disk (if any), returning a list
        # of loaded DrcovData objects (which contain coverage data)
        #

        drcov_list = load_coverage_files(filenames)
        if not drcov_list:
            self.director.metadata.abort_refresh()
            return

        # prompt the user to name the new coverage aggregate
        default_name = "BATCH_%s" % self.director.peek_shorthand()
        ok, coverage_name = prompt_string(
            "Batch Name:", "Please enter a name for this coverage",
            default_name)

        #
        # if user didn't enter a name for the batch (or hit cancel) we should
        # abort the loading process...
        #

        if not (ok and coverage_name):
            lmsg("User failed to enter a name for the loaded batch...")
            self.director.metadata.abort_refresh()
            return

        #
        # to begin mapping the loaded coverage data, we require that the
        # asynchronous database metadata refresh has completed. if it is
        # not done yet, we will block here until it completes.
        #
        # a progress dialog depicts the work remaining in the refresh
        #

        disassembler.show_wait_box("Building database metadata...")
        await_future(future)

        #
        # now that the database metadata is available, we can use the director
        # to normalize and condense (aggregate) all the coverage data
        #

        new_coverage, errors = self.director.aggregate_drcov_batch(drcov_list)

        #
        # finally, we can inject the aggregated coverage data into the
        # director under the user specified batch name
        #

        disassembler.replace_wait_box("Mapping coverage...")
        self.director.create_coverage(coverage_name, new_coverage.data)

        # select the newly created batch coverage
        disassembler.replace_wait_box("Selecting coverage...")
        self.director.select_coverage(coverage_name)

        # all done! pop the coverage overview to show the user their results
        disassembler.hide_wait_box()
        lmsg("Successfully loaded batch %s..." % coverage_name)
        self.open_coverage_overview()

        # finally, emit any notable issues that occurred during load
        warn_errors(errors)

    def interactive_load_file(self):
        """
        Perform the user-interactive loading of individual coverage files.
        """
        self.palette.refresh_colors()

        #
        # kick off an asynchronous metadata refresh. this will run in the
        # background while the user is selecting which coverage files to load
        #

        future = self.director.refresh_metadata(
            progress_callback=metadata_progress)

        #
        # we will now prompt the user with an interactive file dialog so they
        # can select the coverage files they would like to load from disk
        #

        filenames = self._select_coverage_files()

        #
        # load the selected coverage files from disk (if any), returning a list
        # of loaded DrcovData objects (which contain coverage data)
        #

        disassembler.show_wait_box("Loading coverage from disk...")
        drcov_list = load_coverage_files(filenames)
        if not drcov_list:
            disassembler.hide_wait_box()
            self.director.metadata.abort_refresh()
            return

        #
        # to begin mapping the loaded coverage data, we require that the
        # asynchronous database metadata refresh has completed. if it is
        # not done yet, we will block here until it completes.
        #
        # a progress dialog depicts the work remaining in the refresh
        #

        disassembler.replace_wait_box("Building database metadata...")
        await_future(future)

        # insert the loaded drcov data objects into the director
        created_coverage, errors = self.director.create_coverage_from_drcov_list(
            drcov_list)

        #
        # if the director failed to map any coverage, the user probably
        # provided bad files. emit any warnings and bail...
        #

        if not created_coverage:
            lmsg("No coverage files could be loaded...")
            disassembler.hide_wait_box()
            warn_errors(errors)
            return

        #
        # activate the first of the newly loaded coverage file(s). this is the
        # one that will be visible in the coverage overview once opened
        #

        disassembler.replace_wait_box("Selecting coverage...")
        self.director.select_coverage(created_coverage[0])

        # all done! pop the coverage overview to show the user their results
        disassembler.hide_wait_box()
        lmsg("Successfully loaded %u coverage file(s)..." %
             len(created_coverage))
        self.open_coverage_overview()

        # finally, emit any notable issues that occurred during load
        warn_errors(errors)

    #--------------------------------------------------------------------------
    # Internal
    #--------------------------------------------------------------------------

    def _select_coverage_files(self):
        """
        Prompt a file selection dialog, returning file selections.

        NOTE: This saves & reuses the last known directory for subsequent uses.
        """
        if not self._last_directory:
            self._last_directory = disassembler.get_database_directory()

        # create & configure a Qt File Dialog for immediate use
        file_dialog = QtWidgets.QFileDialog(None, 'Open code coverage file',
                                            self._last_directory,
                                            'All Files (*.*)')
        file_dialog.setFileMode(QtWidgets.QFileDialog.ExistingFiles)

        # prompt the user with the file dialog, and await filename(s)
        filenames, _ = file_dialog.getOpenFileNames()

        #
        # remember the last directory we were in (parsed from a selected file)
        # for the next time the user comes to load coverage files
        #

        if filenames:
            self._last_directory = os.path.dirname(filenames[0]) + os.sep

        # log the captured (selected) filenames from the dialog
        logger.debug("Captured filenames from file dialog:")
        for name in filenames:
            logger.debug(" - %s" % name)

        # return the captured filenames
        return filenames
Пример #3
0
class Lighthouse(plugin_t):
    """
    The Lighthouse IDA Plugin.
    """

    flags = idaapi.PLUGIN_PROC | idaapi.PLUGIN_MOD | idaapi.PLUGIN_HIDE
    comment = "Code Coverage Explorer"
    help = ""
    wanted_name = "Lighthouse"
    wanted_hotkey = ""

    def __init__(self):

        # plugin color palette
        self.palette = LighthousePalette()

        # the coverage engine
        self.director = CoverageDirector(self.palette)

        # the coverage painter
        self.painter = CoveragePainter(self.director, self.palette)

        # the coverage overview widget
        self._ui_coverage_overview = None

        # members for the 'Load Code Coverage' menu entry
        self._icon_id_load = idaapi.BADADDR

        # members for the 'Coverage Overview' menu entry
        self._icon_id_overview = idaapi.BADADDR

    #--------------------------------------------------------------------------
    # IDA Plugin Overloads
    #--------------------------------------------------------------------------

    def init(self):
        """
        This is called by IDA when it is loading the plugin.
        """

        # attempt plugin initialization
        try:
            self._install_plugin()

        # failed to initialize or integrate the plugin, log and skip loading
        except Exception as e:
            logger.exception("Failed to initialize")
            return idaapi.PLUGIN_SKIP

        # plugin loaded successfully, print the Lighthouse banner
        self.print_banner()
        logger.info("Successfully initialized")

        # tell IDA to keep the plugin loaded (everything is okay)
        return idaapi.PLUGIN_KEEP

    def run(self, arg):
        """
        This is called by IDA when this file is loaded as a script.
        """
        idaapi.warning("The Lighthouse plugin cannot be run as a script.")

    def term(self):
        """
        This is called by IDA when it is unloading the plugin.
        """

        # attempt to cleanup and uninstall our plugin instance
        try:
            self._uninstall_plugin()

        # failed to cleanly remove the plugin, log failure
        except Exception as e:
            logger.exception("Failed to cleanly unload the plugin")

        logger.info("-" * 75)
        logger.info("Plugin terminated")

    #--------------------------------------------------------------------------
    # Initialization
    #--------------------------------------------------------------------------

    def _install_plugin(self):
        """
        Initialize & integrate the plugin into IDA.
        """
        self._install_ui()

    def print_banner(self):
        """
        Print the Lighthouse plugin banner.
        """

        # build the main banner title
        banner_params = (PLUGIN_VERSION, AUTHORS, DATE)
        banner_title = "Lighthouse v%s - (c) %s - %s" % banner_params

        # print plugin banner
        lmsg("")
        lmsg("-" * 75)
        lmsg("---[ %s" % banner_title)
        lmsg("-" * 75)
        lmsg("")

    #--------------------------------------------------------------------------
    # Initialization - UI
    #--------------------------------------------------------------------------

    def _install_ui(self):
        """
        Initialize & integrate all UI elements.
        """

        # install the 'Load Coverage' file dialog
        self._install_load_file_dialog()
        self._install_open_coverage_overview()

    #--------------------------------------------------------------------------
    # Termination
    #--------------------------------------------------------------------------

    def _uninstall_plugin(self):
        """
        Cleanup & uninstall the plugin from IDA.
        """
        self._uninstall_ui()

    #--------------------------------------------------------------------------
    # Termination - UI
    #--------------------------------------------------------------------------

    def _uninstall_ui(self):
        """
        Cleanup & uninstall the plugin UI from IDA.
        """
        self._uninstall_open_coverage_overview()
        self._uninstall_load_file_dialog()

    #--------------------------------------------------------------------------
    # IDA Actions
    #--------------------------------------------------------------------------

    ACTION_LOAD_COVERAGE = "lighthouse:load_coverage"
    ACTION_COVERAGE_OVERVIEW = "lighthouse:coverage_overview"

    def _install_load_file_dialog(self):
        """
        Install the 'File->Load->Code Coverage File(s)...' menu entry.
        """

        # create a custom IDA icon
        icon_path = plugin_resource(os.path.join("icons", "load.png"))
        icon_data = str(open(icon_path, "rb").read())
        self._icon_id_load = idaapi.load_custom_icon(data=icon_data)

        # describe a custom IDA UI action
        action_desc = idaapi.action_desc_t(
            self.ACTION_LOAD_COVERAGE,  # The action name.
            "~C~ode Coverage File(s)...",  # The action text.
            IDACtxEntry(self.load_coverage),  # The action handler.
            None,  # Optional: action shortcut
            "Load a code coverage file for this IDB",  # Optional: tooltip
            self._icon_id_load  # Optional: the action icon
        )

        # register the action with IDA
        result = idaapi.register_action(action_desc)
        if not result:
            RuntimeError("Failed to register load coverage action with IDA")

        # attach the action to the File-> dropdown menu
        result = idaapi.attach_action_to_menu(
            "File/Load file/",  # Relative path of where to add the action
            self.ACTION_LOAD_COVERAGE,  # The action ID (see above)
            idaapi.SETMENU_APP  # We want to append the action after ^
        )
        if not result:
            RuntimeError("Failed action attach to 'File/Load file/' dropdown")

        logger.info("Installed the 'Load Code Coverage' menu entry")

    def _install_open_coverage_overview(self):
        """
        Install the 'View->Open subviews->Coverage Overview' menu entry.
        """

        # create a custom IDA icon
        icon_path = plugin_resource(os.path.join("icons", "overview.png"))
        icon_data = str(open(icon_path, "rb").read())
        self._icon_id_overview = idaapi.load_custom_icon(data=icon_data)

        # describe a custom IDA UI action
        action_desc = idaapi.action_desc_t(
            self.ACTION_COVERAGE_OVERVIEW,  # The action name.
            "~C~overage Overview",  # The action text.
            IDACtxEntry(self.open_coverage_overview),  # The action handler.
            None,  # Optional: action shortcut
            "Open database code coverage overview",  # Optional: tooltip
            self._icon_id_overview  # Optional: the action icon
        )

        # register the action with IDA
        result = idaapi.register_action(action_desc)
        if not result:
            RuntimeError(
                "Failed to register open coverage overview action with IDA")

        # attach the action to the View-> dropdown menu
        result = idaapi.attach_action_to_menu(
            "View/Open subviews/Hex dump",  # Relative path of where to add the action
            self.ACTION_COVERAGE_OVERVIEW,  # The action ID (see above)
            idaapi.SETMENU_INS  # We want to insert the action before ^
        )
        if not result:
            RuntimeError(
                "Failed action attach to 'View/Open subviews' dropdown")

        logger.info("Installed the 'Coverage Overview' menu entry")

    def _uninstall_load_file_dialog(self):
        """
        Remove the 'File->Load file->Code Coverage File(s)...' menu entry.
        """

        # remove the entry from the File-> menu
        result = idaapi.detach_action_from_menu("File/Load file/",
                                                self.ACTION_LOAD_COVERAGE)
        if not result:
            return False

        # unregister the action
        result = idaapi.unregister_action(self.ACTION_LOAD_COVERAGE)
        if not result:
            return False

        # delete the entry's icon
        idaapi.free_custom_icon(self._icon_id_load)
        self._icon_id_load = idaapi.BADADDR

        logger.info("Uninstalled the 'Load Code Coverage' menu entry")

    def _uninstall_open_coverage_overview(self):
        """
        Remove the 'View->Open subviews->Coverage Overview' menu entry.
        """

        # remove the entry from the View-> menu
        result = idaapi.detach_action_from_menu("View/Open subviews/Hex dump",
                                                self.ACTION_COVERAGE_OVERVIEW)
        if not result:
            return False

        # unregister the action
        result = idaapi.unregister_action(self.ACTION_COVERAGE_OVERVIEW)
        if not result:
            return False

        # delete the entry's icon
        idaapi.free_custom_icon(self._icon_id_overview)
        self._icon_id_overview = idaapi.BADADDR

        logger.info("Uninstalled the 'Coverage Overview' menu entry")

    #--------------------------------------------------------------------------
    # UI - Actions
    #--------------------------------------------------------------------------

    def open_coverage_overview(self):
        """
        Open the 'Coverage Overview' dialog.
        """
        self._ui_coverage_overview = CoverageOverview(self.director)
        self._ui_coverage_overview.show()

    def load_coverage(self):
        """
        An interactive file dialog flow for loading code coverage files.
        """

        #
        # kick off an asynchronous metadata refresh. this collects underlying
        # database metadata while the user will be busy selecting coverage files.
        #
        # the collected metadata enables the director to process, map, and
        # manipulate loaded coverage data in a performant, asynchronous manner.
        #

        future = self.director.metadata.refresh(
            progress_callback=metadata_progress)

        #
        # prompt the user with a QtFileDialog so that they can select any
        # number of coverage files to load at once.
        #
        # if no files are selected, we abort the coverage loading process.
        #

        filenames = self._select_coverage_files()
        if not filenames:
            return

        #
        # load the selected coverage files from disk
        #

        coverage_data = self._load_coverage_files(filenames)

        #
        # refresh the theme aware color palette for lighthouse
        #

        self.palette.refresh_colors()

        #
        # to continue any further, we need the database metadata. hopefully
        # it has finished with its asynchronous collection, otherwise we will
        # block until it completes. the user will be shown a progress dialog.
        #

        idaapi.show_wait_box("Building database metadata...")
        await_future(future)

        #
        # at this point the metadata caching is guaranteed to be complete.
        # the coverage data has been loaded and is ready for mapping and
        # management by the director.
        #

        idaapi.replace_wait_box("Normalizing and mapping coverage data...")

        #
        # start a batch coverage data load for better performance incase we
        # are loading more than one new coverage file / data to the director.
        #

        self.director.start_batch()

        # a list to output the names of successfully loaded coverage files
        loaded = []

        #
        # loop through the coverage data we have loaded from disk, and begin
        # the normalization process to translate / filter / flatten it for
        # insertion into the director (as a list of instruction addresses)
        #

        for i, data in enumerate(coverage_data, 1):

            # keep the user informed about our progress while loading coverage
            idaapi.replace_wait_box("Normalizing and mapping coverage %u/%u" %
                                    (i, len(coverage_data)))

            # TODO: it would be nice to get rid of this try/catch in the long run
            try:

                # normalize coverage data to the database
                addresses = self._normalize_coverage(data,
                                                     self.director.metadata)

                # enlighten the coverage director to this new runtime data
                coverage_name = os.path.basename(data.filepath)
                self.director.add_coverage(coverage_name, addresses)

                # if we made it this far, the coverage must have loaded okay...
                loaded.append(coverage_name)

            except Exception as e:
                lmsg("Failed to load coverage:")
                lmsg("- %s" % e)
                logger.error("Error details:")
                continue

        # collapse the batch job to recompute the director's aggregate coverage set
        self.director.end_batch()

        # select the 'first' coverage file loaded from this round
        if loaded:
            self.director.select_coverage(loaded[0])

        # all done, hide the IDA wait box
        idaapi.hide_wait_box()

        # print a success message to the output window
        lmsg("loaded %u coverage file(s)..." % len(loaded))

        # show the coverage overview
        self.open_coverage_overview()

    def _select_coverage_files(self):
        """
        Open the 'Load Code Coverage' dialog and capture file selections.
        """

        # create & configure a Qt File Dialog for immediate use
        file_dialog = QtWidgets.QFileDialog(None, 'Open Code Coverage File(s)')
        file_dialog.setFileMode(QtWidgets.QFileDialog.ExistingFiles)

        # prompt the user with the file dialog, and await filename(s)
        filenames, _ = file_dialog.getOpenFileNames()

        # log the captured (selected) filenames from the dialog
        logger.debug("Captured filenames from file dialog:")
        logger.debug('\n - ' + '\n - '.join(filenames))

        # return the captured filenames
        return filenames

    #--------------------------------------------------------------------------
    # Misc
    #--------------------------------------------------------------------------
    #
    #   NOTE / FUTURE / TODO
    #
    #    In my vision for Lighthouse, I always imagined that it would be
    #    able to dynamically detect and load coverage data from a variety of
    #    different coverage sources and formats (DR, PIN, an inst trace, etc)
    #
    #    The dream was that Lighthouse would have a folder of loaders to parse
    #    and normalize their data to the database / loaded executable so that
    #    they can be injected into the director for exploration.
    #
    #    I would still like to do this, but really haven't heard many people
    #    asking for additional coverage source support yet... so this feature
    #    keeps getting pushed back.
    #
    #    ...
    #
    #    In the mean time, we have a few random functions that are hardcoded
    #    here to load DrCov files and normalize them to the current databasae.
    #

    def _load_coverage_files(self, filenames):
        """
        Load multiple code coverage files from disk.
        """
        return [self._load_coverage_file(filename) for filename in filenames]

    def _load_coverage_file(self, filename):
        """
        Load a single code coverage file from disk.
        """
        return DrcovData(filename)

    def _normalize_coverage(self, coverage_data, metadata):
        """
        Normalize loaded DrCov data to the database metadata.
        """

        # extract the coverage relevant to this IDB (well, the root binary)
        root_filename = idaapi.get_root_filename()
        coverage_blocks = coverage_data.filter_by_module(root_filename)

        # rebase the basic blocks
        base = idaapi.get_imagebase()
        rebased_blocks = rebase_blocks(base, coverage_blocks)

        # coalesce the blocks into larger contiguous blobs
        condensed_blocks = coalesce_blocks(rebased_blocks)

        # flatten the blobs into individual instructions or addresses
        return metadata.flatten_blocks(condensed_blocks)
Пример #4
0
class Lighthouse(plugin_t):
    """
    The Lighthouse IDA Plugin.
    """

    flags = idaapi.PLUGIN_PROC | idaapi.PLUGIN_MOD
    comment = "Code Coverage Explorer"
    help = ""
    wanted_name = "Lighthouse"
    wanted_hotkey = ""

    def __init__(self):

        # plugin color palette
        self.palette = LighthousePalette()

        # the database coverage data conglomerate
        self.director = CoverageDirector(self.palette)

        # the coverage painter
        self.painter = CoveragePainter(self.director, self.palette)

        # plugin qt elements
        self._ui_coverage_overview = CoverageOverview(self.director)

        # members for the 'Load Code Coverage' menu entry
        self._icon_id_load = idaapi.BADADDR
        self._action_name_load = "lighthouse:load_coverage"

        # members for the 'Coverage Overview' menu entry
        self._icon_id_overview = idaapi.BADADDR
        self._action_name_overview = "lighthouse:coverage_overview"

    #--------------------------------------------------------------------------
    # IDA Plugin Overloads
    #--------------------------------------------------------------------------

    def init(self):
        """
        This is called by IDA when it is loading the plugin.
        """

        # attempt plugin initialization
        try:
            self._install_plugin()

        # failed to initialize or integrate the plugin, log and skip loading
        except Exception as e:
            logger.exception("Failed to initialize")
            return idaapi.PLUGIN_SKIP

        # plugin loaded successfully, print the Lighthouse banner
        self.print_banner()
        logger.info("Successfully initialized")

        # tell IDA to keep the plugin loaded (everything is okay)
        return idaapi.PLUGIN_KEEP

    def run(self, arg):
        """
        This is called by IDA when this file is loaded as a script.
        """
        idaapi.warning("The Lighthouse plugin cannot be run as a script.")

    def term(self):
        """
        This is called by IDA when it is unloading the plugin.
        """

        # attempt to cleanup and uninstall our plugin instance
        try:
            self._uninstall_plugin()

        # failed to cleanly remove the plugin, log failure
        except Exception as e:
            logger.exception("Failed to cleanly unload the plugin")

        logger.info("-" * 75)
        logger.info("Plugin terminated")

    #--------------------------------------------------------------------------
    # Initialization
    #--------------------------------------------------------------------------

    def _install_plugin(self):
        """
        Initialize & integrate the plugin into IDA.
        """
        self._install_ui()

    def print_banner(self):
        """
        Print the Lighthouse plugin banner.
        """

        # build the main banner title
        banner_params = (PLUGIN_VERSION, AUTHORS, DATE)
        banner_title = "Lighthouse v%s - (c) %s - %s" % banner_params

        # print plugin banner
        lmsg("")
        lmsg("-" * 75)
        lmsg("---[ %s" % banner_title)
        lmsg("-" * 75)
        lmsg("")

    #--------------------------------------------------------------------------
    # Initialization - UI
    #--------------------------------------------------------------------------

    def _install_ui(self):
        """
        Initialize & integrate all UI elements.
        """

        # install the 'Load Coverage' file dialog
        self._install_load_file_dialog()
        self._install_open_coverage_overview()

    def _install_load_file_dialog(self):
        """
        Install the 'File->Load->Code Coverage File(s)...' menu entry.
        """

        # create a custom IDA icon
        self._icon_id_load = idaapi.load_custom_icon(
            data=str(open(plugin_resource("icons/load.png"), "rb").read()))

        # describe a custom IDA UI action
        action_desc = idaapi.action_desc_t(
            self._action_name_load,  # The action name.
            "~C~ode Coverage File(s)...",  # The action text.
            IDACtxEntry(self.load_coverage),  # The action handler.
            None,  # Optional: action shortcut
            "Load a code coverage file for this IDB",  # Optional: tooltip
            self._icon_id_load  # Optional: the action icon
        )

        # register the action with IDA
        result = idaapi.register_action(action_desc)
        if not result:
            RuntimeError("Failed to register load coverage action with IDA")

        # attach the action to the File-> dropdown menu
        result = idaapi.attach_action_to_menu(
            "File/Load file/",  # Relative path of where to add the action
            self._action_name_load,  # The action ID (see above)
            idaapi.SETMENU_APP  # We want to append the action after ^
        )
        if not result:
            RuntimeError("Failed action attach to 'File/Load file/' dropdown")

        logger.info("Installed the 'Load Code Coverage' menu entry")

    def _install_open_coverage_overview(self):
        """
        Install the 'View->Open subviews->Coverage Overview' menu entry.
        """

        # create a custom IDA icon
        self._icon_id_overview = idaapi.load_custom_icon(
            data=str(open(plugin_resource("icons/overview.png"), "rb").read()))

        # describe a custom IDA UI action
        action_desc = idaapi.action_desc_t(
            self._action_name_overview,  # The action name.
            "~C~overage Overview",  # The action text.
            IDACtxEntry(self.open_coverage_overview),  # The action handler.
            None,  # Optional: action shortcut
            "Open database code coverage overview",  # Optional: tooltip
            self._icon_id_overview  # Optional: the action icon
        )

        # register the action with IDA
        result = idaapi.register_action(action_desc)
        if not result:
            RuntimeError(
                "Failed to register open coverage overview action with IDA")

        # attach the action to the View-> dropdown menu
        result = idaapi.attach_action_to_menu(
            "View/Open subviews/Hex dump",  # Relative path of where to add the action
            self._action_name_overview,  # The action ID (see above)
            idaapi.SETMENU_INS  # We want to insert the action before ^
        )
        if not result:
            RuntimeError(
                "Failed action attach to 'View/Open subviews' dropdown")

        logger.info("Installed the 'Coverage Overview' menu entry")

    #--------------------------------------------------------------------------
    # Termination
    #--------------------------------------------------------------------------

    def _uninstall_plugin(self):
        """
        Cleanup & uninstall the plugin from IDA.
        """
        self._uninstall_ui()

    #--------------------------------------------------------------------------
    # Termination - UI
    #--------------------------------------------------------------------------

    def _uninstall_ui(self):
        """
        Cleanup & uninstall the plugin UI from IDA.
        """
        self._uninstall_open_coverage_overview()
        self._uninstall_load_file_dialog()

    def _uninstall_load_file_dialog(self):
        """
        Remove the 'File->Load file->Code Coverage File(s)...' menu entry.
        """

        # remove the entry from the File-> menu
        result = idaapi.detach_action_from_menu("File/Load file/",
                                                self._action_name_load)
        if not result:
            return False

        # unregister the action
        result = idaapi.unregister_action(self._action_name_load)
        if not result:
            return False

        # delete the entry's icon
        idaapi.free_custom_icon(self._icon_id_load)
        self._icon_id_load = idaapi.BADADDR

        logger.info("Uninstalled the 'Load Code Coverage' menu entry")

    def _uninstall_open_coverage_overview(self):
        """
        Remove the 'View->Open subviews->Coverage Overview' menu entry.
        """

        # remove the entry from the View-> menu
        result = idaapi.detach_action_from_menu("View/Open subviews/Hex dump",
                                                self._action_name_overview)
        if not result:
            return False

        # unregister the action
        result = idaapi.unregister_action(self._action_name_overview)
        if not result:
            return False

        # delete the entry's icon
        idaapi.free_custom_icon(self._icon_id_overview)
        self._icon_id_overview = idaapi.BADADDR

        logger.info("Uninstalled the 'Coverage Overview' menu entry")

    #--------------------------------------------------------------------------
    # UI - Actions
    #--------------------------------------------------------------------------

    def load_coverage(self):
        """
        An interactive file dialog flow for loading code coverage files.
        """

        #
        # kick off an asynchronous metadata refresh. this collects underlying
        # database metadata while the user will be busy selecting coverage files.
        #
        # the collected metadata enables the director to process, map, and
        # manipulate loaded coverage data in a performant, asynchronous manner.
        #

        future = self.director.metadata.refresh(
            progress_callback=metadata_progress)

        #
        # prompt the user with a QtFileDialog so that they can select any
        # number of coverage files to load at once.
        #
        # if no files are selected, we abort the coverage loading process.
        #

        filenames = self._select_coverage_files()
        if not filenames:
            return

        #
        # load the selected coverage files from disk
        #

        coverage_data = self._load_coverage_files(filenames)

        #
        # refresh the theme aware color palette for lighthouse
        #

        self.palette.refresh_colors()

        #
        # to continue any further, we need the database metadata. hopefully
        # it has finished with its asynchronous collection, otherwise we will
        # block until it completes. the user will be shown a progress dialog.
        #

        idaapi.show_wait_box("Building database metadata...")
        await_future(future)

        #
        # at this point the metadata caching is guaranteed to be complete.
        # the coverage data has been loaded and is ready for mapping.
        #

        idaapi.replace_wait_box("Normalizing and mapping coverage data...")

        #
        # TODO:
        #
        #   I do not hold great confidence in this code yet, so let's wrap
        #   this in a try/catch so the user doesn't get stuck with a wait
        #   box they can't close should things go poorly ;P
        #

        try:

            for data in coverage_data:

                # normalize coverage data to the database
                name = os.path.basename(data.filepath)
                addresses = self._normalize_coverage(data,
                                                     self.director.metadata)

                # enlighten the coverage director to this new runtime data
                self.director.add_coverage(name, addresses)

            # select the 'first' coverage file loaded
            self.director.select_coverage(self.director.coverage_names[0])

            idaapi.hide_wait_box()

        # 'something happened :('
        except Exception as e:
            idaapi.hide_wait_box()
            lmsg("Failed to load coverage:")
            lmsg("- %s" % e)
            logger.exception(e)
            return

        # print a success message to the output window
        lmsg("loaded %u coverage file(s)..." % len(coverage_data))

        # show the coverage overview
        self.open_coverage_overview()

    def open_coverage_overview(self):
        """
        Open the 'Coverage Overview' dialog.
        """
        self._ui_coverage_overview.Show()

    def _select_coverage_files(self):
        """
        Open the 'Load Code Coverage' dialog and capture file selections.
        """

        # create & configure a Qt File Dialog for immediate use
        file_dialog = QtWidgets.QFileDialog(None, 'Open Code Coverage File(s)')
        file_dialog.setFileMode(QtWidgets.QFileDialog.ExistingFiles)

        # prompt the user with the file dialog, and await filename(s)
        filenames, _ = file_dialog.getOpenFileNames()

        # log the captured (selected) filenames from the dialog
        logger.debug("Captured filenames from file dialog:")
        logger.debug(filenames)

        # return the captured filenames
        return filenames

    #--------------------------------------------------------------------------
    # Misc
    #--------------------------------------------------------------------------

    def _load_coverage_files(self, filenames):
        """
        Load multiple code coverage files from disk.
        """
        return [self._load_coverage_file(filename) for filename in filenames]

    def _load_coverage_file(self, filename):
        """
        Load a single code coverage file from disk.

        TODO: Add other formats. Only drcov logs supported for now.
        """
        return DrcovData(filename)

    def _normalize_coverage(self, coverage_data, metadata):
        """
        Normalize loaded coverage data to the database metadata.

        TODO:

          This will probably be moved out and turn into a layer for each unique
          lighthouse coverage parser/loader to implement.

          for example, this effectively translate the DrcovData log to a more
          general / universal format for the director.

        """

        # extract the coverage relevant to this IDB (well, the root binary)
        root_filename = idaapi.get_root_filename()
        coverage_blocks = coverage_data.filter_by_module(root_filename)

        # rebase the basic blocks
        base = idaapi.get_imagebase()
        rebased_blocks = rebase_blocks(base, coverage_blocks)

        # flatten the basic blocks into individual instructions or addresses
        return metadata.flatten_blocks(rebased_blocks)
Пример #5
0
class Lighthouse(plugin_t):
    """
    The Lighthouse IDA Plugin.
    """

    flags = idaapi.PLUGIN_PROC | idaapi.PLUGIN_MOD
    comment = "Code Coverage Explorer"
    help = ""
    wanted_name = "Lighthouse"
    wanted_hotkey = ""

    def __init__(self):

        # plugin color palette
        self.palette = LighthousePalette()

        #----------------------------------------------------------------------

        # the database coverage data conglomerate
        self.director = CoverageDirector(self.palette)

        # hexrays hooks
        self._hxe_events = None

        # plugin qt elements
        self._ui_coverage_overview = CoverageOverview(self.director)

        # members for the 'Load Code Coverage' menu entry
        self._icon_id_load = idaapi.BADADDR
        self._action_name_load = "lighthouse:load_coverage"

        # members for the 'Coverage Overview' menu entry
        self._icon_id_overview = idaapi.BADADDR
        self._action_name_overview = "lighthouse:coverage_overview"

    #--------------------------------------------------------------------------
    # IDA Plugin Overloads
    #--------------------------------------------------------------------------

    def init(self):
        """
        This is called by IDA when it is loading the plugin.
        """

        # attempt plugin initialization
        try:
            self._install_plugin()

        # failed to initialize or integrate the plugin, log and skip loading
        except Exception as e:
            logger.exception("Failed to initialize")
            return idaapi.PLUGIN_SKIP

        # plugin loaded successfully, print the Lighthouse banner
        self.print_banner()
        logger.info("Successfully initialized")

        # tell IDA to keep the plugin loaded (everything is okay)
        return idaapi.PLUGIN_KEEP

    def run(self, arg):
        """
        This is called by IDA when this file is loaded as a script.
        """
        idaapi.warning("The Lighthouse plugin cannot be run as a script.")

    def term(self):
        """
        This is called by IDA when it is unloading the plugin.
        """

        # attempt to cleanup and uninstall our plugin instance
        try:
            self._uninstall_plugin()

        # failed to cleanly remove the plugin, log failure
        except Exception as e:
            logger.exception("Failed to cleanly unload the plugin")

        logger.info("-" * 75)
        logger.info("Plugin terminated")

    #--------------------------------------------------------------------------
    # Initialization
    #--------------------------------------------------------------------------

    def _install_plugin(self):
        """
        Initialize & integrate the plugin into IDA.
        """
        self._install_ui()

        # NOTE: let's delay these till coverage load instead
        #self._install_hexrays_hooks()

    def _install_hexrays_hooks(self):
        """
        Install Hexrays hook listeners.
        """

        # event hooks appear to already be installed for hexrays
        if self._hxe_events:
            return

        # ensure hexrays is loaded & ready for use
        if not idaapi.init_hexrays_plugin():
            raise RuntimeError("HexRays not available for hooking")

        #
        # map our callback function to an actual member since we can't properly
        # remove bindings from IDA callback registrations otherwise. it also
        # makes installation tracking/status easier.
        #

        self._hxe_events = self._hexrays_callback

        # install the callback handler
        assert idaapi.install_hexrays_callback(self._hxe_events)

    def print_banner(self):
        """
        Print the Lighthouse plugin banner.
        """

        # build the main banner title
        banner_params = (PLUGIN_VERSION, AUTHORS, DATE)
        banner_title = "Lighthouse v%s - (c) %s - %s" % banner_params

        # print plugin banner
        lmsg("")
        lmsg("-" * 75)
        lmsg("---[ %s" % banner_title)
        lmsg("-" * 75)
        lmsg("")

    #--------------------------------------------------------------------------
    # Initialization - UI
    #--------------------------------------------------------------------------

    def _install_ui(self):
        """
        Initialize & integrate all UI elements.
        """

        # install the 'Load Coverage' file dialog
        self._install_load_file_dialog()
        self._install_open_coverage_overview()

    def _install_load_file_dialog(self):
        """
        Install the 'File->Load->Code Coverage File(s)...' menu entry.
        """

        # createa a custom IDA icon
        self._icon_id_load = idaapi.load_custom_icon(
            data=str(open(resource_file("icons/load.png"), "rb").read()))

        # describe the action
        # add an menu entry to the options dropdown on the IDA toolbar
        action_desc = idaapi.action_desc_t(
            self._action_name_load,  # The action name.
            "~C~ode Coverage File(s)...",  # The action text.
            IDACtxEntry(self.load_code_coverage),  # The action handler.
            None,  # Optional: action shortcut
            "Load a code coverage file for this IDB",  # Optional: tooltip
            self._icon_id_load  # Optional: the action icon
        )

        # register the action with IDA
        result = idaapi.register_action(action_desc)
        if not result:
            RuntimeError("Failed to register load coverage action with IDA")

        # attach the action to the File-> dropdown menu
        result = idaapi.attach_action_to_menu(
            "File/Load file/",  # Relative path of where to add the action
            self._action_name_load,  # The action ID (see above)
            idaapi.SETMENU_APP  # We want to append the action after ^
        )
        if not result:
            RuntimeError(
                "Failed to attach load action to 'File/Load file/' dropdown")

        logger.info("Installed the 'Load Code Coverage' menu entry")

    def _install_open_coverage_overview(self):
        """
        Install the 'View->Open subviews->Coverage Overview' menu entry.
        """

        # createa a custom IDA icon
        self._icon_id_overview = idaapi.load_custom_icon(
            data=str(open(resource_file("icons/overview.png"), "rb").read()))

        # describe the action
        # add an menu entry to the options dropdown on the IDA toolbar
        action_desc = idaapi.action_desc_t(
            self._action_name_overview,  # The action name.
            "~C~overage Overview",  # The action text.
            IDACtxEntry(self.open_coverage_overview),  # The action handler.
            None,  # Optional: action shortcut
            "Open database code coverage overview",  # Optional: tooltip
            self._icon_id_overview  # Optional: the action icon
        )

        # register the action with IDA
        result = idaapi.register_action(action_desc)
        if not result:
            RuntimeError(
                "Failed to register open coverage overview action with IDA")

        # attach the action to the File-> dropdown menu
        result = idaapi.attach_action_to_menu(
            "View/Open subviews/Hex dump",  # Relative path of where to add the action
            self._action_name_overview,  # The action ID (see above)
            idaapi.SETMENU_INS  # We want to append the action after ^
        )
        if not result:
            RuntimeError("Failed to attach open action to 'subviews' dropdown")

        logger.info("Installed the 'Coverage Overview' menu entry")

    #--------------------------------------------------------------------------
    # Termination
    #--------------------------------------------------------------------------

    def _uninstall_plugin(self):
        """
        Cleanup & uninstall the plugin from IDA.
        """
        self._uninstall_ui()
        self._uninstall_hexrays_hooks()

    def _uninstall_hexrays_hooks(self):
        """
        Cleanup & uninstall Hexrays hook listeners.
        """
        if not self._hxe_events:
            return

        # remove the callbacks
        #    NOTE: w/e IDA removes this anyway.....
        #idaapi.remove_hexrays_callback(self._hxe_events)
        self._hxe_events = None

    #--------------------------------------------------------------------------
    # Termination - UI
    #--------------------------------------------------------------------------

    def _uninstall_ui(self):
        """
        Cleanup & uninstall the plugin UI from IDA.
        """
        self._uninstall_open_coverage_overview()
        self._uninstall_load_file_dialog()

    def _uninstall_load_file_dialog(self):
        """
        Remove the 'File->Load file->Code Coverage File(s)...' menu entry.
        """

        # remove the entry from the File-> menu
        result = idaapi.detach_action_from_menu(
            "File/Load file/",  # Relative path of where we put the action
            self._action_name_load)
        if not result:
            return False

        # unregister the action
        result = idaapi.unregister_action(self._action_name_load)
        if not result:
            return False

        # delete the entry's icon
        idaapi.free_custom_icon(self._icon_id_load)
        self._icon_id_load = idaapi.BADADDR

        logger.info("Uninstalled the 'Load Code Coverage' menu entry")

    def _uninstall_open_coverage_overview(self):
        """
        Remove the 'View->Open subviews->Coverage Overview' menu entry.
        """

        # remove the entry from the View-> menu
        result = idaapi.detach_action_from_menu(
            "View/Open subviews/Hex dump",  # Relative path of where we put the action
            self._action_name_overview)
        if not result:
            return False

        # unregister the action
        result = idaapi.unregister_action(self._action_name_overview)
        if not result:
            return False

        # delete the entry's icon
        idaapi.free_custom_icon(self._icon_id_overview)
        self._icon_id_overview = idaapi.BADADDR

        logger.info("Uninstalled the 'Coverage Overview' menu entry")

    #--------------------------------------------------------------------------
    # UI - Actions
    #--------------------------------------------------------------------------

    def load_code_coverage(self):
        """
        Interactive file dialog based loading of Code Coverage.
        """

        # prompt the user with a QtFileDialog to select coverage files
        coverage_files = self._select_code_coverage_files()
        if not coverage_files:
            return

        # TODO: this is okay here for now, but should probably be moved later
        self.palette.refresh_colors()

        #
        # TODO:
        #
        #   I do not hold great confidence in this code yet, so let's wrap
        #   this in a try/catch so the user doesn't get stuck with a wait
        #   box they can't close should things go poorly ;P
        #

        try:

            #
            # collect underlying database metadata so that the plugin core can
            # process, map, and manipulate coverage data in a performant manner.
            #
            # TODO: do this asynchronously as the user is selecting files
            #

            idaapi.show_wait_box("Building database metadata...")
            self.director.refresh()

            #
            # load the selected code coverage files into the plugin core
            #

            idaapi.replace_wait_box("Loading coverage files from disk...")
            for filename in coverage_files:
                self.load_code_coverage_file(filename)
            idaapi.hide_wait_box()

        # 'something happened :('
        except Exception as e:
            idaapi.hide_wait_box()
            lmsg("Failed to load coverage:")
            lmsg("- %s" % e)
            logger.exception(e)
            return

        # select the 'first' coverage file loaded
        self.director.select_coverage(os.path.basename(coverage_files[0]))

        # install hexrays hooks if available for this arch/install
        try:
            self._install_hexrays_hooks()
        except RuntimeError:
            pass

        # show the coverage overview
        self.open_coverage_overview()

    def open_coverage_overview(self):
        """
        Open the Coverage Overview dialog.
        """
        self._ui_coverage_overview.Show()

    def _select_code_coverage_files(self):
        """
        Open the 'Load Code Coverage' dialog and capture file selections.
        """

        # create & configure a Qt File Dialog for immediate use
        file_dialog = QtWidgets.QFileDialog(None, 'Open Code Coverage File(s)')
        file_dialog.setFileMode(QtWidgets.QFileDialog.ExistingFiles)

        # prompt the user with the file dialog, and await filename(s)
        filenames, _ = file_dialog.getOpenFileNames()

        # log the captured (selected) filenames from the dialog
        logger.debug("Captured filenames from file dialog:")
        logger.debug(filenames)

        # return the captured filenames
        return filenames

    #--------------------------------------------------------------------------
    # Misc
    #--------------------------------------------------------------------------

    def load_code_coverage_file(self, filename):
        """
        Load code coverage file by filename.

        NOTE: At this time only binary drcov logs are supported.
        """
        basename = os.path.basename(filename)

        # load coverage data from file
        coverage_data = DrcovData(filename)

        # extract the coverage relevant to this IDB (well, the root binary)
        root_filename = idaapi.get_root_filename()
        coverage_blocks = coverage_data.filter_by_module(root_filename)

        # index the coverage for use
        base = idaapi.get_imagebase()
        indexed_data = index_coverage(base, coverage_blocks)

        # enlighten the coverage director to this new coverage data
        self.director.add_coverage(basename, base, indexed_data)

    def _hexrays_callback(self, event, *args):
        """
        HexRays callback event handler.
        """

        # decompilation text generation is complete and it is about to be shown
        if event == idaapi.hxe_text_ready:
            vdui = args[0]
            cfunc = vdui.cfunc

            # if there's no coverage data for this function, there's nothing to do
            if not cfunc.entry_ea in self.director.coverage.functions:
                return 0

            # paint the decompilation text for this function
            paint_hexrays(cfunc, self.director.metadata,
                          self.director.coverage, self.palette.ida_coverage)

        return 0