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

Code before:
import capa.render.json
import capa.ghidra.helpers
import capa.features.extractors.ghidra.extractor

logger = logging.getLogger("capa_explorer")


class CAPADATA:
    def __init__(self, namespace, scope, capability, locations, node, attack=None):
        self.namespace = namespace
        self.scope = scope
        self.capability = capability
        self.locations = locations
        self.node = node
        self.attack = attack

    def recurse_node(self, node_dict):
        """pull string data by recursing node dicts"""

        if not node_dict:
            # KeyError
            # ex. {'type':'subscope', 'scope':'basic block'}
            return ""

        if isinstance(node_dict, int):
            node_dict = hex(node_dict)

        if isinstance(node_dict, str):
            return node_dict

        if "description" in node_dict:
            return node_dict.get("description")
        else:
            return self.recurse_node(node_dict.get(node_dict.get("type")))

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

        # handle rules w/o namespace -> capa lib rule
        if self.namespace == "capa":
            lib_str = Namespace.DELIMITER + "lib" + Namespace.DELIMITER
            self.namespace = self.namespace + lib_str + self.capability.replace(" ", "-")
        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


Review comment (suggestion):
Should `node_dict` be returned here? Otherwise, we attempt to access `node_dict` as a Python Dictionary on line 58. 

Code after:
import capa.ghidra.helpers
import capa.features.extractors.ghidra.extractor

logger = logging.getLogger("capa_explorer")


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


def create_namespace(namespace_str):
    """create new Ghidra namespace for each capa namespace"""

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


def create_label(ghidra_addr, name, capa_namespace):
    """custom label cmd to overlay symbols under capa-generated namespaces"""

    # prevent duplicate labels under the same capa-generated namespace
    symbol_table = currentProgram().getSymbolTable()  # type: ignore [name-defined] # noqa: F821
    for sym in symbol_table.getSymbols(ghidra_addr):
        if sym.getName(True) == capa_namespace.getName(True) + Namespace.DELIMITER + name:
            return

    # create SymbolType.LABEL at addr
    # prioritize capa-generated namespace (duplicate match @ new addr), else put under global Ghidra one (new match)
    cmd = AddLabelCmd(ghidra_addr, name, True, SourceType.USER_DEFINED)
    cmd.applyTo(currentProgram())  # type: ignore [name-defined] # noqa: F821

    # assign new match overlay label to capa-generated namespace
    cmd.getSymbol().setNamespace(capa_namespace)
    return


class CapaMatchData:
    def __init__(
        self,
        namespace,
        scope,
        capability,
        matches,
        attack: List[Dict[Any, Any]],
        mbc: List[Dict[Any, Any]],
    ):
        self.namespace = namespace
        self.scope = scope
        self.capability = capability