Repository: mandiant/capa
File: capa/ghidra/capa_explorer.py (Python)

Code before:
                mbc_txt = item.get("objective") + Namespace.DELIMITER + item.get("id")
                add_bookmark(addr, mbc_txt, "CapaExplorer::MBC")

    def bookmark_locations(self):
        """bookmark & label findings at all scopes"""
        symbol_table = currentProgram().getSymbolTable()  # type: ignore [name-defined] # noqa: F821
        name_space = self.create_namespace()

        for addr in self.locations:
            txt = self.capability.replace(" ", "-")
            a = toAddr(hex(addr))  # type: ignore [name-defined] # noqa: F821

            add_bookmark(a, txt)

            # avoid re-naming functions
            # to namespace/rule names
            to_cont = False
            for sym in symbol_table.getSymbols(a):
                if sym.getSymbolType() == SymbolType.FUNCTION:
                    to_cont = True
                    txt = sym.getName()
                    # label to classify function under capa-generated namespace
                    createLabel(a, txt, name_space, True, SourceType.USER_DEFINED)  # type: ignore [name-defined] # noqa: F821

            if to_cont:
                continue

            # greedily create labels at basic block & insn scopes
            node_to_parse = self.node[self.locations.index(addr)]
            if node_to_parse:
                txt = self.recurse_node(node_to_parse)

                if not txt:
                    continue

                txt = txt.replace(" ", "-")
                createLabel(a, txt, name_space, True, SourceType.USER_DEFINED)  # type: ignore [name-defined] # noqa: F821
                continue

            # handle first address of match
            createLabel(a, txt, name_space, True, SourceType.USER_DEFINED)  # type: ignore [name-defined] # noqa: F821


def get_capabilities():
    rules_dir: str = ""
    try:
        selected_dir = askDirectory("Choose capa rules directory", "Ok")  # type: ignore [name-defined] # noqa: F821
        if selected_dir:
            rules_dir = selected_dir.getPath()
    except RuntimeError:
        # RuntimeError thrown when user selects "Cancel"

Review comment (suggestion):
Let's rename `to_cont` to something more descriptive, this check is confusing as it currently stands.

Code after:

                if self.mbc != []:
                    for item in self.mbc:
                        mbc_txt = ""
                        for part in item.get("parts", {}):
                            mbc_txt = mbc_txt + part + Namespace.DELIMITER
                        mbc_txt = mbc_txt + item.get("id", {})
                        add_bookmark(func_addr, mbc_txt, "CapaExplorer::MBC")

    def set_plate_comment(self, ghidra_addr):
        """set plate comments at matched functions"""
        comment = getPlateComment(ghidra_addr)  # type: ignore [name-defined] # noqa: F821
        rule_path = self.namespace.replace(Namespace.DELIMITER, "/")
        # 2 calls to avoid duplicate comments via subsequent script runs
        if comment is None:
            # first comment @ function
            comment = rule_path + "\n"
            setPlateComment(ghidra_addr, comment)  # type: ignore [name-defined] # noqa: F821
        elif rule_path not in comment:
            comment = comment + rule_path + "\n"
            setPlateComment(ghidra_addr, comment)  # type: ignore [name-defined] # noqa: F821
        else:
            return

    def set_pre_comment(self, ghidra_addr, sub_type, description):
        """set pre comments at subscoped matches of main rules"""
        comment = getPreComment(ghidra_addr)  # type: ignore [name-defined] # noqa: F821
        if comment is None:
            comment = "capa: " + sub_type + "(" + description + ")" + ' matched in "' + self.capability + '"\n'
            setPreComment(ghidra_addr, comment)  # type: ignore [name-defined] # noqa: F821
        elif self.capability not in comment:
            comment = (
                comment + "capa: " + sub_type + "(" + description + ")" + ' matched in "' + self.capability + '"\n'
            )
            setPreComment(ghidra_addr, comment)  # type: ignore [name-defined] # noqa: F821
        else:
            return

    def label_matches(self):
        """label findings at function scopes and comment on subscope matches"""
        capa_namespace = create_namespace(self.namespace)
        symbol_table = currentProgram().getSymbolTable()  # type: ignore [name-defined] # noqa: F821

        # handle function main scope of matched rule
        # these will typically contain further matches within
        if self.scope == "function":
            for addr in self.matches.keys():
                ghidra_addr = toAddr(hex(addr))  # type: ignore [name-defined] # noqa: F821

                # classify new function label under capa-generated namespace
                sym = symbol_table.getPrimarySymbol(ghidra_addr)