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

Code before:
            # if no "description" key, the "type" key's value is the
            # expected key.
            # ex. {'type':'api', 'api':'RegisterServiceCtrlHandler'}
            return self.recurse_node(node_dict.get(node_dict.get("type")))

    def create_namespace(self):
        """create new ghidra namespace for each capa namespace"""

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

    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
        func = getFunctionContaining(addr)  # type: ignore [name-defined] # noqa: F821

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

            for item in self.mbc:
                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:
            label_name = self.capability.replace(" ", "-")
            ghidra_addr = toAddr(hex(addr))  # type: ignore [name-defined] # noqa: F821

            add_bookmark(ghidra_addr, label_name)

            # avoid renaming user-defined functions
            # to namespace/rule names, since they
            # may contain matches for many rules at this scope
            is_function = False
            for sym in symbol_table.getSymbols(ghidra_addr):
                if sym.getSymbolType() == SymbolType.FUNCTION:
                    is_function = True
                    label_name = sym.getName()
                    # label to classify function under capa-generated namespace

Review comment (suggestion):
```suggestion
                add_bookmark(addr, attack_txt, "CapaExplorer::MITRE ATT&CK")
```

Code after:
        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"
            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