Metasploit MCP Server

#!/usr/bin/env python3 """ Unit tests for helper functions in MetasploitMCP. """ import pytest import sys import os import asyncio from unittest.mock import Mock, patch, AsyncMock, MagicMock from typing import Dict, Any # Add the parent directory to the path to import MetasploitMCP sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..')) # Mock the dependencies that aren't available in test environment sys.modules['uvicorn'] = Mock() sys.modules['fastapi'] = Mock() sys.modules['mcp.server.fastmcp'] = Mock() sys.modules['mcp.server.sse'] = Mock() sys.modules['pymetasploit3.msfrpc'] = Mock() sys.modules['starlette.applications'] = Mock() sys.modules['starlette.routing'] = Mock() sys.modules['mcp.server.session'] = Mock() # Create mock classes for MSF objects class MockMsfRpcClient: def __init__(self): self.modules = Mock() self.core = Mock() self.sessions = Mock() self.jobs = Mock() self.consoles = Mock() class MockMsfConsole: def __init__(self, cid='test-console-id'): self.cid = cid def read(self): return {'data': 'test output', 'prompt': 'msf6 > ', 'busy': False} def write(self, command): return True class MockMsfRpcError(Exception): pass # Patch the MSF modules sys.modules['pymetasploit3.msfrpc'].MsfRpcClient = MockMsfRpcClient sys.modules['pymetasploit3.msfrpc'].MsfConsole = MockMsfConsole sys.modules['pymetasploit3.msfrpc'].MsfRpcError = MockMsfRpcError # Import after mocking from MetasploitMCP import ( _get_module_object, _set_module_options, initialize_msf_client, get_msf_client, get_msf_console, run_command_safely, find_available_port ) class TestMsfClientFunctions: """Test MSF client initialization and management functions.""" @patch('MetasploitMCP.MSF_PASSWORD', 'test-password') @patch('MetasploitMCP.MSF_SERVER', '127.0.0.1') @patch('MetasploitMCP.MSF_PORT_STR', '55553') @patch('MetasploitMCP.MSF_SSL_STR', 'false') def test_initialize_msf_client_success(self): """Test successful MSF client initialization.""" with patch('MetasploitMCP._msf_client_instance', None): with patch('MetasploitMCP.MsfRpcClient') as mock_client_class: mock_client = Mock() mock_client.core.version = {'version': '6.3.0'} mock_client_class.return_value = mock_client result = initialize_msf_client() assert result is mock_client mock_client_class.assert_called_once_with( password='test-password', server='127.0.0.1', port=55553, ssl=False ) @patch('MetasploitMCP.MSF_PORT_STR', 'invalid-port') def test_initialize_msf_client_invalid_port(self): """Test MSF client initialization with invalid port.""" with patch('MetasploitMCP._msf_client_instance', None): with pytest.raises(ValueError, match="Invalid MSF connection parameters"): initialize_msf_client() def test_get_msf_client_not_initialized(self): """Test get_msf_client when client not initialized.""" with patch('MetasploitMCP._msf_client_instance', None): with pytest.raises(ConnectionError, match="not been initialized"): get_msf_client() def test_get_msf_client_initialized(self): """Test get_msf_client when client is initialized.""" mock_client = Mock() with patch('MetasploitMCP._msf_client_instance', mock_client): result = get_msf_client() assert result is mock_client class TestGetModuleObject: """Test the _get_module_object helper function.""" @pytest.fixture def mock_client(self): """Fixture providing a mock MSF client.""" client = Mock() with patch('MetasploitMCP.get_msf_client', return_value=client): yield client @pytest.mark.asyncio async def test_get_module_object_success(self, mock_client): """Test successful module object retrieval.""" mock_module = Mock() mock_client.modules.use.return_value = mock_module result = await _get_module_object('exploit', 'windows/smb/ms17_010_eternalblue') assert result is mock_module mock_client.modules.use.assert_called_once_with('exploit', 'windows/smb/ms17_010_eternalblue') @pytest.mark.asyncio async def test_get_module_object_full_path(self, mock_client): """Test module object retrieval with full path.""" mock_module = Mock() mock_client.modules.use.return_value = mock_module result = await _get_module_object('exploit', 'exploit/windows/smb/ms17_010_eternalblue') assert result is mock_module # Should strip the module type prefix mock_client.modules.use.assert_called_once_with('exploit', 'windows/smb/ms17_010_eternalblue') @pytest.mark.asyncio async def test_get_module_object_not_found(self, mock_client): """Test module object retrieval when module not found.""" mock_client.modules.use.side_effect = KeyError("Module not found") with pytest.raises(ValueError, match="not found"): await _get_module_object('exploit', 'nonexistent/module') @pytest.mark.asyncio async def test_get_module_object_msf_error(self, mock_client): """Test module object retrieval with MSF RPC error.""" mock_client.modules.use.side_effect = MockMsfRpcError("RPC Error") with pytest.raises(MockMsfRpcError, match="RPC Error"): await _get_module_object('exploit', 'test/module') class TestSetModuleOptions: """Test the _set_module_options helper function.""" @pytest.fixture def mock_module(self): """Fixture providing a mock module object.""" module = Mock() module.fullname = 'exploit/test/module' module.__setitem__ = Mock() return module @pytest.mark.asyncio async def test_set_module_options_basic(self, mock_module): """Test basic option setting.""" options = {'RHOSTS': '192.168.1.1', 'RPORT': '80'} await _set_module_options(mock_module, options) # Should be called twice, once for each option assert mock_module.__setitem__.call_count == 2 mock_module.__setitem__.assert_any_call('RHOSTS', '192.168.1.1') mock_module.__setitem__.assert_any_call('RPORT', 80) # Type conversion: '80' -> 80 @pytest.mark.asyncio async def test_set_module_options_type_conversion(self, mock_module): """Test option setting with type conversion.""" options = { 'RPORT': '80', # String number -> int 'SSL': 'true', # String boolean -> bool 'VERBOSE': 'false', # String boolean -> bool 'THREADS': '10' # String number -> int } await _set_module_options(mock_module, options) # Verify type conversions calls = mock_module.__setitem__.call_args_list call_dict = {call[0][0]: call[0][1] for call in calls} assert call_dict['RPORT'] == 80 assert call_dict['SSL'] is True assert call_dict['VERBOSE'] is False assert call_dict['THREADS'] == 10 @pytest.mark.asyncio async def test_set_module_options_error(self, mock_module): """Test option setting with error.""" mock_module.__setitem__.side_effect = KeyError("Invalid option") options = {'INVALID_OPT': 'value'} with pytest.raises(ValueError, match="Failed to set option"): await _set_module_options(mock_module, options) class TestGetMsfConsole: """Test the get_msf_console context manager.""" @pytest.fixture def mock_client(self): """Fixture providing a mock MSF client.""" client = Mock() with patch('MetasploitMCP.get_msf_client', return_value=client): yield client @pytest.mark.asyncio async def test_get_msf_console_success(self, mock_client): """Test successful console creation and cleanup.""" mock_console = MockMsfConsole('test-console-123') mock_client.consoles.console.return_value = mock_console mock_client.consoles.destroy.return_value = 'destroyed' # Mock the global client instance for cleanup with patch('MetasploitMCP._msf_client_instance', mock_client): async with get_msf_console() as console: assert console is mock_console assert console.cid == 'test-console-123' # Verify cleanup was called mock_client.consoles.destroy.assert_called_once_with('test-console-123') @pytest.mark.asyncio async def test_get_msf_console_creation_error(self, mock_client): """Test console creation error handling.""" mock_client.consoles.console.side_effect = MockMsfRpcError("Console creation failed") with pytest.raises(MockMsfRpcError, match="Console creation failed"): async with get_msf_console() as console: pass @pytest.mark.asyncio async def test_get_msf_console_cleanup_error(self, mock_client): """Test that cleanup errors don't propagate.""" mock_console = MockMsfConsole('test-console-123') mock_client.consoles.console.return_value = mock_console mock_client.consoles.destroy.side_effect = Exception("Cleanup failed") # Should not raise exception even if cleanup fails async with get_msf_console() as console: assert console is mock_console class TestRunCommandSafely: """Test the run_command_safely function.""" @pytest.fixture def mock_console(self): """Fixture providing a mock console.""" console = Mock() console.write = Mock() console.read = Mock() return console @pytest.mark.asyncio async def test_run_command_safely_basic(self, mock_console): """Test basic command execution.""" # Mock console read to return prompt immediately mock_console.read.return_value = { 'data': 'command output\n', 'prompt': '\x01\x02msf6\x01\x02 \x01\x02> \x01\x02', 'busy': False } result = await run_command_safely(mock_console, 'help') mock_console.write.assert_called_once_with('help\n') assert 'command output' in result @pytest.mark.asyncio async def test_run_command_safely_invalid_console(self, mock_console): """Test command execution with invalid console.""" # Remove required methods delattr(mock_console, 'write') with pytest.raises(TypeError, match="Unsupported console object"): await run_command_safely(mock_console, 'help') @pytest.mark.asyncio async def test_run_command_safely_read_error(self, mock_console): """Test command execution with read error - should timeout gracefully.""" mock_console.read.side_effect = Exception("Read failed") # Should not raise exception, but timeout and return empty result result = await run_command_safely(mock_console, 'help') # Should return empty string after timeout assert isinstance(result, str) assert result == "" # Empty result after timeout class TestFindAvailablePort: """Test the find_available_port utility function.""" def test_find_available_port_success(self): """Test finding an available port.""" # This should succeed as it tests real socket binding port = find_available_port(8080, max_attempts=5) assert isinstance(port, int) assert 8080 <= port < 8085 @patch('socket.socket') def test_find_available_port_all_busy(self, mock_socket_class): """Test when all ports in range are busy.""" mock_socket = Mock() mock_socket_class.return_value.__enter__.return_value = mock_socket mock_socket.bind.side_effect = OSError("Port in use") # Should return the start port as fallback port = find_available_port(8080, max_attempts=3) assert port == 8080 @patch('socket.socket') def test_find_available_port_second_attempt(self, mock_socket_class): """Test finding port on second attempt.""" mock_socket = Mock() mock_socket_class.return_value.__enter__.return_value = mock_socket # First call fails, second succeeds mock_socket.bind.side_effect = [OSError("Port in use"), None] port = find_available_port(8080, max_attempts=3) assert port == 8081 if __name__ == "__main__": pytest.main([__file__, "-v"])