From ba7e789a80713dbb5caa536fbe521b752e808158 Mon Sep 17 00:00:00 2001 From: WerWolv Date: Fri, 26 Dec 2025 22:32:48 +0100 Subject: [PATCH] feat: Add support for executing patterns using MCP --- lib/libimhex/source/mcp/client.cpp | 9 ++- .../romfs/mcp/tools/execute_pattern_code.json | 29 +++++++ .../tools/get_pattern_console_content.json | 24 ++++++ .../builtin/romfs/mcp/tools/get_patterns.json | 24 ++++++ .../builtin/romfs/mcp/tools/open_file.json | 30 +++++++- plugins/builtin/source/content/mcp_tools.cpp | 23 ++++-- .../content/views/view_pattern_editor.cpp | 77 +++++++++++++++++++ 7 files changed, 206 insertions(+), 10 deletions(-) create mode 100644 plugins/builtin/romfs/mcp/tools/execute_pattern_code.json create mode 100644 plugins/builtin/romfs/mcp/tools/get_pattern_console_content.json create mode 100644 plugins/builtin/romfs/mcp/tools/get_patterns.json diff --git a/lib/libimhex/source/mcp/client.cpp b/lib/libimhex/source/mcp/client.cpp index e09b3830f..46f85cf98 100644 --- a/lib/libimhex/source/mcp/client.cpp +++ b/lib/libimhex/source/mcp/client.cpp @@ -10,16 +10,17 @@ #include #include #include +#include #include namespace hex::mcp { + using namespace wolv::literals; + int Client::run(std::istream &input, std::ostream &output) { wolv::net::SocketClient client(wolv::net::SocketClient::Type::TCP, true); client.connect("127.0.0.1", Server::McpInternalPort); - fprintf(stderr, "Established connection to main ImHex instance!\n"); - while (true) { std::string request; std::getline(input, request); @@ -32,12 +33,14 @@ namespace hex::mcp { } client.writeString(request); - auto response = client.readString(); + auto response = client.readString(4_MiB); if (!response.empty() && response.front() != 0x00) output << response << '\n'; if (!client.isConnected()) break; + + std::this_thread::sleep_for(std::chrono::milliseconds(100)); } return EXIT_SUCCESS; diff --git a/plugins/builtin/romfs/mcp/tools/execute_pattern_code.json b/plugins/builtin/romfs/mcp/tools/execute_pattern_code.json new file mode 100644 index 000000000..1e6664b3e --- /dev/null +++ b/plugins/builtin/romfs/mcp/tools/execute_pattern_code.json @@ -0,0 +1,29 @@ +{ + "name": "execute_pattern_code", + "title": "Execute Pattern Code", + "description": "Executes code written in the Pattern Language on the currently opened data source. This will parse the binary data using the defined patterns and you can then get the console output using the get_pattern_console_content tool and the content of the generated patterns using the get_patterns tool. The runtime returns a result_code in form of an integer. If that value is 0, then patterns should have been generated. If it's non-zero, an error occurred during execution and you can check the console output for more details. This command may take some time to execute depending on the complexity of the code and the size of the data source.", + "inputSchema": { + "type": "object", + "properties": { + "source_code": { + "type": "string", + "description": "Code written in the Pattern Language to be executed on the currently opened data source." + } + }, + "required": ["source_code"] + }, + "outputSchema": { + "type": "object", + "properties": { + "handle": { + "type": "number", + "description": "Handle of the data source the data was read from" + }, + "result_code": { + "type": "number", + "description": "Return code of the Pattern Language runtime after executing the provided source code. A value of 0 indicates successful execution, while a non-zero value indicates an error occurred." + } + }, + "required": ["handle", "result_code"] + } +} \ No newline at end of file diff --git a/plugins/builtin/romfs/mcp/tools/get_pattern_console_content.json b/plugins/builtin/romfs/mcp/tools/get_pattern_console_content.json new file mode 100644 index 000000000..5f1018de7 --- /dev/null +++ b/plugins/builtin/romfs/mcp/tools/get_pattern_console_content.json @@ -0,0 +1,24 @@ +{ + "name": "get_pattern_console_content", + "title": "Get Pattern Console Content", + "description": "Reads the console output generated by the Pattern Language. You can read this output after a pattern has been executed previously.", + "inputSchema": { + "type": "object", + "properties": { + } + }, + "outputSchema": { + "type": "object", + "properties": { + "handle": { + "type": "number", + "description": "Handle of the data source the data was read from" + }, + "content": { + "type": "string", + "description": "The console output generated by the Pattern Language" + } + }, + "required": ["handle", "content"] + } +} \ No newline at end of file diff --git a/plugins/builtin/romfs/mcp/tools/get_patterns.json b/plugins/builtin/romfs/mcp/tools/get_patterns.json new file mode 100644 index 000000000..ed0d90a3a --- /dev/null +++ b/plugins/builtin/romfs/mcp/tools/get_patterns.json @@ -0,0 +1,24 @@ +{ + "name": "get_patterns", + "title": "Get Patterns", + "description": "This command retrieves the patterns generated by the last executed Pattern Language code. Patterns will only be returned if the last execution was successful (i.e., the result_code was 0). If the execution failed, no patterns will be returned.", + "inputSchema": { + "type": "object", + "properties": { + } + }, + "outputSchema": { + "type": "object", + "properties": { + "handle": { + "type": "number", + "description": "Handle of the data source the data was read from" + }, + "patterns": { + "type": "object", + "description": "The patterns generated by the Pattern Language" + } + }, + "required": ["handle", "patterns"] + } +} \ No newline at end of file diff --git a/plugins/builtin/romfs/mcp/tools/open_file.json b/plugins/builtin/romfs/mcp/tools/open_file.json index b12c053e2..a052172d3 100644 --- a/plugins/builtin/romfs/mcp/tools/open_file.json +++ b/plugins/builtin/romfs/mcp/tools/open_file.json @@ -14,6 +14,34 @@ }, "outputSchema": { "type": "object", - "properties": {} + "properties": { + "name": { + "type": "string", + "description": "Name of the data source" + }, + "type": { + "type": "string", + "description": "Internal type of the data source" + }, + "size": { + "type": "number", + "description": "Size of the data source in bytes" + }, + "is_writable": { + "type": "boolean", + "description": "Whether the data source is writable" + }, + "handle": { + "type": "number", + "description": "Handle of the data source to reference it in other tools" + } + }, + "required": [ + "name", + "type", + "size", + "is_writable", + "handle" + ] } } \ No newline at end of file diff --git a/plugins/builtin/source/content/mcp_tools.cpp b/plugins/builtin/source/content/mcp_tools.cpp index f9c220a79..e224a2dcb 100644 --- a/plugins/builtin/source/content/mcp_tools.cpp +++ b/plugins/builtin/source/content/mcp_tools.cpp @@ -20,8 +20,17 @@ namespace hex::plugin::builtin { ImHexApi::Provider::openProvider(provider); } - return mcp::TextContent { - .text = "File opened" + nlohmann::json result = { + { "handle", provider->getID() }, + { "name", provider->getName() }, + { "type", provider->getTypeName().get() }, + { "size", provider->getSize() }, + { "is_writable", provider->isWritable() } + }; + + return mcp::StructuredContent { + .text = result.dump(), + .data = result }; }); @@ -59,7 +68,9 @@ namespace hex::plugin::builtin { } } - nlohmann::json result = { "selected_handle", ImHexApi::Provider::get()->getID() }; + nlohmann::json result = { + { "selected_handle", ImHexApi::Provider::get()->getID() } + }; return mcp::StructuredContent { .text = result.dump(), .data = result @@ -79,9 +90,9 @@ namespace hex::plugin::builtin { auto base64 = crypt::encode64(buffer); nlohmann::json result = { - "handle", provider->getID(), - "data", std::string(base64.begin(), base64.end()), - "data_size", buffer.size() + { "handle", provider->getID() }, + { "data", std::string(base64.begin(), base64.end()) }, + { "data_size", buffer.size() } }; return mcp::StructuredContent { .text = result.dump(), diff --git a/plugins/builtin/source/content/views/view_pattern_editor.cpp b/plugins/builtin/source/content/views/view_pattern_editor.cpp index dabed1f45..6d672d986 100644 --- a/plugins/builtin/source/content/views/view_pattern_editor.cpp +++ b/plugins/builtin/source/content/views/view_pattern_editor.cpp @@ -43,9 +43,11 @@ #include #include +#include #include #include #include +#include namespace hex::plugin::builtin { @@ -2551,6 +2553,81 @@ namespace hex::plugin::builtin { return result; }); + + ContentRegistry::MCP::registerTool(romfs::get("mcp/tools/execute_pattern_code.json").string(), [this](const nlohmann::json &data) -> nlohmann::json { + auto provider = ImHexApi::Provider::get(); + + auto sourceCode = data.at("source_code").get(); + + this->evaluatePattern(sourceCode, provider); + + // Wait until evaluation has finished + while (m_runningEvaluators > 0) { + std::this_thread::sleep_for(std::chrono::milliseconds(100)); + } + + auto lock = std::scoped_lock(ContentRegistry::PatternLanguage::getRuntimeLock()); + + auto evaluationResult = m_lastEvaluationResult.load(); + + nlohmann::json result = { + { "handle", provider->getID() }, + { "result_code", evaluationResult } + }; + return mcp::StructuredContent { + .text = result.dump(), + .data = result + }; + }); + + ContentRegistry::MCP::registerTool(romfs::get("mcp/tools/get_pattern_console_content.json").string(), [this](const nlohmann::json &data) -> nlohmann::json { + std::ignore = data; + + auto provider = ImHexApi::Provider::get(); + + // Wait until evaluation has finished + while (m_runningEvaluators > 0) { + std::this_thread::sleep_for(std::chrono::milliseconds(100)); + } + + auto lock = std::scoped_lock(ContentRegistry::PatternLanguage::getRuntimeLock()); + + auto consoleOutput = m_console.get(provider); + + nlohmann::json result = { + { "handle", provider->getID() }, + { "content", wolv::util::combineStrings(consoleOutput, "\n") } + }; + return mcp::StructuredContent { + .text = result.dump(), + .data = result + }; + }); + + ContentRegistry::MCP::registerTool(romfs::get("mcp/tools/get_patterns.json").string(), [this](const nlohmann::json &data) -> nlohmann::json { + std::ignore = data; + + auto provider = ImHexApi::Provider::get(); + + // Wait until evaluation has finished + while (m_runningEvaluators > 0) { + std::this_thread::sleep_for(std::chrono::milliseconds(100)); + } + + auto lock = std::scoped_lock(ContentRegistry::PatternLanguage::getRuntimeLock()); + + pl::gen::fmt::FormatterJson formatter; + auto formattedPatterns = formatter.format(ContentRegistry::PatternLanguage::getRuntime()); + + nlohmann::json result = { + { "handle", provider->getID() }, + { "patterns", nlohmann::json::parse(std::string(formattedPatterns.begin(), formattedPatterns.end())) } + }; + return mcp::StructuredContent { + .text = result.dump(), + .data = result + }; + }); } void ViewPatternEditor::handleFileChange(prv::Provider *provider) {