137 lines
5.1 KiB
Python
Executable File
137 lines
5.1 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
"""
|
|
Test script for the comment functionality in GhydraMCP.
|
|
|
|
Tests both HTTP API and MCP bridge interfaces for setting and retrieving
|
|
different types of comments in Ghidra, including plate, pre, post, EOL,
|
|
repeatable, and decompiler comments.
|
|
"""
|
|
import json
|
|
import logging
|
|
import sys
|
|
import time
|
|
from urllib.parse import quote
|
|
|
|
import anyio
|
|
import requests
|
|
from mcp.client.session import ClientSession
|
|
from mcp.client.stdio import StdioServerParameters, stdio_client
|
|
|
|
logging.basicConfig(level=logging.INFO)
|
|
logger = logging.getLogger("comment_test")
|
|
|
|
def test_http_api_comments(port=8192, address="08000200"):
|
|
"""
|
|
Test setting and retrieving comments using direct HTTP API.
|
|
|
|
Args:
|
|
port: Ghidra HTTP API port
|
|
address: Memory address for comments
|
|
"""
|
|
logger.info("===== Testing HTTP API Comments =====")
|
|
base_url = f"http://localhost:{port}"
|
|
|
|
comment_types = ["plate", "pre", "post", "eol", "repeatable"]
|
|
|
|
for comment_type in comment_types:
|
|
comment_text = f"TEST {comment_type.upper()} COMMENT {int(time.time())}"
|
|
logger.info(f"Setting {comment_type} comment: {comment_text}")
|
|
|
|
url = f"{base_url}/memory/{address}/comments/{comment_type}"
|
|
payload = {"comment": comment_text}
|
|
|
|
try:
|
|
r = requests.post(url, json=payload, timeout=10)
|
|
logger.info(f"Status code: {r.status_code}")
|
|
logger.info(f"Response: {r.text}")
|
|
|
|
if r.status_code == 200:
|
|
r_get = requests.get(url, timeout=10)
|
|
logger.info(f"GET Status code: {r_get.status_code}")
|
|
logger.info(f"GET Response: {r_get.text}")
|
|
except Exception as e:
|
|
logger.error(f"Error setting {comment_type} comment: {e}")
|
|
|
|
async def test_bridge_comments():
|
|
"""
|
|
Test MCP bridge comment functionality.
|
|
|
|
Sets and clears both plate comments and decompiler comments using the
|
|
MCP bridge interface.
|
|
"""
|
|
logger.info("===== Testing MCP Bridge Comments =====")
|
|
|
|
server_parameters = StdioServerParameters(
|
|
command=sys.executable,
|
|
args=["bridge_mcp_hydra.py"],
|
|
)
|
|
|
|
logger.info("Connecting to bridge...")
|
|
async with stdio_client(server_parameters) as (read_stream, write_stream):
|
|
logger.info("Creating session...")
|
|
async with ClientSession(read_stream, write_stream) as session:
|
|
logger.info("Initializing session...")
|
|
await session.initialize()
|
|
|
|
# First set the current instance
|
|
logger.info("Setting current Ghidra instance...")
|
|
await session.call_tool(
|
|
"instances_use",
|
|
arguments={"port": 8192}
|
|
)
|
|
|
|
logger.info("Getting current address...")
|
|
addr_result = await session.call_tool("ui_get_current_address")
|
|
addr_data = json.loads(addr_result.content[0].text)
|
|
|
|
if not addr_data.get("success", False):
|
|
logger.error("Failed to get current address")
|
|
return
|
|
|
|
address = addr_data.get("result", {}).get("address", "08000200")
|
|
logger.info(f"Using address: {address}")
|
|
|
|
logger.info("Testing comments_set with plate type...")
|
|
comment_text = f"MCP PLATE COMMENT {int(time.time())}"
|
|
result = await session.call_tool("comments_set",
|
|
arguments={"address": address,
|
|
"comment": comment_text,
|
|
"comment_type": "plate"})
|
|
logger.info(f"comments_set result: {result}")
|
|
|
|
logger.info("Testing functions_set_comment...")
|
|
decompiler_comment = f"MCP DECOMPILER COMMENT {int(time.time())}"
|
|
decompile_result = await session.call_tool("functions_set_comment",
|
|
arguments={"address": address,
|
|
"comment": decompiler_comment})
|
|
logger.info(f"functions_set_comment result: {decompile_result}")
|
|
|
|
await anyio.sleep(5)
|
|
|
|
logger.info("Clearing comments...")
|
|
await session.call_tool("comments_set",
|
|
arguments={"address": address,
|
|
"comment": "",
|
|
"comment_type": "plate"})
|
|
|
|
await session.call_tool("functions_set_comment",
|
|
arguments={"address": address,
|
|
"comment": ""})
|
|
|
|
def main():
|
|
"""
|
|
Main entry point for comment tests.
|
|
|
|
Runs both HTTP API and MCP bridge tests sequentially.
|
|
"""
|
|
try:
|
|
test_http_api_comments()
|
|
anyio.run(test_bridge_comments)
|
|
logger.info("All comment tests completed successfully")
|
|
return True
|
|
except Exception as e:
|
|
logger.error(f"Error in comment tests: {e}")
|
|
sys.exit(1)
|
|
|
|
if __name__ == "__main__":
|
|
main() |