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

Code before:
        cmd = CreateNamespacesCmd(self.namespace, SourceType.USER_DEFINED)
        cmd.applyTo(currentProgram())  # type: ignore [name-defined] # noqa: F821

        return cmd.getNamespace()

    def add_bookmark(self, addr, txt, category="CapaExplorer"):
        """create bookmark at addr"""
        currentProgram().getBookmarkManager().setBookmark(addr, "Info", category, txt)  # type: ignore [name-defined] # noqa: F821

    def tag_functions(self):
        """create function tags for capabilities"""

        # self.locations[0] will always be the largest
        # scoped offset yielded i.e. closest to an entrypoint
        addr = toAddr(hex(self.locations[0]))  # type: ignore [name-defined] # noqa: F821
        f = getFunctionContaining(addr)  # type: ignore [name-defined] # noqa: F821

        # bookmark Mitre ATT&CK tactics @ function scope
        if f:
            f.addTag(self.capability)
            if self.attack:
                for item in self.attack:
                    attack_txt = item.get("tactic") + Namespace.DELIMITER + item.get("id")
                    self.add_bookmark(addr, attack_txt, "CapaExplorer::Mitre ATT&CK")

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

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

            self.add_bookmark(a, txt)

            # avoid re-naming functions
            to_cont = False
            for s in st.getSymbols(a):
                if s.getSymbolType() == SymbolType.FUNCTION:
                    to_cont = True
                    txt = s.getName()
                    createLabel(a, txt, ns, True, SourceType.USER_DEFINED)  # type: ignore [name-defined] # noqa: F821

            if to_cont:
                continue

            # 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)

Review comment (suggestion):
Let's use more verbose local variable names in this function to improve readability e.g. `st` -> `symbol_table`

Code after:
        scope,
        capability,
        matches,
        attack: List[Dict[Any, Any]],
        mbc: List[Dict[Any, Any]],
    ):
        self.namespace = namespace
        self.scope = scope
        self.capability = capability
        self.matches = matches
        self.attack = attack
        self.mbc = mbc

    def bookmark_functions(self):
        """create bookmarks for MITRE ATT&CK & MBC mappings"""

        if self.attack == [] and self.mbc == []:
            return

        for key in self.matches.keys():
            addr = toAddr(hex(key))  # type: ignore [name-defined] # noqa: F821
            func = getFunctionContaining(addr)  # type: ignore [name-defined] # noqa: F821

            # bookmark & tag MITRE ATT&CK tactics & MBC @ function scope
            if func is not None:
                func_addr = func.getEntryPoint()

                if self.attack != []:
                    for item in self.attack:
                        attack_txt = ""
                        for part in item.get("parts", {}):
                            attack_txt = attack_txt + part + Namespace.DELIMITER
                        attack_txt = attack_txt + item.get("id", {})
                        add_bookmark(func_addr, attack_txt, "CapaExplorer::MITRE ATT&CK")

                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"