From 3a44b840be04312b22b94493546025977a4e9120 Mon Sep 17 00:00:00 2001 From: WerWolv Date: Thu, 22 Feb 2024 20:49:21 +0100 Subject: [PATCH] impr: Restructure yara rule handling again --- .../include/content/views/view_yara.hpp | 13 +- .../yara_rules/include/content/yara_rule.hpp | 11 +- .../content/data_information_sections.cpp | 37 ++++-- .../source/content/views/view_yara.cpp | 111 +++++------------- .../yara_rules/source/content/yara_rule.cpp | 37 ++++-- 5 files changed, 94 insertions(+), 115 deletions(-) diff --git a/plugins/yara_rules/include/content/views/view_yara.hpp b/plugins/yara_rules/include/content/views/view_yara.hpp index ad7967b52..e2dcd7689 100644 --- a/plugins/yara_rules/include/content/views/view_yara.hpp +++ b/plugins/yara_rules/include/content/views/view_yara.hpp @@ -17,21 +17,14 @@ namespace hex::plugin::yara { void drawContent() override; - private: - struct YaraMatch { - YaraRule::Match match; - - mutable u32 highlightId; - mutable u32 tooltipId; - }; - private: PerProvider>> m_rulePaths; - PerProvider> m_matches; - PerProvider> m_sortedMatches; + PerProvider> m_matchedRules; PerProvider> m_consoleMessages; PerProvider m_selectedRule; + PerProvider> m_tooltipIds, m_highlightingIds; + TaskHolder m_matcherTask; void applyRules(); diff --git a/plugins/yara_rules/include/content/yara_rule.hpp b/plugins/yara_rules/include/content/yara_rule.hpp index f0910968d..5738fd1a4 100644 --- a/plugins/yara_rules/include/content/yara_rule.hpp +++ b/plugins/yara_rules/include/content/yara_rule.hpp @@ -18,14 +18,21 @@ namespace hex::plugin::yara { static void cleanup(); struct Match { - std::string identifier; std::string variable; Region region; bool wholeDataMatch; }; - struct Result { + struct Rule { + std::string identifier; + std::map metadata; + std::vector tags; + std::vector matches; + }; + + struct Result { + std::vector matchedRules; std::vector consoleMessages; }; diff --git a/plugins/yara_rules/source/content/data_information_sections.cpp b/plugins/yara_rules/source/content/data_information_sections.cpp index dd778ceb4..77124d4db 100644 --- a/plugins/yara_rules/source/content/data_information_sections.cpp +++ b/plugins/yara_rules/source/content/data_information_sections.cpp @@ -16,6 +16,16 @@ namespace hex::plugin::yara { InformationYaraRules() : InformationSection("hex.yara.information_section.yara_rules") { } ~InformationYaraRules() override = default; + struct Category { + struct Comperator { + bool operator()(const YaraRule::Rule &a, const YaraRule::Rule &b) const { + return a.identifier < b.identifier; + } + }; + + std::set matchedRules; + }; + void process(Task &task, prv::Provider *provider, Region region) override { const auto &ruleFilePaths = romfs::list("rules"); task.setMaxValue(ruleFilePaths.size()); @@ -24,10 +34,16 @@ namespace hex::plugin::yara { for (const auto &ruleFilePath : ruleFilePaths) { const std::string fileContent = romfs::get(ruleFilePath).data(); - YaraRule rule(fileContent); - auto result = rule.match(provider, region); + YaraRule yaraRule(fileContent); + const auto result = yaraRule.match(provider, region); if (result.has_value()) { - m_matches[ruleFilePath.filename().string()] = result.value().matches; + const auto &rules = result.value().matchedRules; + for (const auto &rule : rules) { + if (!rule.metadata.contains("category")) continue; + + const auto &categoryName = rule.metadata.at("category"); + m_categories[categoryName].matchedRules.insert(rule); + } } task.update(progress); @@ -36,7 +52,7 @@ namespace hex::plugin::yara { } void reset() override { - m_matches.clear(); + m_categories.clear(); } void drawContent() override { @@ -46,15 +62,16 @@ namespace hex::plugin::yara { ImGui::TableNextRow(); - for (auto &[name, matches] : m_matches) { - if (matches.empty()) + for (auto &[categoryName, category] : m_categories) { + if (category.matchedRules.empty()) continue; ImGui::TableNextColumn(); - ImGuiExt::BeginSubWindow(name.c_str()); + ImGuiExt::BeginSubWindow(categoryName.c_str()); { - for (const auto &match : matches) { - ImGui::TextUnformatted(match.identifier.c_str()); + for (const auto &match : category.matchedRules) { + const auto &ruleName = match.metadata.contains("name") ? match.metadata.at("name") : match.identifier; + ImGui::TextUnformatted(ruleName.c_str()); } } ImGuiExt::EndSubWindow(); @@ -65,7 +82,7 @@ namespace hex::plugin::yara { } private: - std::map> m_matches; + std::map m_categories; }; void registerDataInformationSections() { diff --git a/plugins/yara_rules/source/content/views/view_yara.cpp b/plugins/yara_rules/source/content/views/view_yara.cpp index 3a518fdcb..b4f039b44 100644 --- a/plugins/yara_rules/source/content/views/view_yara.cpp +++ b/plugins/yara_rules/source/content/views/view_yara.cpp @@ -48,7 +48,7 @@ namespace hex::plugin::yara { if (!rules.is_array()) return false; - m_matches.get(provider).clear(); + m_matchedRules.get(provider).clear(); for (auto &rule : rules) { if (!rule.contains("name") || !rule.contains("path")) @@ -130,7 +130,6 @@ namespace hex::plugin::yara { ImGui::NewLine(); if (ImGui::Button("hex.yara_rules.view.yara.match"_lang)) this->applyRules(); - ImGui::SameLine(); if (m_matcherTask.isRunning()) { ImGui::SameLine(); @@ -143,76 +142,34 @@ namespace hex::plugin::yara { matchesTableSize.y *= 3.75 / 5.0; matchesTableSize.y -= ImGui::GetTextLineHeightWithSpacing(); - if (ImGui::BeginTable("matches", 4, ImGuiTableFlags_Borders | ImGuiTableFlags_Resizable | ImGuiTableFlags_Sortable | ImGuiTableFlags_Reorderable | ImGuiTableFlags_RowBg | ImGuiTableFlags_ScrollY, matchesTableSize)) { + if (ImGui::BeginTable("matches", 3, ImGuiTableFlags_Borders | ImGuiTableFlags_Resizable | ImGuiTableFlags_Sortable | ImGuiTableFlags_Reorderable | ImGuiTableFlags_RowBg | ImGuiTableFlags_ScrollY, matchesTableSize)) { ImGui::TableSetupScrollFreeze(0, 1); - ImGui::TableSetupColumn("hex.yara_rules.view.yara.matches.identifier"_lang, ImGuiTableColumnFlags_PreferSortAscending, 0, ImGui::GetID("identifier")); ImGui::TableSetupColumn("hex.yara_rules.view.yara.matches.variable"_lang, ImGuiTableColumnFlags_PreferSortAscending, 0, ImGui::GetID("variable")); ImGui::TableSetupColumn("hex.ui.common.address"_lang, ImGuiTableColumnFlags_PreferSortAscending, 0, ImGui::GetID("address")); ImGui::TableSetupColumn("hex.ui.common.size"_lang, ImGuiTableColumnFlags_PreferSortAscending, 0, ImGui::GetID("size")); ImGui::TableHeadersRow(); - auto sortSpecs = ImGui::TableGetSortSpecs(); - if (!m_matches->empty() && (sortSpecs->SpecsDirty || m_sortedMatches->empty())) { - m_sortedMatches->clear(); - std::transform(m_matches->begin(), m_matches->end(), std::back_inserter(*m_sortedMatches), [](auto &match) { - return &match; - }); - - std::sort(m_sortedMatches->begin(), m_sortedMatches->end(), [&sortSpecs](const YaraMatch *left, const YaraMatch *right) -> bool { - if (sortSpecs->Specs->ColumnUserID == ImGui::GetID("identifier")) - return left->match.identifier < right->match.identifier; - else if (sortSpecs->Specs->ColumnUserID == ImGui::GetID("variable")) - return left->match.variable < right->match.variable; - else if (sortSpecs->Specs->ColumnUserID == ImGui::GetID("address")) - return left->match.region.getStartAddress() < right->match.region.getStartAddress(); - else if (sortSpecs->Specs->ColumnUserID == ImGui::GetID("size")) - return left->match.region.getSize() < right->match.region.getSize(); - else - return false; - }); - - if (sortSpecs->Specs->SortDirection == ImGuiSortDirection_Descending) - std::reverse(m_sortedMatches->begin(), m_sortedMatches->end()); - - sortSpecs->SpecsDirty = false; - } - if (!m_matcherTask.isRunning()) { - ImGuiListClipper clipper; - clipper.Begin(m_sortedMatches->size()); + for (const auto &rule : *m_matchedRules) { + if (rule.matches.empty()) continue; - while (clipper.Step()) { - for (int i = clipper.DisplayStart; i < clipper.DisplayEnd; i++) { - auto &[match, highlightId, tooltipId] = *(*m_sortedMatches)[i]; - auto &[identifier, variableName, region, wholeDataMatch] = match; - ImGui::TableNextRow(); - ImGui::TableNextColumn(); - ImGui::PushID(i); - if (ImGui::Selectable("match", false, ImGuiSelectableFlags_SpanAllColumns | ImGuiSelectableFlags_AllowOverlap)) { - ImHexApi::HexEditor::setSelection(region.getStartAddress(), region.getSize()); - } - ImGui::PopID(); - ImGui::SameLine(); - ImGui::TextUnformatted(identifier.c_str()); - ImGui::TableNextColumn(); - ImGui::TextUnformatted(variableName.c_str()); + ImGui::TableNextRow(); + ImGui::TableNextColumn(); - if (!wholeDataMatch) { + if (ImGui::TreeNode(rule.identifier.c_str())) { + for (const auto &match : rule.matches) { + ImGui::TableNextRow(); ImGui::TableNextColumn(); - ImGuiExt::TextFormatted("0x{0:X} : 0x{1:X}", region.getStartAddress(), region.getEndAddress()); + ImGui::TextUnformatted(match.variable.c_str()); ImGui::TableNextColumn(); - ImGuiExt::TextFormatted("0x{0:X}", region.getSize()); - } else { + ImGui::TextUnformatted(hex::format("0x{0:08X}", match.region.getStartAddress()).c_str()); ImGui::TableNextColumn(); - ImGuiExt::TextFormattedColored(ImVec4(0.92F, 0.25F, 0.2F, 1.0F), "{}", "hex.yara_rules.view.yara.whole_data"_lang); - ImGui::TableNextColumn(); - ImGui::TextUnformatted(""); + ImGui::TextUnformatted(hex::format("0x{0:08X}", match.region.getSize()).c_str()); } + ImGui::TreePop(); } } - - clipper.End(); } ImGui::EndTable(); @@ -237,12 +194,14 @@ namespace hex::plugin::yara { } void ViewYara::clearResult() { - for (const auto &match : *m_matches) { - ImHexApi::HexEditor::removeBackgroundHighlight(match.highlightId); - ImHexApi::HexEditor::removeTooltip(match.tooltipId); - } + for (const auto &id : *m_highlightingIds) + ImHexApi::HexEditor::removeBackgroundHighlight(id); + for (const auto &id : *m_tooltipIds) + ImHexApi::HexEditor::removeTooltip(id); - m_matches->clear(); + m_highlightingIds->clear(); + m_tooltipIds->clear(); + m_matchedRules->clear(); m_consoleMessages->clear(); } @@ -280,33 +239,19 @@ namespace hex::plugin::yara { } TaskManager::doLater([this, results = std::move(results)] { - for (const auto &match : *m_matches) { - ImHexApi::HexEditor::removeBackgroundHighlight(match.highlightId); - ImHexApi::HexEditor::removeTooltip(match.tooltipId); - } + this->clearResult(); for (auto &result : results) { - for (auto &match : result.matches) { - m_matches->emplace_back(YaraMatch { std::move(match), 0, 0 }); - } - - m_consoleMessages->append_range(result.consoleMessages); + m_matchedRules->append_range(std::move(result.matchedRules)); + m_consoleMessages->append_range(std::move(result.consoleMessages)); } - auto uniques = std::set(m_matches->begin(), m_matches->end(), [](const auto &leftMatch, const auto &rightMatch) { - const auto &l = leftMatch.match; - const auto &r = rightMatch.match; - return std::tie(l.region.address, l.wholeDataMatch, l.identifier, l.variable) < - std::tie(r.region.address, r.wholeDataMatch, r.identifier, r.variable); - }); - - m_matches->clear(); - std::move(uniques.begin(), uniques.end(), std::back_inserter(*m_matches)); - constexpr static color_t YaraColor = 0x70B4771F; - for (const YaraMatch &yaraMatch : uniques) { - yaraMatch.highlightId = ImHexApi::HexEditor::addBackgroundHighlight(yaraMatch.match.region, YaraColor); - yaraMatch.tooltipId = ImHexApi::HexEditor::addTooltip(yaraMatch.match.region, hex::format("{0} [{1}]", yaraMatch.match.identifier, yaraMatch.match.variable), YaraColor); + for (YaraRule::Rule &rule : *m_matchedRules) { + for (auto &match : rule.matches) { + m_highlightingIds->push_back(ImHexApi::HexEditor::addBackgroundHighlight(match.region, YaraColor)); + m_tooltipIds->push_back(ImHexApi::HexEditor::addTooltip(match.region, hex::format("{0} : {1} [{2}]", rule.identifier, fmt::join(rule.tags, ", "), match.variable), YaraColor)); + } } }); }); diff --git a/plugins/yara_rules/source/content/yara_rule.cpp b/plugins/yara_rules/source/content/yara_rule.cpp index fb8aa5c58..0b5404310 100644 --- a/plugins/yara_rules/source/content/yara_rule.cpp +++ b/plugins/yara_rules/source/content/yara_rule.cpp @@ -17,7 +17,7 @@ namespace hex::plugin::yara { struct ResultContext { YaraRule *rule; - std::vector newMatches; + std::vector matchedRules; std::vector consoleMessages; std::string includeBuffer; @@ -42,35 +42,52 @@ namespace hex::plugin::yara { } static int scanFunction(YR_SCAN_CONTEXT *context, int message, void *data, void *userData) { - auto &results = *static_cast(userData); + auto &resultContext = *static_cast(userData); switch (message) { case CALLBACK_MSG_RULE_MATCHING: { - auto rule = static_cast(data); + const auto *rule = static_cast(data); if (rule->strings != nullptr) { - YR_STRING *string = nullptr; - YR_MATCH *match = nullptr; + YR_STRING *string; yr_rule_strings_foreach(rule, string) { + YaraRule::Rule newRule; + newRule.identifier = rule->identifier; + + YR_MATCH *match; yr_string_matches_foreach(context, string, match) { - results.newMatches.push_back({ rule->identifier, string->identifier, Region { u64(match->offset), size_t(match->match_length) }, false }); + newRule.matches.push_back({ string->identifier, Region { u64(match->offset), size_t(match->match_length) }, false }); } + + YR_META *meta; + yr_rule_metas_foreach(rule, meta) { + newRule.metadata[meta->identifier] = meta->string; + } + + const char *tag; + yr_rule_tags_foreach(rule, tag) { + newRule.tags.emplace_back(tag); + } + + resultContext.matchedRules.emplace_back(std::move(newRule)); } } else { - results.newMatches.push_back({ rule->identifier, "", Region::Invalid(), true }); + YaraRule::Rule newRule; + newRule.identifier = rule->identifier; + newRule.matches.push_back({ "", Region::Invalid(), true }); } break; } case CALLBACK_MSG_CONSOLE_LOG: { - results.consoleMessages.emplace_back(static_cast(data)); + resultContext.consoleMessages.emplace_back(static_cast(data)); break; } default: break; } - return results.rule->isInterrupted() ? CALLBACK_ABORT : CALLBACK_CONTINUE; + return resultContext.rule->isInterrupted() ? CALLBACK_ABORT : CALLBACK_CONTINUE; } wolv::util::Expected YaraRule::match(prv::Provider *provider, Region region) { @@ -183,7 +200,7 @@ namespace hex::plugin::yara { if (m_interrupted) return wolv::util::Unexpected(Error { Error::Type::Interrupted, "" }); - return Result { resultContext.newMatches, resultContext.consoleMessages }; + return Result { resultContext.matchedRules, resultContext.consoleMessages }; } void YaraRule::interrupt() {