ubuntu_server.py•4.91 kB
#!/usr/bin/env python3
"""
Simple Ubuntu VM Control MCP Server - An MCP server to interact with an Ubuntu VM via SSH.
"""
import os
import sys
import logging
import subprocess
import tempfile
import base64
from mcp.server.fastmcp import FastMCP
# Configure logging to stderr
logging.basicConfig(
level=logging.INFO,
format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
stream=sys.stderr
)
logger = logging.getLogger("ubuntu-server")
# Initialize MCP server
mcp = FastMCP("ubuntu")
# Configuration
UBUNTU_HOST = os.environ.get("UBUNTU_HOST", "")
UBUNTU_USER = os.environ.get("UBUNTU_USER", "user")
SSH_KEY_CONTENT = os.environ.get("UBUNTU_SSH_KEY", "")
# Handle SSH Key
SSH_KEY_PATH = None
if SSH_KEY_CONTENT:
with tempfile.NamedTemporaryFile(delete=False, mode='w') as temp_key_file:
temp_key_file.write(SSH_KEY_CONTENT)
SSH_KEY_PATH = temp_key_file.name
os.chmod(SSH_KEY_PATH, 0o600)
def run_ssh_command(command):
"""Helper function to run a command on the remote VM."""
if not UBUNTU_HOST:
return "❌ Error: UBUNTU_HOST environment variable not set."
if not SSH_KEY_PATH:
return "❌ Error: UBUNTU_SSH_KEY secret not set or empty."
ssh_command = [
"ssh",
"-i", SSH_KEY_PATH,
"-o", "StrictHostKeyChecking=no",
"-o", "UserKnownHostsFile=/dev/null",
f"{UBUNTU_USER}@{UBUNTU_HOST}",
command
]
try:
result = subprocess.run(
ssh_command,
capture_output=True,
text=True,
timeout=30,
check=True
)
return result.stdout.strip()
except subprocess.TimeoutExpired:
return "❌ Error: SSH command timed out."
except subprocess.CalledProcessError as e:
return f"❌ Error executing command: {e.stderr}"
except Exception as e:
return f"❌ Error: {str(e)}"
@mcp.tool()
async def execute_command(command: str = "") -> str:
"""Executes a shell command on the Ubuntu VM."""
logger.info(f"Executing command: {command}")
if not command:
return "❌ Error: Command cannot be empty."
output = run_ssh_command(command)
return f"✅ Output:\n{output}"
@mcp.tool()
async def type_text(text: str = "") -> str:
"""Types the given text on the Ubuntu VM."""
logger.info(f"Typing text: {text}")
if not text:
return "❌ Error: Text cannot be empty."
# It's safer to quote the text to handle special characters
remote_command = f"xdotool type --clearmodifiers -- '{text}'"
output = run_ssh_command(remote_command)
if "Error" in output:
return output
return "✅ Text typed successfully."
@mcp.tool()
async def click(x: str = "", y: str = "", button: str = "1") -> str:
"""Simulates a mouse click at the given coordinates (x, y)."""
logger.info(f"Clicking at ({x}, {y}) with button {button}")
if not x or not y:
return "❌ Error: X and Y coordinates are required."
remote_command = f"xdotool mousemove {x} {y} click {button}"
output = run_ssh_command(remote_command)
if "Error" in output:
return output
return f"✅ Clicked at ({x}, {y})."
@mcp.tool()
async def get_active_window_title() -> str:
"""Gets the title of the currently active window."""
logger.info("Getting active window title.")
remote_command = "xdotool getactivewindow getwindowname"
title = run_ssh_command(remote_command)
if "Error" in title:
return title
return f"✅ Active window title: {title}"
@mcp.tool()
async def take_screenshot() -> str:
"""Takes a screenshot of the entire screen and returns it as a base64 encoded string."""
logger.info("Taking screenshot.")
# Taking a screenshot, saving it to a temp file, encoding it, and cleaning up.
remote_command = "mktemp --suffix=.png"
tmp_file = run_ssh_command(remote_command)
if "Error" in tmp_file:
return f"❌ Error creating temp file on remote: {tmp_file}"
screenshot_command = f"scrot -o '{tmp_file}' && base64 '{tmp_file}' && rm '{tmp_file}'"
base64_image = run_ssh_command(screenshot_command)
if "Error" in base64_image:
return f"❌ Error taking screenshot: {base64_image}"
if not base64_image:
return "❌ Error: Screenshot command returned no data. Is 'scrot' installed on the VM?"
return f"✅ Screenshot (base64): {base64_image}"
if __name__ == "__main__":
if not all([UBUNTU_HOST, SSH_KEY_PATH]):
logger.error("UBUNTU_HOST and UBUNTU_SSH_KEY must be set.")
sys.exit(1)
logger.info("Starting Ubuntu VM Control MCP server...")
try:
mcp.run(transport='stdio')
except Exception as e:
logger.error(f"Server error: {e}", exc_info=True)
sys.exit(1)
finally:
if SSH_KEY_PATH:
os.remove(SSH_KEY_PATH)