259 lines
16 KiB
Python
259 lines
16 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
Test script for the GhydraMCP bridge using the MCP client.
|
|
This script tests the bridge by sending MCP requests and handling responses.
|
|
"""
|
|
import json
|
|
import logging
|
|
import os
|
|
import sys
|
|
from typing import Any
|
|
|
|
import anyio
|
|
|
|
# Get host and port from environment variables or use defaults
|
|
GHYDRAMCP_TEST_HOST = os.getenv('GHYDRAMCP_TEST_HOST', 'localhost')
|
|
GHYDRAMCP_TEST_PORT = int(os.getenv('GHYDRAMCP_TEST_PORT', '8192'))
|
|
from mcp.client.session import ClientSession
|
|
from mcp.client.stdio import StdioServerParameters, stdio_client
|
|
|
|
# Set up logging
|
|
logging.basicConfig(level=logging.INFO)
|
|
logger = logging.getLogger("mcp_client_test")
|
|
|
|
async def assert_standard_mcp_success_response(response_content, expected_result_type=None):
|
|
"""Helper to assert the standard success response structure for MCP tool calls."""
|
|
assert response_content, "Response content is empty"
|
|
try:
|
|
data = json.loads(response_content[0].text)
|
|
except (json.JSONDecodeError, IndexError) as e:
|
|
assert False, f"Failed to parse JSON response: {e} - Content: {response_content}"
|
|
|
|
assert "success" in data, "Response missing 'success' field"
|
|
assert data["success"] is True, f"API call failed: {data.get('error', 'Unknown error')}"
|
|
assert "timestamp" in data, "Response missing 'timestamp' field"
|
|
assert isinstance(data["timestamp"], (int, float)), "'timestamp' should be a number"
|
|
assert "port" in data, "Response missing 'port' field"
|
|
# We don't strictly check port number here as it might vary in MCP tests
|
|
assert "result" in data, "Response missing 'result' field"
|
|
if expected_result_type:
|
|
assert isinstance(data["result"], expected_result_type), \
|
|
f"'result' field type mismatch: expected {expected_result_type}, got {type(data['result'])}"
|
|
return data # Return parsed data for further checks if needed
|
|
|
|
async def test_bridge():
|
|
"""Test the bridge using the MCP client"""
|
|
# Configure the server parameters
|
|
server_parameters = StdioServerParameters(
|
|
command=sys.executable,
|
|
args=["bridge_mcp_hydra.py"],
|
|
)
|
|
|
|
# Connect to the bridge
|
|
logger.info("Connecting to bridge...")
|
|
async with stdio_client(server_parameters) as (read_stream, write_stream):
|
|
# Create a session
|
|
logger.info("Creating session...")
|
|
async with ClientSession(read_stream, write_stream) as session:
|
|
# Initialize the session
|
|
logger.info("Initializing session...")
|
|
init_result = await session.initialize()
|
|
logger.info(f"Initialization result: {init_result}")
|
|
|
|
# List tools
|
|
logger.info("Listing tools...")
|
|
tools_result = await session.list_tools()
|
|
logger.info(f"Tools result: {tools_result}")
|
|
|
|
# Call the list_instances tool
|
|
logger.info("Calling list_instances tool...")
|
|
list_instances_result = await session.call_tool("list_instances")
|
|
logger.info(f"List instances result: {list_instances_result}")
|
|
|
|
# Call the discover_instances tool
|
|
logger.info("Calling discover_instances tool...")
|
|
discover_instances_result = await session.call_tool("discover_instances")
|
|
logger.info(f"Discover instances result: {discover_instances_result}")
|
|
|
|
# Call the list_functions tool
|
|
logger.info("Calling list_functions tool...")
|
|
list_functions_result = await session.call_tool(
|
|
"list_functions",
|
|
arguments={"port": GHYDRAMCP_TEST_PORT, "offset": 0, "limit": 5}
|
|
)
|
|
logger.info(f"List functions result: {list_functions_result}")
|
|
|
|
# Test mutating operations by changing and reverting
|
|
logger.info("Testing mutating operations...")
|
|
|
|
try:
|
|
# Get a known function to test with from list_functions result
|
|
list_funcs = await session.call_tool(
|
|
"list_functions",
|
|
arguments={"port": GHYDRAMCP_TEST_PORT, "offset": 0, "limit": 5}
|
|
)
|
|
|
|
if not list_funcs or not list_funcs.content:
|
|
logger.warning("No functions found - skipping mutating tests")
|
|
return
|
|
|
|
# Parse the JSON response from list_functions using helper
|
|
try:
|
|
list_funcs_data = await assert_standard_mcp_success_response(list_funcs.content, expected_result_type=list)
|
|
func_list = list_funcs_data.get("result", [])
|
|
if not func_list:
|
|
logger.warning("No functions in list_functions result - skipping mutating tests")
|
|
return
|
|
|
|
# Get first function's name and address
|
|
first_func = func_list[0]
|
|
func_name = first_func.get("name", "")
|
|
func_address = first_func.get("address", "")
|
|
|
|
if not func_name or not func_address:
|
|
logger.warning("No function name/address found in list_functions result - skipping mutating tests")
|
|
return
|
|
except AssertionError as e:
|
|
logger.warning(f"Error processing list_functions data: {e} - skipping mutating tests")
|
|
return
|
|
|
|
# Test function renaming
|
|
original_name = func_name
|
|
test_name = f"{func_name}_test"
|
|
|
|
# Test successful rename operations (These return simple success/message, not full result)
|
|
rename_args = {"port": GHYDRAMCP_TEST_PORT, "name": original_name, "new_name": test_name}
|
|
logger.info(f"Calling update_function with args: {rename_args}")
|
|
rename_result = await session.call_tool("update_function", arguments=rename_args)
|
|
rename_data = json.loads(rename_result.content[0].text) # Parse simple response
|
|
assert rename_data.get("success") is True, f"Rename failed: {rename_data}"
|
|
logger.info(f"Rename result: {rename_result}")
|
|
|
|
# Verify rename by getting the function
|
|
renamed_func = await session.call_tool("get_function", arguments={"port": 8192, "name": test_name})
|
|
renamed_data = await assert_standard_mcp_success_response(renamed_func.content, expected_result_type=dict)
|
|
assert renamed_data.get("result", {}).get("name") == test_name, f"Renamed function has wrong name: {renamed_data}"
|
|
logger.info(f"Renamed function result: {renamed_func}")
|
|
|
|
# Rename back to original
|
|
revert_args = {"port": GHYDRAMCP_TEST_PORT, "name": test_name, "new_name": original_name}
|
|
logger.info(f"Calling update_function with args: {revert_args}")
|
|
revert_result = await session.call_tool("update_function", arguments=revert_args)
|
|
revert_data = json.loads(revert_result.content[0].text) # Parse simple response
|
|
assert revert_data.get("success") is True, f"Revert rename failed: {revert_data}"
|
|
logger.info(f"Revert rename result: {revert_result}")
|
|
|
|
# Verify revert by getting the function
|
|
original_func = await session.call_tool("get_function", arguments={"port": GHYDRAMCP_TEST_PORT, "name": original_name})
|
|
original_data = await assert_standard_mcp_success_response(original_func.content, expected_result_type=dict)
|
|
assert original_data.get("result", {}).get("name") == original_name, f"Original function has wrong name: {original_data}"
|
|
logger.info(f"Original function result: {original_func}")
|
|
|
|
# Test get_function_by_address
|
|
logger.info(f"Calling get_function_by_address with address: {func_address}")
|
|
get_by_addr_result = await session.call_tool("get_function_by_address", arguments={"port": GHYDRAMCP_TEST_PORT, "address": func_address})
|
|
get_by_addr_data = await assert_standard_mcp_success_response(get_by_addr_result.content, expected_result_type=dict)
|
|
result_data = get_by_addr_data.get("result", {})
|
|
assert "name" in result_data, "Missing name field in get_function_by_address result"
|
|
assert "address" in result_data, "Missing address field in get_function_by_address result"
|
|
assert "signature" in result_data, "Missing signature field in get_function_by_address result"
|
|
assert "decompilation" in result_data, "Missing decompilation field in get_function_by_address result"
|
|
assert result_data.get("name") == original_name, f"Wrong name in get_function_by_address: {result_data.get('name')}"
|
|
logger.info(f"Get function by address result: {get_by_addr_result}")
|
|
|
|
# Test decompile_function_by_address
|
|
logger.info(f"Calling decompile_function_by_address with address: {func_address}")
|
|
decompile_result = await session.call_tool("decompile_function_by_address", arguments={"port": GHYDRAMCP_TEST_PORT, "address": func_address})
|
|
decompile_data = await assert_standard_mcp_success_response(decompile_result.content, expected_result_type=dict)
|
|
assert "decompilation" in decompile_data.get("result", {}), f"Decompile result missing 'decompilation': {decompile_data}"
|
|
assert isinstance(decompile_data.get("result", {}).get("decompilation", ""), str), f"Decompilation is not a string: {decompile_data}"
|
|
assert len(decompile_data.get("result", {}).get("decompilation", "")) > 0, f"Decompilation result is empty: {decompile_data}"
|
|
logger.info(f"Decompile function by address result: {decompile_result}")
|
|
|
|
# Test list_variables
|
|
logger.info("Calling list_variables tool...")
|
|
list_vars_result = await session.call_tool("list_variables", arguments={"port": 8192, "limit": 10})
|
|
list_vars_data = await assert_standard_mcp_success_response(list_vars_result.content, expected_result_type=list)
|
|
variables_list = list_vars_data.get("result", [])
|
|
if variables_list: # Only validate structure if we get results
|
|
for var in variables_list:
|
|
assert "name" in var, f"Variable missing name: {var}"
|
|
assert "type" in var, f"Variable missing type: {var}"
|
|
assert "dataType" in var, f"Variable missing dataType: {var}"
|
|
logger.info(f"List variables result: {list_vars_result}")
|
|
|
|
# Test successful comment operations (These return simple success/message)
|
|
test_comment = "Test comment from MCP client"
|
|
comment_args = {"port": 8192, "address": func_address, "comment": test_comment}
|
|
logger.info(f"Calling set_decompiler_comment with args: {comment_args}")
|
|
comment_result = await session.call_tool("set_decompiler_comment", arguments=comment_args)
|
|
comment_data = json.loads(comment_result.content[0].text)
|
|
assert comment_data.get("success") is True, f"Add comment failed: {comment_data}"
|
|
logger.info(f"Add comment result: {comment_result}")
|
|
|
|
# Remove comment
|
|
remove_comment_args = {"port": 8192, "address": func_address, "comment": ""}
|
|
logger.info(f"Calling set_decompiler_comment with args: {remove_comment_args}")
|
|
remove_comment_result = await session.call_tool("set_decompiler_comment", arguments=remove_comment_args)
|
|
remove_data = json.loads(remove_comment_result.content[0].text)
|
|
assert remove_data.get("success") is True, f"Remove comment failed: {remove_data}"
|
|
logger.info(f"Remove comment result: {remove_comment_result}")
|
|
|
|
# Test expected failure cases
|
|
# Try to rename non-existent function
|
|
bad_rename_args = {"port": 8192, "name": "nonexistent_function", "new_name": "should_fail"}
|
|
logger.info(f"Calling update_function with args: {bad_rename_args}")
|
|
bad_rename_result = await session.call_tool("update_function", arguments=bad_rename_args)
|
|
logger.info(f"Bad rename result: {bad_rename_result}") # Log the response
|
|
bad_rename_data = json.loads(bad_rename_result.content[0].text)
|
|
assert bad_rename_data.get("success") is False, f"Renaming non-existent function should fail, but got: {bad_rename_data}"
|
|
|
|
# Try to get non-existent function
|
|
bad_get_result = await session.call_tool(
|
|
"get_function",
|
|
arguments={"port": 8192, "name": "nonexistent_function"}
|
|
)
|
|
logger.info(f"Bad get result: {bad_get_result}") # Log the response
|
|
bad_get_data = json.loads(bad_get_result.content[0].text)
|
|
assert bad_get_data.get("success") is False, f"Getting non-existent function should fail, but got: {bad_get_data}"
|
|
|
|
# Try to comment on invalid address
|
|
bad_comment_args = {"port": 8192, "address": "0xinvalid", "comment": "should fail"}
|
|
logger.info(f"Calling set_decompiler_comment with args: {bad_comment_args}")
|
|
bad_comment_result = await session.call_tool("set_decompiler_comment", arguments=bad_comment_args)
|
|
bad_comment_data = json.loads(bad_comment_result.content[0].text)
|
|
assert bad_comment_data.get("success") is False, "Commenting on invalid address should fail"
|
|
|
|
# Test get_current_address
|
|
logger.info("Calling get_current_address tool...")
|
|
current_addr_result = await session.call_tool("get_current_address", arguments={"port": 8192})
|
|
current_addr_data = await assert_standard_mcp_success_response(current_addr_result.content, expected_result_type=dict)
|
|
assert "address" in current_addr_data.get("result", {}), "Missing address in get_current_address result"
|
|
assert isinstance(current_addr_data.get("result", {}).get("address", ""), str), "Address should be a string"
|
|
logger.info(f"Get current address result: {current_addr_result}")
|
|
|
|
# Test get_current_function
|
|
logger.info("Calling get_current_function tool...")
|
|
current_func_result = await session.call_tool("get_current_function", arguments={"port": 8192})
|
|
current_func_data = await assert_standard_mcp_success_response(current_func_result.content, expected_result_type=dict)
|
|
result_data = current_func_data.get("result", {})
|
|
assert "name" in result_data, "Missing name in get_current_function result"
|
|
assert "address" in result_data, "Missing address in get_current_function result"
|
|
assert "signature" in result_data, "Missing signature in get_current_function result"
|
|
logger.info(f"Get current function result: {current_func_result}")
|
|
|
|
except Exception as e:
|
|
logger.error(f"Error testing mutating operations: {e}")
|
|
raise
|
|
|
|
def main():
|
|
"""Main entry point"""
|
|
try:
|
|
anyio.run(test_bridge)
|
|
except Exception as e:
|
|
logger.error(f"Error: {e}")
|
|
sys.exit(1)
|
|
|
|
if __name__ == "__main__":
|
|
main()
|