commit 82a11e104c38d16ba890abd8450dbf705b11cdd2 Author: Steve White Date: Tue Apr 8 17:45:14 2025 -0500 initial commit diff --git a/README.md b/README.md new file mode 100644 index 0000000..c9474cc --- /dev/null +++ b/README.md @@ -0,0 +1,113 @@ +# MCP SSH Server + +An Anthropic Model Context Protocol (MCP) server that provides SSH access to remote systems, enabling full access to remote virtual machines in a sandbox environment. + +## Features + +- Uses stdio for MCP communication +- Provides SSH connection to remote servers +- Enables command execution on remote systems +- Secure public key authentication +- Configurable via environment variables or MCP config + +## Installation + +```bash +pip install -e . +``` + +## Configuration + +The SSH server can be configured using environment variables or the MCP JSON configuration: + +| Environment Variable | Description | Default | +|----------------------|-------------|---------| +| `MCP_SSH_HOSTNAME` | SSH server hostname or IP address | None | +| `MCP_SSH_PORT` | SSH server port | 22 | +| `MCP_SSH_USERNAME` | SSH username | None | +| `MCP_SSH_KEY_FILENAME` | Path to SSH private key file | None | + +### Claude Desktop MCP Configuration + +Add the following to your Claude Desktop MCP configuration file: + +```json +{ + "tools": [ + { + "name": "mcpssh", + "path": "python -m mcpssh", + "environment": { + "MCP_SSH_HOSTNAME": "example.com", + "MCP_SSH_PORT": "22", + "MCP_SSH_USERNAME": "user", + "MCP_SSH_KEY_FILENAME": "/path/to/private_key" + } + } + ] +} +``` + +## Usage + +This server implements the Anthropic MCP protocol and provides the following tools: + +- `ssh_connect`: Connect to an SSH server using public key authentication (using config or explicit parameters) +- `ssh_execute`: Execute a command on the SSH server +- `ssh_disconnect`: Disconnect from the SSH server + +### Example + +```python +from mcp import ClientSession, StdioServerParameters +from mcpssh.server import SSHServerMCP + +# Start the server in a subprocess +server_params = StdioServerParameters( + command="python", + args=["-m", "mcpssh"], + env={ + "MCP_SSH_HOSTNAME": "example.com", + "MCP_SSH_PORT": "22", + "MCP_SSH_USERNAME": "user", + "MCP_SSH_KEY_FILENAME": "/path/to/private_key" + } +) + +# Use with an MCP client +with ClientSession(server_params) as client: + # Connect to SSH server + client.ssh_connect() + + # Execute a command + result = client.ssh_execute(command="ls -la") + print(result["stdout"]) + + # Disconnect + client.ssh_disconnect() +``` + +### Direct Server Usage + +```python +from mcpssh.server import SSHServerMCP + +# Initialize and run the server +server = SSHServerMCP( + hostname="example.com", + port=22, + username="user", + key_filename="/path/to/private_key" +) + +# Run the server with stdio transport +server.run(transport="stdio") +``` + +## Security Note + +This tool provides full access to a remote system. It should only be used with virtual machines in sandbox environments where security implications are well understood. + +## License + +MIT \ No newline at end of file diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..f423e1c --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,24 @@ +[build-system] +requires = ["setuptools>=42", "wheel"] +build-backend = "setuptools.build_meta" + +[project] +name = "mcpssh" +version = "0.1.0" +description = "Anthropic MCP server that provides SSH access to remote systems" +requires-python = ">=3.8" +dependencies = [ + "paramiko>=3.0.0", + "pydantic>=2.0.0", + "mcp>=1.6.0", +] + +[project.optional-dependencies] +dev = [ + "pytest>=7.0.0", + "black>=23.0.0", + "mypy>=1.0.0", +] + +[tool.setuptools] +packages = ["mcpssh"] \ No newline at end of file diff --git a/setup.py b/setup.py new file mode 100644 index 0000000..fc639a0 --- /dev/null +++ b/setup.py @@ -0,0 +1,8 @@ +"""Setup script for mcpssh.""" + +from setuptools import setup + +setup( + name="mcpssh", + package_dir={"": "src"}, +) \ No newline at end of file diff --git a/src/mcpssh.egg-info/PKG-INFO b/src/mcpssh.egg-info/PKG-INFO new file mode 100644 index 0000000..93596b6 --- /dev/null +++ b/src/mcpssh.egg-info/PKG-INFO @@ -0,0 +1,12 @@ +Metadata-Version: 2.4 +Name: mcpssh +Version: 0.1.0 +Summary: Anthropic MCP server that provides SSH access to remote systems +Requires-Python: >=3.8 +Requires-Dist: paramiko>=3.0.0 +Requires-Dist: pydantic>=2.0.0 +Requires-Dist: mcp>=1.6.0 +Provides-Extra: dev +Requires-Dist: pytest>=7.0.0; extra == "dev" +Requires-Dist: black>=23.0.0; extra == "dev" +Requires-Dist: mypy>=1.0.0; extra == "dev" diff --git a/src/mcpssh.egg-info/SOURCES.txt b/src/mcpssh.egg-info/SOURCES.txt new file mode 100644 index 0000000..275472c --- /dev/null +++ b/src/mcpssh.egg-info/SOURCES.txt @@ -0,0 +1,12 @@ +README.md +pyproject.toml +setup.py +src/mcpssh/__init__.py +src/mcpssh/__main__.py +src/mcpssh/server.py +src/mcpssh.egg-info/PKG-INFO +src/mcpssh.egg-info/SOURCES.txt +src/mcpssh.egg-info/dependency_links.txt +src/mcpssh.egg-info/requires.txt +src/mcpssh.egg-info/top_level.txt +tests/test_server.py \ No newline at end of file diff --git a/src/mcpssh.egg-info/dependency_links.txt b/src/mcpssh.egg-info/dependency_links.txt new file mode 100644 index 0000000..8b13789 --- /dev/null +++ b/src/mcpssh.egg-info/dependency_links.txt @@ -0,0 +1 @@ + diff --git a/src/mcpssh.egg-info/requires.txt b/src/mcpssh.egg-info/requires.txt new file mode 100644 index 0000000..8600b64 --- /dev/null +++ b/src/mcpssh.egg-info/requires.txt @@ -0,0 +1,8 @@ +paramiko>=3.0.0 +pydantic>=2.0.0 +mcp>=1.6.0 + +[dev] +pytest>=7.0.0 +black>=23.0.0 +mypy>=1.0.0 diff --git a/src/mcpssh.egg-info/top_level.txt b/src/mcpssh.egg-info/top_level.txt new file mode 100644 index 0000000..5c7cea7 --- /dev/null +++ b/src/mcpssh.egg-info/top_level.txt @@ -0,0 +1 @@ +mcpssh diff --git a/src/mcpssh/__init__.py b/src/mcpssh/__init__.py new file mode 100644 index 0000000..859b8b7 --- /dev/null +++ b/src/mcpssh/__init__.py @@ -0,0 +1,3 @@ +"""MCP SSH Server module.""" + +__version__ = "0.1.0" \ No newline at end of file diff --git a/src/mcpssh/__main__.py b/src/mcpssh/__main__.py new file mode 100644 index 0000000..2865c83 --- /dev/null +++ b/src/mcpssh/__main__.py @@ -0,0 +1,6 @@ +"""Main entry point for the MCP SSH server.""" + +from .server import main + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/src/mcpssh/__pycache__/__init__.cpython-312.pyc b/src/mcpssh/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000..e8ba371 Binary files /dev/null and b/src/mcpssh/__pycache__/__init__.cpython-312.pyc differ diff --git a/src/mcpssh/__pycache__/__init__.cpython-313.pyc b/src/mcpssh/__pycache__/__init__.cpython-313.pyc new file mode 100644 index 0000000..5e4ac0e Binary files /dev/null and b/src/mcpssh/__pycache__/__init__.cpython-313.pyc differ diff --git a/src/mcpssh/__pycache__/__main__.cpython-312.pyc b/src/mcpssh/__pycache__/__main__.cpython-312.pyc new file mode 100644 index 0000000..211e0e8 Binary files /dev/null and b/src/mcpssh/__pycache__/__main__.cpython-312.pyc differ diff --git a/src/mcpssh/__pycache__/__main__.cpython-313.pyc b/src/mcpssh/__pycache__/__main__.cpython-313.pyc new file mode 100644 index 0000000..ca2e466 Binary files /dev/null and b/src/mcpssh/__pycache__/__main__.cpython-313.pyc differ diff --git a/src/mcpssh/__pycache__/server.cpython-312.pyc b/src/mcpssh/__pycache__/server.cpython-312.pyc new file mode 100644 index 0000000..1576b0f Binary files /dev/null and b/src/mcpssh/__pycache__/server.cpython-312.pyc differ diff --git a/src/mcpssh/__pycache__/server.cpython-313.pyc b/src/mcpssh/__pycache__/server.cpython-313.pyc new file mode 100644 index 0000000..40952b7 Binary files /dev/null and b/src/mcpssh/__pycache__/server.cpython-313.pyc differ diff --git a/src/mcpssh/server.py b/src/mcpssh/server.py new file mode 100644 index 0000000..6e09bb3 --- /dev/null +++ b/src/mcpssh/server.py @@ -0,0 +1,222 @@ +"""MCP SSH Server implementation.""" + +import json +import logging +import os +import sys +from typing import Dict, List, Optional, Any, Iterator + +import paramiko +from pydantic import BaseModel, Field +from mcp import types +from mcp.server import FastMCP + +logging.basicConfig(level=logging.INFO) +logger = logging.getLogger(__name__) + +class SSHSession: + """SSH Session that connects to a remote server.""" + + def __init__(self, hostname: str, port: int, username: str, key_filename: str): + """Initialize SSH session. + + Args: + hostname: Remote server hostname + port: SSH port + username: SSH username + key_filename: Path to SSH key file + """ + self.hostname = hostname + self.port = port + self.username = username + self.key_filename = key_filename + self.client = paramiko.SSHClient() + self.client.set_missing_host_key_policy(paramiko.AutoAddPolicy()) + self.connected = False + + def connect(self) -> bool: + """Connect to SSH server. + + Returns: + True if connection successful, False otherwise + """ + try: + self.client.connect( + hostname=self.hostname, + port=self.port, + username=self.username, + key_filename=self.key_filename + ) + self.connected = True + return True + except Exception as e: + logger.error(f"SSH connection error: {e}") + return False + + def execute_command(self, command: str) -> Dict[str, Any]: + """Execute a command on the remote server. + + Args: + command: Command to execute + + Returns: + Dictionary with stdout, stderr, and exit_code + """ + if not self.connected: + if not self.connect(): + return {"stdout": "", "stderr": "Not connected to SSH server", "exit_code": -1} + + try: + stdin, stdout, stderr = self.client.exec_command(command) + exit_code = stdout.channel.recv_exit_status() + + return { + "stdout": stdout.read().decode("utf-8"), + "stderr": stderr.read().decode("utf-8"), + "exit_code": exit_code + } + except Exception as e: + logger.error(f"Command execution error: {e}") + return {"stdout": "", "stderr": str(e), "exit_code": -1} + + def close(self) -> None: + """Close SSH connection.""" + if self.connected: + self.client.close() + self.connected = False + + +class SSHConnectionParams(BaseModel): + """Parameters for SSH connection.""" + + hostname: Optional[str] = Field(None, description="SSH server hostname or IP address") + port: Optional[int] = Field(None, description="SSH server port") + username: Optional[str] = Field(None, description="SSH username") + key_filename: Optional[str] = Field(None, description="Path to SSH private key file") + + +class CommandParams(BaseModel): + """Parameters for executing a command.""" + + command: str = Field(..., description="Command to execute on remote server") + + +class SSHServerMCP(FastMCP): + """MCP server that provides SSH access to remote systems.""" + + def __init__(self, hostname=None, port=None, username=None, key_filename=None, server_name=None, tool_prefix=None): + """Initialize SSH server. + + Args: + hostname: SSH server hostname from environment or config + port: SSH server port from environment or config + username: SSH username from environment or config + key_filename: Path to SSH key file from environment or config + server_name: Custom name for the server instance (for multiple servers) + tool_prefix: Prefix for tool names (e.g., 'server1_' for 'server1_ssh_connect') + """ + # Get server name from environment variable or parameter + server_name = server_name or os.environ.get("MCP_SSH_SERVER_NAME", "SSH Server") + super().__init__(name=server_name) + self.ssh_session: Optional[SSHSession] = None + + # Get configuration from environment variables or passed parameters + self.default_hostname = hostname or os.environ.get("MCP_SSH_HOSTNAME") + self.default_port = port or int(os.environ.get("MCP_SSH_PORT", 22)) + self.default_username = username or os.environ.get("MCP_SSH_USERNAME") + self.default_key_filename = key_filename or os.environ.get("MCP_SSH_KEY_FILENAME") + + # Get tool prefix from parameter or environment variable + self.tool_prefix = tool_prefix or os.environ.get("MCP_SSH_TOOL_PREFIX", "") + + # Register tools with optional prefix + connect_name = f"{self.tool_prefix}ssh_connect" + execute_name = f"{self.tool_prefix}ssh_execute" + disconnect_name = f"{self.tool_prefix}ssh_disconnect" + + self.add_tool(self.ssh_connect, name=connect_name, description=f"Connect to SSH server: {server_name}") + self.add_tool(self.ssh_execute, name=execute_name, description=f"Execute a command on SSH server: {server_name}") + self.add_tool(self.ssh_disconnect, name=disconnect_name, description=f"Disconnect from SSH server: {server_name}") + + def ssh_connect(self, params: SSHConnectionParams) -> Dict[str, Any]: + """Connect to an SSH server. + + Args: + params: SSH connection parameters (ignored for security-critical fields) + + Returns: + Result of connection attempt + """ + # Always use config/environment variables for security-critical fields + hostname = self.default_hostname + username = self.default_username + key_filename = self.default_key_filename + + # Only allow port to be optionally specified in the request + port = self.default_port if self.default_port else params.port + + # Validate that we have all required parameters + if not all([hostname, username, key_filename]): + missing = [] + if not hostname: missing.append("hostname") + if not username: missing.append("username") + if not key_filename: missing.append("key_filename") + return {"success": False, "message": f"Missing required parameters: {', '.join(missing)}"} + + self.ssh_session = SSHSession( + hostname=hostname, + port=port, + username=username, + key_filename=key_filename + ) + + if self.ssh_session.connect(): + return {"success": True, "message": f"Connected to {hostname}"} + else: + return {"success": False, "message": "Failed to connect to SSH server"} + + def ssh_execute(self, params: CommandParams) -> Dict[str, Any]: + """Execute a command on the SSH server. + + Args: + params: Command parameters + + Returns: + Command execution result + """ + if not self.ssh_session or not self.ssh_session.connected: + return {"success": False, "message": "Not connected to SSH server"} + + return self.ssh_session.execute_command(params.command) + + def ssh_disconnect(self) -> Dict[str, Any]: + """Disconnect from the SSH server. + + Returns: + Disconnect result + """ + if not self.ssh_session: + return {"success": True, "message": "Not connected to SSH server"} + + self.ssh_session.close() + return {"success": True, "message": "Disconnected from SSH server"} + + +def main(): + """Run the MCP SSH server.""" + # Get custom server name and tool prefix from environment if specified + server_name = os.environ.get("MCP_SSH_SERVER_NAME") + tool_prefix = os.environ.get("MCP_SSH_TOOL_PREFIX") + + # Create server with config from environment variables + server = SSHServerMCP( + server_name=server_name, + tool_prefix=tool_prefix + ) + + # Run the server with stdio transport + server.run(transport="stdio") + + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/test_mcp.py b/test_mcp.py new file mode 100644 index 0000000..f41b745 --- /dev/null +++ b/test_mcp.py @@ -0,0 +1,85 @@ +#!/usr/bin/env python +""" +Test script that simulates an MCP client using the SSH server. +This helps verify that the server properly responds to MCP requests. +""" +import json +import os +import sys +import subprocess +import tempfile + +# Set required environment variables for the subprocess +env = os.environ.copy() +env["MCP_SSH_HOSTNAME"] = "10.0.1.232" +env["MCP_SSH_USERNAME"] = "stwhite" +env["MCP_SSH_KEY_FILENAME"] = "~/.ssh/id_ed25519" + +# Create temporary files for input/output +with tempfile.NamedTemporaryFile('w+') as input_file, tempfile.NamedTemporaryFile('w+') as output_file: + # Write an MCP request to connect to SSH server + mcp_request = { + "type": "request", + "id": "1", + "method": "ssh_connect", + "params": {} + } + input_file.write(json.dumps(mcp_request) + "\n") + input_file.flush() + + # Run the MCP server process with stdin/stdout redirected + try: + print("Starting MCP SSH server...") + process = subprocess.Popen( + [sys.executable, "-m", "mcpssh"], + env=env, + stdin=open(input_file.name, 'r'), + stdout=open(output_file.name, 'w'), + stderr=subprocess.PIPE, + text=True + ) + + # Give it a moment to process + stderr_output = process.stderr.read() + if stderr_output: + print("Error output from server:") + print(stderr_output) + + # Kill the process (we're just testing initialization) + process.terminate() + process.wait() + + # Read server's response + output_file.seek(0) + response_lines = output_file.readlines() + + if not response_lines: + print("No response received from server. This likely indicates an initialization issue.") + else: + for line in response_lines: + try: + response = json.loads(line.strip()) + print("Server response:") + print(json.dumps(response, indent=2)) + except json.JSONDecodeError: + print(f"Non-JSON response: {line.strip()}") + + except Exception as e: + print(f"Error running MCP server: {e}") + +print("\nTest complete.") +print("For Claude Desktop, use this configuration:") +print(""" +"mcpssh": { + "command": "/Volumes/SAM2/CODE/MCP/mcpssh/venv/bin/python", + "args": [ + "-m", + "mcpssh" + ], + "env": { + "MCP_SSH_HOSTNAME": "10.0.1.232", + "MCP_SSH_USERNAME": "stwhite", + "MCP_SSH_KEY_FILENAME": "~/.ssh/id_ed25519" + } +} +""") diff --git a/test_ssh.py b/test_ssh.py new file mode 100644 index 0000000..b4addd0 --- /dev/null +++ b/test_ssh.py @@ -0,0 +1,23 @@ +#!/usr/bin/env python +""" +Test script for the MCP SSH server. +""" +import os +import sys +from mcpssh.server import SSHServerMCP + +# Set required environment variables for testing +os.environ["MCP_SSH_HOSTNAME"] = "10.0.1.232" # Replace with your server +os.environ["MCP_SSH_USERNAME"] = "stwhite" # Replace with your username +os.environ["MCP_SSH_KEY_FILENAME"] = "~/.ssh/id_ed25519" # Replace with your key path +# os.environ["MCP_SSH_PORT"] = "22" # Optional, defaults to 22 + +# Create and run the server +server = SSHServerMCP() +print("Server created with configuration:") +print(f"Hostname: {server.default_hostname}") +print(f"Port: {server.default_port}") +print(f"Username: {server.default_username}") +print(f"Key filename: {server.default_key_filename}") +print("\nServer ready. In a real scenario, server.run(transport='stdio') would be called.") +print("To fully integrate with Claude, add the configuration to Claude Desktop.") diff --git a/tests/__init__.py b/tests/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/__pycache__/__init__.cpython-312.pyc b/tests/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000..401ae7c Binary files /dev/null and b/tests/__pycache__/__init__.cpython-312.pyc differ diff --git a/tests/__pycache__/__init__.cpython-313.pyc b/tests/__pycache__/__init__.cpython-313.pyc new file mode 100644 index 0000000..1637e70 Binary files /dev/null and b/tests/__pycache__/__init__.cpython-313.pyc differ diff --git a/tests/__pycache__/test_server.cpython-312-pytest-7.4.4.pyc b/tests/__pycache__/test_server.cpython-312-pytest-7.4.4.pyc new file mode 100644 index 0000000..c7c9490 Binary files /dev/null and b/tests/__pycache__/test_server.cpython-312-pytest-7.4.4.pyc differ diff --git a/tests/__pycache__/test_server.cpython-313-pytest-8.3.5.pyc b/tests/__pycache__/test_server.cpython-313-pytest-8.3.5.pyc new file mode 100644 index 0000000..ed252fb Binary files /dev/null and b/tests/__pycache__/test_server.cpython-313-pytest-8.3.5.pyc differ diff --git a/tests/test_server.py b/tests/test_server.py new file mode 100644 index 0000000..84e8476 --- /dev/null +++ b/tests/test_server.py @@ -0,0 +1,198 @@ +"""Tests for the MCP SSH server.""" + +import unittest +from unittest.mock import patch, MagicMock + +from mcpssh.server import SSHServerMCP, SSHSession, SSHConnectionParams, CommandParams + + +class TestSSHSession(unittest.TestCase): + """Test SSH session class.""" + + @patch('paramiko.SSHClient') + def test_connect_success(self, mock_ssh_client): + """Test successful SSH connection.""" + # Setup + mock_client = MagicMock() + mock_ssh_client.return_value = mock_client + + # Execute + session = SSHSession("example.com", 22, "username", "/path/to/key") + result = session.connect() + + # Verify + self.assertTrue(result) + self.assertTrue(session.connected) + mock_client.connect.assert_called_once_with( + hostname="example.com", + port=22, + username="username", + key_filename="/path/to/key" + ) + + @patch('paramiko.SSHClient') + def test_connect_failure(self, mock_ssh_client): + """Test failed SSH connection.""" + # Setup + mock_client = MagicMock() + mock_client.connect.side_effect = Exception("Connection failed") + mock_ssh_client.return_value = mock_client + + # Execute + session = SSHSession("example.com", 22, "username", "/path/to/key") + result = session.connect() + + # Verify + self.assertFalse(result) + self.assertFalse(session.connected) + + @patch('paramiko.SSHClient') + def test_execute_command_success(self, mock_ssh_client): + """Test successful command execution.""" + # Setup + mock_client = MagicMock() + mock_stdout = MagicMock() + mock_stdout.read.return_value = b"command output" + mock_stdout.channel.recv_exit_status.return_value = 0 + mock_stderr = MagicMock() + mock_stderr.read.return_value = b"" + mock_client.exec_command.return_value = (None, mock_stdout, mock_stderr) + mock_ssh_client.return_value = mock_client + + # Execute + session = SSHSession("example.com", 22, "username", "/path/to/key") + session.connected = True # Skip connection + session.client = mock_client + result = session.execute_command("ls -la") + + # Verify + self.assertEqual(result["stdout"], "command output") + self.assertEqual(result["stderr"], "") + self.assertEqual(result["exit_code"], 0) + + @patch('paramiko.SSHClient') + def test_close(self, mock_ssh_client): + """Test closing SSH connection.""" + # Setup + mock_client = MagicMock() + mock_ssh_client.return_value = mock_client + + # Execute + session = SSHSession("example.com", 22, "username", "/path/to/key") + session.connected = True + session.client = mock_client + session.close() + + # Verify + self.assertFalse(session.connected) + mock_client.close.assert_called_once() + + +class TestSSHServerMCP(unittest.TestCase): + """Test SSH server MCP implementation.""" + + def setUp(self): + """Set up test environment.""" + self.server = SSHServerMCP() + + @patch('mcpssh.server.SSHSession') + def test_ssh_connect_success(self, mock_ssh_session): + """Test successful SSH connection.""" + # Setup + mock_session = MagicMock() + mock_session.connect.return_value = True + mock_ssh_session.return_value = mock_session + + # Execute + params = { + "hostname": "example.com", + "port": 22, + "username": "username", + "key_filename": "/path/to/key" + } + result = self.server.ssh_connect(SSHConnectionParams(**params)) + + # Verify + self.assertTrue(result["success"]) + self.assertEqual(result["message"], "Connected to example.com") + + @patch('mcpssh.server.SSHSession') + def test_ssh_connect_failure(self, mock_ssh_session): + """Test failed SSH connection.""" + # Setup + mock_session = MagicMock() + mock_session.connect.return_value = False + mock_ssh_session.return_value = mock_session + + # Execute + params = { + "hostname": "example.com", + "port": 22, + "username": "username", + "key_filename": "/path/to/key" + } + result = self.server.ssh_connect(SSHConnectionParams(**params)) + + # Verify + self.assertFalse(result["success"]) + self.assertEqual(result["message"], "Failed to connect to SSH server") + + def test_ssh_execute_not_connected(self): + """Test command execution when not connected.""" + # Execute + params = {"command": "ls -la"} + result = self.server.ssh_execute(CommandParams(**params)) + + # Verify + self.assertFalse(result["success"]) + self.assertEqual(result["message"], "Not connected to SSH server") + + @patch('mcpssh.server.SSHSession') + def test_ssh_execute_success(self, mock_ssh_session): + """Test successful command execution.""" + # Setup + mock_session = MagicMock() + mock_session.connected = True + mock_session.execute_command.return_value = { + "stdout": "command output", + "stderr": "", + "exit_code": 0 + } + + self.server.ssh_session = mock_session + + # Execute + params = {"command": "ls -la"} + result = self.server.ssh_execute(CommandParams(**params)) + + # Verify + self.assertEqual(result["stdout"], "command output") + self.assertEqual(result["stderr"], "") + self.assertEqual(result["exit_code"], 0) + + def test_ssh_disconnect_not_connected(self): + """Test disconnection when not connected.""" + # Execute + result = self.server.ssh_disconnect() + + # Verify + self.assertTrue(result["success"]) + self.assertEqual(result["message"], "Not connected to SSH server") + + def test_ssh_disconnect_success(self): + """Test successful disconnection.""" + # Setup + mock_session = MagicMock() + self.server.ssh_session = mock_session + + # Execute + result = self.server.ssh_disconnect() + + # Verify + self.assertTrue(result["success"]) + self.assertEqual(result["message"], "Disconnected from SSH server") + mock_session.close.assert_called_once() + + +if __name__ == "__main__": + unittest.main() \ No newline at end of file