"""Generic SSH MCP server.

Read-only by default. Mutating commands are blocked unless SSH_ALLOW_WRITES=true.
Targets can be set with ssh_connect and reused, or passed per-call.

Conventions mirror the other device MCPs in this stack (saved active target,
*_ALLOW_WRITES gate for writes).
"""
from __future__ import annotations

import os
import re
import shlex

import paramiko
from mcp.server.fastmcp import FastMCP

mcp = FastMCP("ssh-generic")

ALLOW_WRITES = os.environ.get("SSH_ALLOW_WRITES", "false").lower() == "true"

# Heuristic denylist for mutating operations when writes are not allowed.
_WRITE_TOKENS = (
    "rm ", "mv ", "dd ", "mkfs", "systemctl start", "systemctl stop", "systemctl restart",
    "systemctl disable", "systemctl enable", "iptables", "nft ", "ip route add",
    "ip route del", "ip addr add", "ip link set", "apt", "yum", "dnf", "tee ",
    "chmod", "chown", "passwd", "useradd", "userdel", "mount", "umount",
)
# Whole-word matches only (avoid false positives e.g. "halt" inside "head").
_WRITE_WORDS = re.compile(r"\b(reboot|shutdown|poweroff|halt|kill|pkill)\b")
# Shell output redirect to a file (allow stderr redirects like 2>/dev/null).
_WRITE_REDIRECT = re.compile(r"(?<![0-9])>(?!>)|>>")

# Active target: {"host","port","username","password"}
_active: dict | None = None


def _resolve(host, port, username, password) -> dict:
    if host:
        return {
            "host": host,
            "port": port or 22,
            "username": username or "root",
            "password": password,
        }
    if _active:
        return _active
    raise RuntimeError("No SSH target. Call ssh_connect first or pass host/username.")


def _client(target: dict) -> paramiko.SSHClient:
    client = paramiko.SSHClient()
    client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
    client.connect(
        hostname=target["host"],
        port=int(target.get("port", 22)),
        username=target.get("username"),
        password=target.get("password"),
        timeout=15,
        allow_agent=False,
        look_for_keys=bool(not target.get("password")),
    )
    return client


def _is_write(command: str) -> bool:
    low = command.lower()
    if _WRITE_WORDS.search(low):
        return True
    if _WRITE_REDIRECT.search(command):
        return True
    return any(tok in low for tok in _WRITE_TOKENS)


@mcp.tool()
def ssh_connect(host: str, username: str = "root", password: str | None = None, port: int = 22) -> str:
    """Set the active SSH target and verify connectivity."""
    global _active
    target = {"host": host, "port": port, "username": username, "password": password}
    try:
        c = _client(target)
        c.close()
    except Exception as exc:  # noqa: BLE001
        return f"connect failed: {exc}"
    _active = target
    return f"connected to {username}@{host}:{port}"


@mcp.tool()
def ssh_test_connection(host: str | None = None, username: str | None = None,
                        password: str | None = None, port: int = 22) -> str:
    """Test SSH connectivity to a host."""
    target = _resolve(host, port, username, password)
    try:
        c = _client(target)
        c.close()
        return "ok"
    except Exception as exc:  # noqa: BLE001
        return f"failed: {exc}"


@mcp.tool()
def ssh_run(command: str, host: str | None = None, username: str | None = None,
            password: str | None = None, port: int = 22, timeout: int = 30) -> str:
    """Run a shell command over SSH. Mutating commands require SSH_ALLOW_WRITES=true."""
    if _is_write(command) and not ALLOW_WRITES:
        return (
            "blocked: this looks like a mutating command and SSH_ALLOW_WRITES is false. "
            "Use the approval flow / enable writes to proceed."
        )
    target = _resolve(host, port, username, password)
    client = _client(target)
    try:
        _, stdout, stderr = client.exec_command(command, timeout=timeout)
        out = stdout.read().decode(errors="replace")
        err = stderr.read().decode(errors="replace")
        code = stdout.channel.recv_exit_status()
        result = out
        if err:
            result += f"\n[stderr]\n{err}"
        return f"[exit {code}]\n{result}".strip()
    finally:
        client.close()


@mcp.tool()
def ssh_service_status(service: str, host: str | None = None, username: str | None = None,
                       password: str | None = None, port: int = 22) -> str:
    """Read-only systemd service status."""
    return ssh_run(f"systemctl status {shlex.quote(service)} --no-pager",
                   host=host, username=username, password=password, port=port)


if __name__ == "__main__":
    mcp.run()