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

Code before:
        ghidra_data.append(CAPADATA(this_namespace, this_scope, rule, this_locs, this_node, this_attack))

    return ghidra_data


def main():
    logging.basicConfig(level=logging.INFO)
    logging.getLogger().setLevel(logging.INFO)

    if not capa.ghidra.helpers.is_supported_ghidra_version():
        return capa.main.E_UNSUPPORTED_GHIDRA_VERSION

    if not capa.ghidra.helpers.is_supported_file_type():
        return capa.main.E_INVALID_FILE_TYPE

    if not capa.ghidra.helpers.is_supported_arch_type():
        return capa.main.E_INVALID_FILE_ARCH

    if isRunningHeadless():  # type: ignore [name-defined] # noqa: F821
        return 0
    else:
        capa_data = json.loads(get_capabilities())
        for item in parse_json(capa_data):
            item.tag_functions()
            item.bookmark_locations()
        return 0


if __name__ == "__main__":
    if sys.version_info < (3, 8):
        from capa.exceptions import UnsupportedRuntimeError

        raise UnsupportedRuntimeError("This version of capa can only be used with Python 3.8+")
    sys.exit(main())


Review comment (suggestion):
```suggestion
    if not isRunningHeadless():  # type: ignore [name-defined] # noqa: F821
        capa_data = json.loads(get_capabilities())
        for item in parse_json(capa_data):
            item.tag_functions()
            item.bookmark_locations()
    return 0
```

Code after:
                                    # Ex. See "Reference Base64 String" rule
                                    for sub_type, description in parse_node(node):
                                        # in many cases, these will be ghidra-labeled data, so just add the existing
                                        # label symbol to the capa namespace
                                        for sym in symbol_table.getSymbols(sub_ghidra_addr):
                                            if sym.getSymbolType() == SymbolType.LABEL:
                                                sym.setNamespace(capa_namespace)
                                        self.set_pre_comment(sub_ghidra_addr, sub_type, description)


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"
        pass

    if not rules_dir:
        logger.info("You must choose a capa rules directory before running capa.")
        return ""  # return empty str to avoid handling both int and str types

    rules_path: pathlib.Path = pathlib.Path(rules_dir)
    logger.info("running capa using rules from %s", str(rules_path))

    rules = capa.main.get_rules([rules_path])
    meta = capa.ghidra.helpers.collect_metadata([rules_path])
    extractor = capa.features.extractors.ghidra.extractor.GhidraFeatureExtractor()

    capabilities, counts = capa.main.find_capabilities(rules, extractor, True)

    if capa.main.has_file_limitation(rules, capabilities, is_standalone=False):
        popup("capa explorer encountered warnings during analysis. Please check the console output for more information.")  # type: ignore [name-defined] # noqa: F821
        logger.info("capa encountered warnings during analysis")

    return capa.render.json.render(meta, rules, capabilities)


def get_locations(match_dict):
    """recursively collect match addresses and associated nodes"""

    for loc in match_dict.get("locations", {}):
        # either an rva (absolute)
        # or an offset into a file (file)
        if loc.get("type", "") in ("absolute", "file"):
            yield loc.get("value"), match_dict.get("node")

    for child in match_dict.get("children", {}):
        yield from get_locations(child)