ssh-mcp-server/implementation/ssh-mcp-server.py

#!/usr/bin/env python3
"""
SSH-MCP Server for OBNH/OBR Infrastructure Management

This MCP server provides SSH-based tools for infrastructure access.
"""

import asyncio
import json
import subprocess
from typing import Any, Dict, List
from pathlib import Path

# MCP Server implementation using stdio transport
# Based on Model Context Protocol 2025 specification

class SSHMCPServer:
    """SSH-MCP Server providing infrastructure access tools"""

    def __init__(self):
        self.hosts = {
            "photon": {
                "hostname": "photon.obnh.io",
                "ip": "46.247.109.251",
                "description": "Debian server with Traefik, Gitea, Mastodon"
            },
            "proton": {
                "hostname": "proton.obr.sh",
                "ip": "72.61.83.117",
                "description": "Ubuntu server with Traefik, Gitea"
            },
            "fry": {
                "hostname": "fry.obr.sh",
                "ip": "v48682",
                "description": "Ubuntu server (photon replacement)"
            },
            "proxmox": {
                "hostname": "10.50.0.72",
                "ip": "10.50.0.72",
                "description": "Proxmox host"
            }
        }

    def list_tools(self) -> List[Dict[str, Any]]:
        """List available MCP tools"""
        return [
            {
                "name": "ssh_exec",
                "description": "Execute a command on a remote host via SSH",
                "inputSchema": {
                    "type": "object",
                    "properties": {
                        "host": {
                            "type": "string",
                            "description": "Host to connect to (photon, proton, fry, proxmox)"
                        },
                        "command": {
                            "type": "string",
                            "description": "Command to execute"
                        }
                    },
                    "required": ["host", "command"]
                }
            },
            {
                "name": "ssh_list_hosts",
                "description": "List all available hosts for SSH access",
                "inputSchema": {
                    "type": "object",
                    "properties": {}
                }
            },
            {
                "name": "ssh_get_file",
                "description": "Read a file from a remote host",
                "inputSchema": {
                    "type": "object",
                    "properties": {
                        "host": {
                            "type": "string",
                            "description": "Host to connect to"
                        },
                        "path": {
                            "type": "string",
                            "description": "Path to the file"
                        }
                    },
                    "required": ["host", "path"]
                }
            },
            {
                "name": "ssh_put_file",
                "description": "Write a file to a remote host",
                "inputSchema": {
                    "type": "object",
                    "properties": {
                        "host": {
                            "type": "string",
                            "description": "Host to connect to"
                        },
                        "path": {
                            "type": "string",
                            "description": "Path where to write the file"
                        },
                        "content": {
                            "type": "string",
                            "description": "File content"
                        }
                    },
                    "required": ["host", "path", "content"]
                }
            },
            {
                "name": "ssh_docker_ps",
                "description": "List Docker containers on a remote host",
                "inputSchema": {
                    "type": "object",
                    "properties": {
                        "host": {
                            "type": "string",
                            "description": "Host to connect to"
                        }
                    },
                    "required": ["host"]
                }
            },
            {
                "name": "lxc_list",
                "description": "List LXC containers on Proxmox host",
                "inputSchema": {
                    "type": "object",
                    "properties": {
                        "proxmox_host": {
                            "type": "string",
                            "description": "Proxmox host (default: proxmox)",
                            "default": "proxmox"
                        }
                    }
                }
            },
            {
                "name": "lxc_exec",
                "description": "Execute command in an LXC container on Proxmox",
                "inputSchema": {
                    "type": "object",
                    "properties": {
                        "proxmox_host": {
                            "type": "string",
                            "description": "Proxmox host",
                            "default": "proxmox"
                        },
                        "container_id": {
                            "type": "string",
                            "description": "LXC container ID"
                        },
                        "command": {
                            "type": "string",
                            "description": "Command to execute"
                        }
                    },
                    "required": ["container_id", "command"]
                }
            }
        ]

    async def call_tool(self, tool_name: str, arguments: Dict[str, Any]) -> Dict[str, Any]:
        """Execute a tool and return results"""
        try:
            if tool_name == "ssh_list_hosts":
                return await self._ssh_list_hosts()
            elif tool_name == "ssh_exec":
                return await self._ssh_exec(arguments["host"], arguments["command"])
            elif tool_name == "ssh_get_file":
                return await self._ssh_get_file(arguments["host"], arguments["path"])
            elif tool_name == "ssh_put_file":
                return await self._ssh_put_file(arguments["host"], arguments["path"], arguments["content"])
            elif tool_name == "ssh_docker_ps":
                return await self._ssh_docker_ps(arguments["host"])
            elif tool_name == "lxc_list":
                proxmox_host = arguments.get("proxmox_host", "proxmox")
                return await self._lxc_list(proxmox_host)
            elif tool_name == "lxc_exec":
                proxmox_host = arguments.get("proxmox_host", "proxmox")
                return await self._lxc_exec(proxmox_host, arguments["container_id"], arguments["command"])
            else:
                return {
                    "isError": True,
                    "content": [{"type": "text", "text": f"Unknown tool: {tool_name}"}]
                }
        except Exception as e:
            return {
                "isError": True,
                "content": [{"type": "text", "text": f"Error executing {tool_name}: {str(e)}"}]
            }

    async def _ssh_list_hosts(self) -> Dict[str, Any]:
        """List available hosts"""
        hosts_info = "\n".join([
            f"- {name}: {info['hostname']} ({info['description']})"
            for name, info in self.hosts.items()
        ])
        return {
            "content": [{"type": "text", "text": f"Available hosts:\n{hosts_info}"}]
        }

    async def _ssh_exec(self, host: str, command: str) -> Dict[str, Any]:
        """Execute SSH command"""
        if host not in self.hosts:
            return {
                "isError": True,
                "content": [{"type": "text", "text": f"Unknown host: {host}"}]
            }

        hostname = self.hosts[host]["hostname"]
        ssh_command = ["ssh", "-o", "StrictHostKeyChecking=no", f"root@{hostname}", command]

        try:
            process = await asyncio.create_subprocess_exec(
                *ssh_command,
                stdout=asyncio.subprocess.PIPE,
                stderr=asyncio.subprocess.PIPE
            )
            stdout, stderr = await process.communicate()

            result = stdout.decode() if process.returncode == 0 else stderr.decode()
            return {
                "content": [{"type": "text", "text": result}]
            }
        except Exception as e:
            return {
                "isError": True,
                "content": [{"type": "text", "text": f"SSH error: {str(e)}"}]
            }

    async def _ssh_get_file(self, host: str, path: str) -> Dict[str, Any]:
        """Get file from remote host"""
        if host not in self.hosts:
            return {
                "isError": True,
                "content": [{"type": "text", "text": f"Unknown host: {host}"}]
            }

        hostname = self.hosts[host]["hostname"]
        result = await self._ssh_exec(host, f"cat {path}")
        return result

    async def _ssh_put_file(self, host: str, path: str, content: str) -> Dict[str, Any]:
        """Write file to remote host"""
        if host not in self.hosts:
            return {
                "isError": True,
                "content": [{"type": "text", "text": f"Unknown host: {host}"}]
            }

        hostname = self.hosts[host]["hostname"]
        # Use heredoc to write file
        command = f"cat > {path} << 'EOFMCP'\n{content}\nEOFMCP"
        result = await self._ssh_exec(host, command)
        return {
            "content": [{"type": "text", "text": f"File written to {path}"}]
        }

    async def _ssh_docker_ps(self, host: str) -> Dict[str, Any]:
        """List Docker containers"""
        result = await self._ssh_exec(host, "docker ps --format '{{.Names}}\\t{{.Status}}\\t{{.Image}}'")
        return result

    async def _lxc_list(self, proxmox_host: str) -> Dict[str, Any]:
        """List LXC containers"""
        if proxmox_host not in self.hosts:
            return {
                "isError": True,
                "content": [{"type": "text", "text": f"Unknown Proxmox host: {proxmox_host}"}]
            }

        result = await self._ssh_exec(proxmox_host, "pct list")
        return result

    async def _lxc_exec(self, proxmox_host: str, container_id: str, command: str) -> Dict[str, Any]:
        """Execute command in LXC container"""
        if proxmox_host not in self.hosts:
            return {
                "isError": True,
                "content": [{"type": "text", "text": f"Unknown Proxmox host: {proxmox_host}"}]
            }

        lxc_command = f"pct exec {container_id} -- {command}"
        result = await self._ssh_exec(proxmox_host, lxc_command)
        return result


async def main():
    """Run MCP server using stdio transport"""
    server = SSHMCPServer()

    # Read JSON-RPC messages from stdin, write to stdout
    # This follows MCP stdio transport protocol
    # Note: Do not print anything to stdout except JSON-RPC messages

    while True:
        try:
            line = input()
            if not line:
                continue

            request = json.loads(line)

            if request.get("method") == "tools/list":
                response = {
                    "jsonrpc": "2.0",
                    "id": request.get("id"),
                    "result": {
                        "tools": server.list_tools()
                    }
                }
            elif request.get("method") == "tools/call":
                params = request.get("params", {})
                tool_name = params.get("name")
                arguments = params.get("arguments", {})

                result = await server.call_tool(tool_name, arguments)
                response = {
                    "jsonrpc": "2.0",
                    "id": request.get("id"),
                    "result": result
                }
            else:
                response = {
                    "jsonrpc": "2.0",
                    "id": request.get("id"),
                    "error": {
                        "code": -32601,
                        "message": f"Method not found: {request.get('method')}"
                    }
                }

            print(json.dumps(response), flush=True)

        except EOFError:
            break
        except Exception as e:
            error_response = {
                "jsonrpc": "2.0",
                "id": None,
                "error": {
                    "code": -32603,
                    "message": f"Internal error: {str(e)}"
                }
            }
            print(json.dumps(error_response), flush=True)


if __name__ == "__main__":
    asyncio.run(main())