diff --git a/plugins/yara_rules/CMakeLists.txt b/plugins/yara_rules/CMakeLists.txt index 6bec8ab83..b0f5b8f67 100644 --- a/plugins/yara_rules/CMakeLists.txt +++ b/plugins/yara_rules/CMakeLists.txt @@ -15,6 +15,7 @@ add_imhex_plugin( SOURCES source/plugin_yara.cpp + source/content/yara_rule.cpp source/content/views/view_yara.cpp INCLUDES include diff --git a/plugins/yara_rules/include/content/views/view_yara.hpp b/plugins/yara_rules/include/content/views/view_yara.hpp index 8136afc6f..ad7967b52 100644 --- a/plugins/yara_rules/include/content/views/view_yara.hpp +++ b/plugins/yara_rules/include/content/views/view_yara.hpp @@ -6,6 +6,8 @@ #include +#include + namespace hex::plugin::yara { class ViewYara : public View::Window { @@ -17,26 +19,21 @@ namespace hex::plugin::yara { private: struct YaraMatch { - std::string identifier; - std::string variable; - u64 address; - size_t size; - bool wholeDataMatch; + YaraRule::Match match; mutable u32 highlightId; mutable u32 tooltipId; }; private: - PerProvider>> m_rules; + PerProvider>> m_rulePaths; PerProvider> m_matches; PerProvider> m_sortedMatches; + PerProvider> m_consoleMessages; + PerProvider m_selectedRule; - u32 m_selectedRule = 0; TaskHolder m_matcherTask; - std::vector m_consoleMessages; - void applyRules(); void clearResult(); }; diff --git a/plugins/yara_rules/include/content/yara_rule.hpp b/plugins/yara_rules/include/content/yara_rule.hpp new file mode 100644 index 000000000..ea81494e5 --- /dev/null +++ b/plugins/yara_rules/include/content/yara_rule.hpp @@ -0,0 +1,52 @@ +#pragma once + +#include + +#include +#include +#include + +namespace hex::plugin::yara { + + class YaraRule { + public: + YaraRule() = default; + explicit YaraRule(const std::string& content); + explicit YaraRule(const std::fs::path& path); + + static void init(); + static void cleanup(); + + struct Match { + std::string identifier; + std::string variable; + Region region; + bool wholeDataMatch; + }; + + struct Result { + std::vector matches; + std::vector consoleMessages; + }; + + struct Error { + enum class Type { + CompileError, + RuntimeError, + Interrupted + } type; + std::string message; + }; + + wolv::util::Expected match(prv::Provider *provider, u64 address, size_t size); + void interrupt(); + [[nodiscard]] bool isInterrupted() const; + + private: + std::string m_content; + std::fs::path m_filePath; + + std::atomic m_interrupted = false; + }; + +} diff --git a/plugins/yara_rules/source/content/views/view_yara.cpp b/plugins/yara_rules/source/content/views/view_yara.cpp index 0e43210c5..956678dc6 100644 --- a/plugins/yara_rules/source/content/views/view_yara.cpp +++ b/plugins/yara_rules/source/content/views/view_yara.cpp @@ -8,17 +8,9 @@ #include #include -// 's RE type has a zero-sized array, which is not allowed in ISO C++. -#pragma GCC diagnostic push -#pragma GCC diagnostic ignored "-Wpedantic" -#include -#pragma GCC diagnostic pop - #include -#include #include -#include #include namespace hex::plugin::yara { @@ -26,7 +18,7 @@ namespace hex::plugin::yara { using namespace wolv::literals; ViewYara::ViewYara() : View::Window("hex.yara_rules.view.yara.name", ICON_VS_BUG) { - yr_initialize(); + YaraRule::init(); ContentRegistry::FileHandler::add({ ".yar", ".yara" }, [](const auto &path) { for (const auto &destPath : fs::getDefaultPaths(fs::ImHexPath::Yara)) { @@ -68,7 +60,7 @@ namespace hex::plugin::yara { if (!name.is_string() || !path.is_string()) return false; - m_rules.get(provider).emplace_back(std::fs::path(name.get()), std::fs::path(path.get())); + m_rulePaths.get(provider).emplace_back(std::fs::path(name.get()), std::fs::path(path.get())); } return true; @@ -78,7 +70,7 @@ namespace hex::plugin::yara { data["rules"] = nlohmann::json::array(); - for (auto &[name, path] : m_rules.get(provider)) { + for (auto &[name, path] : m_rulePaths.get(provider)) { data["rules"].push_back({ { "name", wolv::util::toUTF8String(name) }, { "path", wolv::util::toUTF8String(path) } @@ -93,17 +85,17 @@ namespace hex::plugin::yara { } ViewYara::~ViewYara() { - yr_finalize(); + YaraRule::cleanup(); } void ViewYara::drawContent() { ImGuiExt::Header("hex.yara_rules.view.yara.header.rules"_lang, true); if (ImGui::BeginListBox("##rules", ImVec2(-FLT_MIN, ImGui::GetTextLineHeightWithSpacing() * 5))) { - for (u32 i = 0; i < m_rules->size(); i++) { - const bool selected = (m_selectedRule == i); - if (ImGui::Selectable(wolv::util::toUTF8String((*m_rules)[i].first).c_str(), selected)) { - m_selectedRule = i; + for (u32 i = 0; i < m_rulePaths->size(); i++) { + const bool selected = (*m_selectedRule == i); + if (ImGui::Selectable(wolv::util::toUTF8String((*m_rulePaths)[i].first).c_str(), selected)) { + *m_selectedRule = i; } } ImGui::EndListBox(); @@ -124,15 +116,15 @@ namespace hex::plugin::yara { ui::PopupFileChooser::open(basePaths, paths, std::vector{ { "Yara File", "yara" }, { "Yara File", "yar" } }, true, [&](const auto &path) { - m_rules->push_back({ path.filename(), path }); + m_rulePaths->push_back({ path.filename(), path }); }); } ImGui::SameLine(); if (ImGuiExt::IconButton(ICON_VS_REMOVE, ImGui::GetStyleColorVec4(ImGuiCol_Text))) { - if (m_selectedRule < m_rules->size()) { - m_rules->erase(m_rules->begin() + m_selectedRule); - m_selectedRule = std::min(m_selectedRule, u32(m_rules->size() - 1)); + if (*m_selectedRule < m_rulePaths->size()) { + m_rulePaths->erase(m_rulePaths->begin() + *m_selectedRule); + m_selectedRule = std::min(*m_selectedRule, u32(m_rulePaths->size() - 1)); } } @@ -169,13 +161,13 @@ namespace hex::plugin::yara { 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->identifier < right->identifier; + return left->match.identifier < right->match.identifier; else if (sortSpecs->Specs->ColumnUserID == ImGui::GetID("variable")) - return left->variable < right->variable; + return left->match.variable < right->match.variable; else if (sortSpecs->Specs->ColumnUserID == ImGui::GetID("address")) - return left->address < right->address; + return left->match.region.getStartAddress() < right->match.region.getStartAddress(); else if (sortSpecs->Specs->ColumnUserID == ImGui::GetID("size")) - return left->size < right->size; + return left->match.region.getSize() < right->match.region.getSize(); else return false; }); @@ -192,12 +184,13 @@ namespace hex::plugin::yara { while (clipper.Step()) { for (int i = clipper.DisplayStart; i < clipper.DisplayEnd; i++) { - auto &[identifier, variableName, address, size, wholeDataMatch, highlightId, tooltipId] = *(*m_sortedMatches)[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(address, size); + ImHexApi::HexEditor::setSelection(region.getStartAddress(), region.getSize()); } ImGui::PopID(); ImGui::SameLine(); @@ -207,9 +200,9 @@ namespace hex::plugin::yara { if (!wholeDataMatch) { ImGui::TableNextColumn(); - ImGuiExt::TextFormatted("0x{0:X} : 0x{1:X}", address, address + size - 1); + ImGuiExt::TextFormatted("0x{0:X} : 0x{1:X}", region.getStartAddress(), region.getEndAddress()); ImGui::TableNextColumn(); - ImGuiExt::TextFormatted("0x{0:X}", size); + ImGuiExt::TextFormatted("0x{0:X}", region.getSize()); } else { ImGui::TableNextColumn(); ImGuiExt::TextFormattedColored(ImVec4(0.92F, 0.25F, 0.2F, 1.0F), "{}", "hex.yara_rules.view.yara.whole_data"_lang); @@ -230,10 +223,10 @@ namespace hex::plugin::yara { if (ImGui::BeginChild("##console", consoleSize, true, ImGuiWindowFlags_AlwaysVerticalScrollbar | ImGuiWindowFlags_HorizontalScrollbar)) { ImGuiListClipper clipper; - clipper.Begin(m_consoleMessages.size()); + clipper.Begin(m_consoleMessages->size()); while (clipper.Step()) { for (int i = clipper.DisplayStart; i < clipper.DisplayEnd; i++) { - const auto &message = m_consoleMessages[i]; + const auto &message = m_consoleMessages->at(i); if (ImGui::Selectable(message.c_str())) ImGui::SetClipboardText(message.c_str()); @@ -250,193 +243,70 @@ namespace hex::plugin::yara { } m_matches->clear(); - m_consoleMessages.clear(); + m_consoleMessages->clear(); } void ViewYara::applyRules() { this->clearResult(); - m_matcherTask = TaskManager::createTask("hex.yara_rules.view.yara.matching", 0, [this](auto &task) { - if (!ImHexApi::Provider::isValid()) return; + auto provider = ImHexApi::Provider::get(); + if (provider == nullptr) + return; - struct ResultContext { - Task *task = nullptr; - std::vector newMatches; - std::vector consoleMessages; - }; + m_matcherTask = TaskManager::createTask("hex.yara_rules.view.yara.matching", 0, [this, provider](auto &task) { + u32 progress = 0; - ResultContext resultContext; - resultContext.task = &task; + std::vector results; + for (const auto &[fileName, filePath] : *m_rulePaths) { + YaraRule rule(filePath); - for (const auto &[fileName, filePath] : *m_rules) { - YR_COMPILER *compiler = nullptr; - yr_compiler_create(&compiler); - ON_SCOPE_EXIT { - yr_compiler_destroy(compiler); - }; + task.setInterruptCallback([&rule] { + rule.interrupt(); + }); - auto currFilePath = wolv::util::toUTF8String(wolv::io::fs::toShortPath(filePath)); - - yr_compiler_set_include_callback( - compiler, - [](const char *includeName, const char *, const char *, void *userData) -> const char * { - wolv::io::File file(std::fs::path(static_cast(userData)).parent_path() / includeName, wolv::io::File::Mode::Read); - if (!file.isValid()) - return nullptr; - - auto size = file.getSize(); - char *buffer = new char[size + 1]; - file.readBuffer(reinterpret_cast(buffer), size); - buffer[size] = 0x00; - - return buffer; - }, - [](const char *ptr, void *userData) { - hex::unused(userData); - - delete[] ptr; - }, - currFilePath.data() - ); - - wolv::io::File file((*m_rules)[m_selectedRule].second, wolv::io::File::Mode::Read); - if (!file.isValid()) return; - - if (yr_compiler_add_file(compiler, file.getHandle(), nullptr, nullptr) != 0) { - std::string errorMessage(0xFFFF, '\x00'); - yr_compiler_get_error_message(compiler, errorMessage.data(), errorMessage.size()); - - TaskManager::doLater([this, errorMessage = wolv::util::trim(errorMessage)] { - this->clearResult(); - - m_consoleMessages.push_back(hex::format("hex.yara_rules.view.yara.error"_lang, errorMessage)); + auto result = rule.match(provider, provider->getBaseAddress(), provider->getSize()); + if (!result.has_value()) { + TaskManager::doLater([this, error = result.error()] { + m_consoleMessages->emplace_back(error.message); }); - return; + break; } - YR_RULES *yaraRules; - yr_compiler_get_rules(compiler, &yaraRules); - ON_SCOPE_EXIT { yr_rules_destroy(yaraRules); }; - - YR_MEMORY_BLOCK_ITERATOR iterator; - - struct ScanContext { - Task *task = nullptr; - std::vector buffer; - YR_MEMORY_BLOCK currBlock = {}; - }; - - ScanContext context; - context.task = &task; - context.currBlock.base = 0; - context.currBlock.fetch_data = [](auto *block) -> const u8 * { - auto &context = *static_cast(block->context); - auto provider = ImHexApi::Provider::get(); - - context.buffer.resize(context.currBlock.size); - - if (context.buffer.empty()) - return nullptr; - - block->size = context.currBlock.size; - provider->read(context.currBlock.base + provider->getBaseAddress(), context.buffer.data(), context.buffer.size()); - - return context.buffer.data(); - }; - iterator.file_size = [](auto *iterator) -> u64 { - hex::unused(iterator); - - return ImHexApi::Provider::get()->getActualSize(); - }; - - iterator.context = &context; - iterator.first = [](YR_MEMORY_BLOCK_ITERATOR *iterator) -> YR_MEMORY_BLOCK *{ - auto &context = *static_cast(iterator->context); - - context.currBlock.base = 0; - context.currBlock.size = 0; - context.buffer.clear(); - iterator->last_error = ERROR_SUCCESS; - - return iterator->next(iterator); - }; - iterator.next = [](YR_MEMORY_BLOCK_ITERATOR *iterator) -> YR_MEMORY_BLOCK * { - auto &context = *static_cast(iterator->context); - - u64 address = context.currBlock.base + context.currBlock.size; - - iterator->last_error = ERROR_SUCCESS; - context.currBlock.base = address; - context.currBlock.size = std::min(ImHexApi::Provider::get()->getActualSize() - address, 10_MiB); - context.currBlock.context = &context; - context.task->update(address); - - if (context.currBlock.size == 0) return nullptr; - - return &context.currBlock; - }; - - yr_rules_scan_mem_blocks( - yaraRules, &iterator, 0, [](YR_SCAN_CONTEXT *context, int message, void *data, void *userData) -> int { - auto &results = *static_cast(userData); - - switch (message) { - case CALLBACK_MSG_RULE_MATCHING: - { - auto rule = static_cast(data); - - - if (rule->strings != nullptr) { - YR_STRING *string = nullptr; - YR_MATCH *match = nullptr; - yr_rule_strings_foreach(rule, string) { - yr_string_matches_foreach(context, string, match) { - results.newMatches.push_back({ rule->identifier, string->identifier, u64(match->offset), size_t(match->match_length), false, 0, 0 }); - } - } - } else { - results.newMatches.push_back({ rule->identifier, "", 0, 0, true, 0, 0 }); - } - } - break; - case CALLBACK_MSG_CONSOLE_LOG: - { - results.consoleMessages.emplace_back(static_cast(data)); - } - break; - default: - break; - } - - return results.task->shouldInterrupt() ? CALLBACK_ABORT : CALLBACK_CONTINUE; - }, - &resultContext, - 0); + results.emplace_back(result.value()); + task.update(progress); + progress += 1; } - TaskManager::doLater([this, resultContext] { + + TaskManager::doLater([this, results = std::move(results)] { for (const auto &match : *m_matches) { ImHexApi::HexEditor::removeBackgroundHighlight(match.highlightId); ImHexApi::HexEditor::removeTooltip(match.tooltipId); } - m_consoleMessages = resultContext.consoleMessages; + for (auto &result : results) { + for (auto &match : result.matches) { + m_matches->emplace_back(YaraMatch { std::move(match), 0, 0 }); + } - std::move(resultContext.newMatches.begin(), resultContext.newMatches.end(), std::back_inserter(*m_matches)); + m_consoleMessages->append_range(result.consoleMessages); + } - auto uniques = std::set(m_matches->begin(), m_matches->end(), [](const auto &l, const auto &r) { - return std::tie(l.address, l.size, l.wholeDataMatch, l.identifier, l.variable) < - std::tie(r.address, r.size, r.wholeDataMatch, r.identifier, r.variable); + 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 (auto &match : uniques) { - match.highlightId = ImHexApi::HexEditor::addBackgroundHighlight({ match.address, match.size }, YaraColor); - match.tooltipId = ImHexApi::HexEditor::addTooltip({ match. address, match.size }, hex::format("{0} [{1}]", match.identifier, match.variable), YaraColor); + 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); } }); }); diff --git a/plugins/yara_rules/source/content/yara_rule.cpp b/plugins/yara_rules/source/content/yara_rule.cpp new file mode 100644 index 000000000..cd4542862 --- /dev/null +++ b/plugins/yara_rules/source/content/yara_rule.cpp @@ -0,0 +1,197 @@ +#include + +#include +#include +#include +#include + +// 's RE type has a zero-sized array, which is not allowed in ISO C++. +#pragma GCC diagnostic push +#pragma GCC diagnostic ignored "-Wpedantic" +#include +#pragma GCC diagnostic pop + +namespace hex::plugin::yara { + + using namespace wolv::literals; + + struct ResultContext { + YaraRule *rule; + std::vector newMatches; + std::vector consoleMessages; + + std::string includeBuffer; + }; + + void YaraRule::init() { + yr_initialize(); + } + + void YaraRule::cleanup() { + yr_finalize(); + } + + YaraRule::YaraRule(const std::string &content) : m_content(content) { } + + YaraRule::YaraRule(const std::fs::path &path) : m_filePath(path) { + wolv::io::File file(path, wolv::io::File::Mode::Read); + if (!file.isValid()) + return; + + m_content = file.readString(); + } + + static int scanFunction(YR_SCAN_CONTEXT *context, int message, void *data, void *userData) { + auto &results = *static_cast(userData); + + switch (message) { + case CALLBACK_MSG_RULE_MATCHING: { + auto rule = static_cast(data); + + if (rule->strings != nullptr) { + YR_STRING *string = nullptr; + YR_MATCH *match = nullptr; + yr_rule_strings_foreach(rule, string) { + 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 }); + } + } + } else { + results.newMatches.push_back({ rule->identifier, "", Region::Invalid(), true }); + } + + break; + } + case CALLBACK_MSG_CONSOLE_LOG: { + results.consoleMessages.emplace_back(static_cast(data)); + break; + } + default: + break; + } + + return results.rule->isInterrupted() ? CALLBACK_ABORT : CALLBACK_CONTINUE; + } + + wolv::util::Expected YaraRule::match(prv::Provider *provider, u64 address, size_t size) { + YR_COMPILER *compiler = nullptr; + yr_compiler_create(&compiler); + ON_SCOPE_EXIT { + yr_compiler_destroy(compiler); + }; + + m_interrupted = false; + + ResultContext resultContext = {}; + resultContext.rule = this; + + yr_compiler_set_include_callback( + compiler, + [](const char *includeName, const char *, const char *, void *userData) -> const char * { + auto context = static_cast(userData); + + wolv::io::File file(context->rule->m_filePath.parent_path() / includeName, wolv::io::File::Mode::Read); + if (!file.isValid()) + return nullptr; + + context->includeBuffer = file.readString(); + return context->includeBuffer.c_str(); + }, + [](const char *ptr, void *userData) { + hex::unused(ptr, userData); + }, + &resultContext + ); + + if (yr_compiler_add_bytes(compiler, m_content.c_str(), m_content.size(), nullptr) != ERROR_SUCCESS) { + std::string errorMessage(0xFFFF, '\x00'); + yr_compiler_get_error_message(compiler, errorMessage.data(), errorMessage.size()); + + return wolv::util::Unexpected(Error { Error::Type::CompileError, errorMessage.c_str() }); + } + + YR_RULES *yaraRules; + yr_compiler_get_rules(compiler, &yaraRules); + ON_SCOPE_EXIT { yr_rules_destroy(yaraRules); }; + + YR_MEMORY_BLOCK_ITERATOR iterator; + + struct ScanContext { + prv::Provider *provider; + Region region; + std::vector buffer; + YR_MEMORY_BLOCK currBlock = {}; + }; + + ScanContext context; + + context.provider = provider; + context.region = { address, size }; + context.currBlock.base = 0; + context.currBlock.fetch_data = [](YR_MEMORY_BLOCK *block) -> const u8 * { + auto &context = *static_cast(block->context); + + context.buffer.resize(context.currBlock.size); + + if (context.buffer.empty()) + return nullptr; + + block->size = context.currBlock.size; + context.provider->read(context.provider->getBaseAddress() + context.currBlock.base, context.buffer.data(), context.buffer.size()); + + return context.buffer.data(); + }; + iterator.file_size = [](YR_MEMORY_BLOCK_ITERATOR *iterator) -> u64 { + auto &context = *static_cast(iterator->context); + + return context.region.size; + }; + + iterator.context = &context; + iterator.first = [](YR_MEMORY_BLOCK_ITERATOR *iterator) -> YR_MEMORY_BLOCK *{ + auto &context = *static_cast(iterator->context); + + context.currBlock.base = context.region.address; + context.currBlock.size = 0; + context.buffer.clear(); + iterator->last_error = ERROR_SUCCESS; + + return iterator->next(iterator); + }; + iterator.next = [](YR_MEMORY_BLOCK_ITERATOR *iterator) -> YR_MEMORY_BLOCK * { + auto &context = *static_cast(iterator->context); + + u64 address = context.currBlock.base + context.currBlock.size; + + iterator->last_error = ERROR_SUCCESS; + context.currBlock.base = address; + context.currBlock.size = std::min(context.region.size - address, 10_MiB); + context.currBlock.context = &context; + + if (context.currBlock.size == 0) return nullptr; + + return &context.currBlock; + }; + + if (yr_rules_scan_mem_blocks(yaraRules, &iterator, 0, scanFunction, &resultContext, 0) != ERROR_SUCCESS) { + std::string errorMessage(0xFFFF, '\x00'); + yr_compiler_get_error_message(compiler, errorMessage.data(), errorMessage.size()); + + return wolv::util::Unexpected(Error { Error::Type::RuntimeError, errorMessage.c_str() }); + } + + if (m_interrupted) + return wolv::util::Unexpected(Error { Error::Type::Interrupted, "" }); + + return Result { resultContext.newMatches, resultContext.consoleMessages }; + } + + void YaraRule::interrupt() { + m_interrupted = true; + } + + bool YaraRule::isInterrupted() const { + return m_interrupted; + } + +} \ No newline at end of file