Cool_Tools/ImHex
1
0
Fork 0
mirror of synced 2025-01-18 00:56:49 +01:00
Code Releases Wiki Activity

impr: Refactor Yara view

Browse Source
This commit is contained in:
WerWolv 2024-02-21 22:08:26 +01:00
parent 56e7c15064
commit 4bd24a4ffe
5 changed files with 315 additions and 198 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
plugins/yara_rules/CMakeLists.txt
View File

@ -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

15
plugins/yara_rules/include/content/views/view_yara.hpp
View File

@ -6,6 +6,8 @@
#include <hex/api/task_manager.hpp>
#include <content/yara_rule.hpp>
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<std::vector<std::pair<std::fs::path, std::fs::path>>> m_rules;
PerProvider<std::vector<std::pair<std::fs::path, std::fs::path>>> m_rulePaths;
PerProvider<std::vector<YaraMatch>> m_matches;
PerProvider<std::vector<YaraMatch*>> m_sortedMatches;
PerProvider<std::vector<std::string>> m_consoleMessages;
PerProvider<u32> m_selectedRule;
u32 m_selectedRule = 0;
TaskHolder m_matcherTask;
std::vector<std::string> m_consoleMessages;
void applyRules();
void clearResult();
};

52
plugins/yara_rules/include/content/yara_rule.hpp Normal file
View File

@ -0,0 +1,52 @@
#pragma once
#include <hex/providers/provider.hpp>
#include <string>
#include <vector>
#include <wolv/utils/expected.hpp>
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<Match> matches;
std::vector<std::string> consoleMessages;
};
struct Error {
enum class Type {
CompileError,
RuntimeError,
Interrupted
} type;
std::string message;
};
wolv::util::Expected<Result, Error> 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<bool> m_interrupted = false;
};
}

248
plugins/yara_rules/source/content/views/view_yara.cpp
View File

@ -8,17 +8,9 @@
#include <toasts/toast_notification.hpp>
#include <popups/popup_file_chooser.hpp>
// <yara/types.h>'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 <yara.h>
#pragma GCC diagnostic pop
#include <filesystem>
#include <wolv/io/file.hpp>
#include <wolv/io/fs.hpp>
#include <wolv/utils/guards.hpp>
#include <wolv/literals.hpp>
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::string>()), std::fs::path(path.get<std::string>()));
m_rulePaths.get(provider).emplace_back(std::fs::path(name.get<std::string>()), std::fs::path(path.get<std::string>()));
}
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<hex::fs::ItemFilter>{ { "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<YaraMatch> newMatches;
std::vector<std::string> 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<YaraRule::Result> 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<const char *>(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<u8 *>(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<u8> 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<ScanContext *>(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<ScanContext *>(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<ScanContext *>(iterator->context);
u64 address = context.currBlock.base + context.currBlock.size;
iterator->last_error = ERROR_SUCCESS;
context.currBlock.base = address;
context.currBlock.size = std::min<size_t>(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<ResultContext *>(userData);
switch (message) {
case CALLBACK_MSG_RULE_MATCHING:
{
auto rule = static_cast<YR_RULE *>(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<const char *>(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);
}
});
});

197
plugins/yara_rules/source/content/yara_rule.cpp Normal file
View File

@ -0,0 +1,197 @@
#include <content/yara_rule.hpp>
#include <wolv/literals.hpp>
#include <wolv/utils/guards.hpp>
#include <wolv/utils/string.hpp>
#include <wolv/io/file.hpp>
// <yara/types.h>'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 <yara.h>
#pragma GCC diagnostic pop
namespace hex::plugin::yara {
using namespace wolv::literals;
struct ResultContext {
YaraRule *rule;
std::vector<YaraRule::Match> newMatches;
std::vector<std::string> 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<ResultContext *>(userData);
switch (message) {
case CALLBACK_MSG_RULE_MATCHING: {
auto rule = static_cast<YR_RULE *>(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<const char *>(data));
break;
}
default:
break;
}
return results.rule->isInterrupted() ? CALLBACK_ABORT : CALLBACK_CONTINUE;
}
wolv::util::Expected<YaraRule::Result, YaraRule::Error> 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<ResultContext *>(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<u8> 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<ScanContext *>(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<ScanContext *>(iterator->context);
return context.region.size;
};
iterator.context = &context;
iterator.first = [](YR_MEMORY_BLOCK_ITERATOR *iterator) -> YR_MEMORY_BLOCK *{
auto &context = *static_cast<ScanContext *>(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<ScanContext *>(iterator->context);
u64 address = context.currBlock.base + context.currBlock.size;
iterator->last_error = ERROR_SUCCESS;
context.currBlock.base = address;
context.currBlock.size = std::min<size_t>(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;
}
}