tailscale_foreground.py

"""Foreground app entrypoint — MCP tool server for Tailscale."""

from __future__ import annotations

import atexit
import asyncio
import logging
from typing import Any

import httpx
from app_runtime.mcp import create_mcp_server, run_mcp_server

from client import CLIError, TailscaleCLI, TailscaleAPI, ensure_daemon
from config import TAILSCALE_AUTH_KEY, TAILSCALE_HOSTNAME

logger = logging.getLogger("tailscale.foreground")
logger.setLevel(logging.INFO)

mcp = create_mcp_server("tailscale")

_cli: TailscaleCLI | None = None
_api: TailscaleAPI | None = None
_daemon_started: bool = False


def _error(message: str, **extra: Any) -> dict[str, Any]:
    payload: dict[str, Any] = {"status": "error", "message": message}
    payload.update(extra)
    return payload


def _success(message: str, **extra: Any) -> dict[str, Any]:
    payload: dict[str, Any] = {"status": "success", "message": message}
    payload.update(extra)
    return payload


async def _get_cli() -> TailscaleCLI:
    global _cli, _daemon_started
    if not _daemon_started:
        await ensure_daemon()
        _daemon_started = True
    if _cli is None:
        _cli = TailscaleCLI()
    return _cli


def _get_api() -> TailscaleAPI:
    global _api
    if _api is None:
        _api = TailscaleAPI()
    return _api


def _require_api() -> TailscaleAPI:
    api = _get_api()
    if not api.available:
        raise ValueError("TAILSCALE_API_KEY is required for this operation")
    return api


# --- CLI-based tools ---


@mcp.tool(
    "get_status",
    description=(
        "Get Tailscale network status: backend state, self info, peer summary. "
        "Parameters: none. "
        "Returns: JSON with backend_state, self_node, peer_count, online_peers."
    ),
)
async def get_status() -> dict[str, Any]:
    try:
        cli = await _get_cli()
        status = await cli.status()
        self_node = status.get("Self", {})
        peer_map = status.get("Peer") or {}
        online = sum(1 for p in peer_map.values() if p.get("Online"))
        return _success(
            f"Tailscale {status.get('BackendState', 'Unknown')}",
            backend_state=status.get("BackendState"),
            self_node={
                "hostname": self_node.get("HostName"),
                "dns_name": self_node.get("DNSName"),
                "tailscale_ips": self_node.get("TailscaleIPs", []),
                "os": self_node.get("OS"),
                "online": self_node.get("Online"),
            },
            peer_count=len(peer_map),
            online_peers=online,
            version=status.get("Version"),
            magic_dns_suffix=status.get("MagicDNSSuffix"),
        )
    except CLIError as e:
        return _error(str(e))
    except Exception as e:
        return _error(str(e))


@mcp.tool(
    "get_peers",
    description=(
        "List all Tailscale peers with name, IP, OS, online/offline, last seen. "
        "Parameters: none. "
        "Returns: JSON with peers array."
    ),
)
async def get_peers() -> dict[str, Any]:
    try:
        cli = await _get_cli()
        peer_list = await cli.peers()
        peers = [
            {
                "id": p.get("ID"),
                "hostname": p.get("HostName"),
                "dns_name": p.get("DNSName"),
                "tailscale_ips": p.get("TailscaleIPs", []),
                "os": p.get("OS"),
                "online": p.get("Online"),
                "last_seen": p.get("LastSeen"),
                "exit_node": p.get("ExitNode"),
                "exit_node_option": p.get("ExitNodeOption"),
                "relay": p.get("Relay"),
                "rx_bytes": p.get("RxBytes"),
                "tx_bytes": p.get("TxBytes"),
            }
            for p in peer_list
        ]
        online = sum(1 for p in peers if p.get("online"))
        return _success(
            f"{online}/{len(peers)} peers online",
            peers=peers,
            total=len(peers),
            online_count=online,
        )
    except CLIError as e:
        return _error(str(e))
    except Exception as e:
        return _error(str(e))


@mcp.tool(
    "ping_peer",
    description=(
        "Ping a Tailscale peer by hostname or IP. "
        "Parameters: target (str, required), count (int, default 1, max 10). "
        "Returns: ping output text."
    ),
)
async def ping_peer(target: str, count: int = 1) -> dict[str, Any]:
    try:
        if count < 1 or count > 10:
            return _error("count must be between 1 and 10")
        cli = await _get_cli()
        result = await cli.ping(target, count=count)
        return _success(f"Ping to {target}", output=result)
    except CLIError as e:
        return _error(str(e))
    except Exception as e:
        return _error(str(e))


@mcp.tool(
    "get_ip",
    description=(
        "Get this device's Tailscale IPv4/IPv6 addresses. "
        "Parameters: none. "
        "Returns: JSON with ips array."
    ),
)
async def get_ip() -> dict[str, Any]:
    try:
        cli = await _get_cli()
        ips = await cli.ip()
        return _success(f"IPs: {', '.join(ips)}", ips=ips)
    except CLIError as e:
        return _error(str(e))
    except Exception as e:
        return _error(str(e))


@mcp.tool(
    "connect",
    description=(
        "Join the Tailscale network with an auth key. "
        "Parameters: auth_key (str, optional — uses configured key if omitted), "
        "hostname (str, optional). "
        "Returns: connection result."
    ),
)
async def connect(
    auth_key: str | None = None,
    hostname: str | None = None,
) -> dict[str, Any]:
    try:
        key = auth_key or TAILSCALE_AUTH_KEY
        if not key:
            return _error("No auth key provided and TAILSCALE_AUTH_KEY not set")
        cli = await _get_cli()
        result = await cli.up(key, hostname=hostname or TAILSCALE_HOSTNAME or None)
        return _success("Connected to tailnet", output=result)
    except CLIError as e:
        return _error(str(e))
    except Exception as e:
        return _error(str(e))


@mcp.tool(
    "disconnect",
    description=(
        "Disconnect from the Tailscale network. "
        "Parameters: none. "
        "Returns: disconnection result."
    ),
)
async def disconnect() -> dict[str, Any]:
    try:
        cli = await _get_cli()
        result = await cli.down()
        return _success("Disconnected from tailnet", output=result)
    except CLIError as e:
        return _error(str(e))
    except Exception as e:
        return _error(str(e))


@mcp.tool(
    "set_exit_node",
    description=(
        "Route all traffic through a Tailscale exit node peer. "
        "Parameters: node (str, required — hostname or IP of exit node). "
        "Returns: result."
    ),
)
async def set_exit_node(node: str) -> dict[str, Any]:
    try:
        cli = await _get_cli()
        result = await cli.set_exit_node(node)
        return _success(f"Exit node set to {node}", output=result)
    except CLIError as e:
        return _error(str(e))
    except Exception as e:
        return _error(str(e))


@mcp.tool(
    "clear_exit_node",
    description=(
        "Stop routing traffic through an exit node. "
        "Parameters: none. "
        "Returns: result."
    ),
)
async def clear_exit_node() -> dict[str, Any]:
    try:
        cli = await _get_cli()
        result = await cli.clear_exit_node()
        return _success("Exit node cleared", output=result)
    except CLIError as e:
        return _error(str(e))
    except Exception as e:
        return _error(str(e))


@mcp.tool(
    "get_version",
    description=(
        "Get Tailscale version information. "
        "Parameters: none. "
        "Returns: version string."
    ),
)
async def get_version() -> dict[str, Any]:
    try:
        cli = await _get_cli()
        version = await cli.version()
        return _success("Version fetched", version=version)
    except CLIError as e:
        return _error(str(e))
    except Exception as e:
        return _error(str(e))


@mcp.tool(
    "tailscale_health",
    description=(
        "Check if the Tailscale daemon is running and connected. "
        "Parameters: none. "
        "Returns: JSON with healthy bool, backend_state, version."
    ),
)
async def tailscale_health() -> dict[str, Any]:
    try:
        cli = await _get_cli()
        status = await cli.status()
        state = status.get("BackendState", "Unknown")
        version = status.get("Version", "")
        healthy = state == "Running"
        return _success(
            f"Tailscale {'healthy' if healthy else 'unhealthy'}: {state}",
            healthy=healthy,
            backend_state=state,
            version=version,
        )
    except CLIError as e:
        return _error(str(e), healthy=False)
    except Exception as e:
        return _error(str(e), healthy=False)


@mcp.tool(
    "file_send",
    description=(
        "Send a file to a Tailscale peer using Taildrop. "
        "Parameters: filepath (str, required — path to local file), "
        "target (str, required — peer hostname or IP). "
        "Returns: result of the file transfer."
    ),
)
async def file_send(filepath: str, target: str) -> dict[str, Any]:
    try:
        cli = await _get_cli()
        result = await cli.file_send(filepath, target)
        return _success(f"File sent to {target}", output=result)
    except CLIError as e:
        return _error(str(e))
    except Exception as e:
        return _error(str(e))


@mcp.tool(
    "file_receive",
    description=(
        "Receive pending files from other Tailscale peers via Taildrop. "
        "Parameters: directory (str, required — local directory to save files to), "
        "wait (bool, default false — if true, block until a file arrives). "
        "Returns: result of the file receive."
    ),
)
async def file_receive(directory: str, wait: bool = False) -> dict[str, Any]:
    try:
        cli = await _get_cli()
        result = await cli.file_get(directory, wait=wait)
        return _success(f"Files received to {directory}", output=result)
    except CLIError as e:
        return _error(str(e))
    except Exception as e:
        return _error(str(e))


# --- API-based tools (require TAILSCALE_API_KEY) ---


@mcp.tool(
    "list_devices",
    description=(
        "List all devices in the tailnet. Requires TAILSCALE_API_KEY. "
        "Parameters: none. "
        "Returns: JSON with devices array."
    ),
)
async def list_devices() -> dict[str, Any]:
    try:
        api = _require_api()
        devices = await api.list_devices()
        return _success(
            f"{len(devices)} devices in tailnet",
            devices=[
                {
                    "id": d.get("id"),
                    "name": d.get("name"),
                    "hostname": d.get("hostname"),
                    "os": d.get("os"),
                    "addresses": d.get("addresses", []),
                    "authorized": d.get("authorized"),
                    "last_seen": d.get("lastSeen"),
                    "client_version": d.get("clientVersion"),
                }
                for d in devices
            ],
            total=len(devices),
        )
    except httpx.HTTPStatusError as e:
        return _error(f"Tailscale API error: {e.response.status_code}")
    except Exception as e:
        return _error(str(e))


@mcp.tool(
    "device_action",
    description=(
        "Authorize or delete a device in the tailnet. Requires TAILSCALE_API_KEY. "
        "Parameters: device_id (str, required), action (str: authorize|delete). "
        "Returns: result of the action."
    ),
)
async def device_action(device_id: str, action: str) -> dict[str, Any]:
    try:
        if action not in {"authorize", "delete"}:
            return _error("action must be 'authorize' or 'delete'")
        api = _require_api()
        if action == "authorize":
            result = await api.authorize_device(device_id)
            return _success(f"Device {device_id} authorized", result=result)
        else:
            await api.delete_device(device_id)
            return _success(f"Device {device_id} deleted")
    except httpx.HTTPStatusError as e:
        return _error(f"Tailscale API error: {e.response.status_code}")
    except Exception as e:
        return _error(str(e))


@mcp.tool(
    "manage_routes",
    description=(
        "Get or set subnet routes for a device. Requires TAILSCALE_API_KEY. "
        "Parameters: device_id (str, required), action (str: get|set), "
        "routes (array of CIDR strings, required for set). "
        "Returns: current or updated routes."
    ),
)
async def manage_routes(
    device_id: str,
    action: str,
    routes: list[str] | None = None,
) -> dict[str, Any]:
    try:
        if action not in {"get", "set"}:
            return _error("action must be 'get' or 'set'")
        api = _require_api()
        if action == "get":
            result = await api.get_device_routes(device_id)
            return _success("Routes fetched", routes=result)
        else:
            if routes is None:
                return _error("routes parameter required for set action")
            result = await api.set_device_routes(device_id, routes)
            return _success("Routes updated", routes=result)
    except httpx.HTTPStatusError as e:
        return _error(f"Tailscale API error: {e.response.status_code}")
    except Exception as e:
        return _error(str(e))


def _cleanup() -> None:
    global _api
    if _api is None:
        return
    try:
        asyncio.run(_api.close())
    except Exception:
        pass
    _api = None


def main() -> None:
    atexit.register(_cleanup)
    run_mcp_server(mcp, logger)


if __name__ == "__main__":
    main()