client.py

"""Dual Tailscale client: CLI wrapper + REST API."""

from __future__ import annotations

import asyncio
import json
import logging
from typing import Any

import httpx

from config import TAILSCALE_API_BASE, TAILSCALE_API_KEY, TAILSCALE_TAILNET

logger = logging.getLogger("tailscale.client")


class CLIError(Exception):
    """Raised when a tailscale CLI command fails."""

    def __init__(self, message: str, returncode: int = 1) -> None:
        super().__init__(message)
        self.returncode = returncode


async def _daemon_ready() -> bool:
    """Check if tailscaled is accepting commands."""
    try:
        proc = await asyncio.create_subprocess_exec(
            "tailscale", "status", "--json",
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE,
        )
        stdout, stderr = await asyncio.wait_for(proc.communicate(), timeout=5.0)
        # returncode 0 = daemon is up (connected or not)
        # returncode 1 with JSON = daemon is up but not logged in (that's fine)
        if proc.returncode == 0:
            return True
        if stdout and stdout.strip().startswith(b"{"):
            return True
        return False
    except Exception:
        return False


async def ensure_daemon() -> None:
    """Start tailscaled if it's not already running."""
    if await _daemon_ready():
        logger.info("tailscaled already running")
        return

    import os
    state_dir = "/var/lib/tailscale"
    sock_dir = "/var/run/tailscale"
    for d in (state_dir, sock_dir):
        os.makedirs(d, exist_ok=True)

    logger.info("Starting tailscaled...")
    proc = await asyncio.create_subprocess_exec(
        "tailscaled",
        f"--statedir={state_dir}",
        f"--socket={sock_dir}/tailscaled.sock",
        "--tun=userspace-networking",
        stdout=asyncio.subprocess.PIPE,
        stderr=asyncio.subprocess.PIPE,
    )

    # Wait for daemon to be ready
    for i in range(15):
        await asyncio.sleep(1)
        if proc.returncode is not None:
            # Process exited — read stderr for diagnostics
            _, stderr = await proc.communicate()
            err = stderr.decode().strip() if stderr else "unknown error"
            logger.error("tailscaled exited with code %s: %s", proc.returncode, err)
            raise CLIError(f"tailscaled exited: {err}", returncode=proc.returncode or 1)
        if await _daemon_ready():
            logger.info("tailscaled is ready")
            return

    logger.error("tailscaled did not become ready in 15 seconds")
    raise CLIError("tailscaled failed to start within 15 seconds")


class TailscaleCLI:
    """Wraps the tailscale CLI via asyncio.create_subprocess_exec."""

    async def _run(self, *args: str, timeout: float = 30.0) -> str:
        proc = await asyncio.create_subprocess_exec(
            "tailscale",
            *args,
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE,
        )
        try:
            stdout, stderr = await asyncio.wait_for(
                proc.communicate(), timeout=timeout
            )
        except asyncio.TimeoutError:
            proc.kill()
            await proc.communicate()
            raise CLIError(f"tailscale {args[0]} timed out after {timeout}s")

        out = stdout.decode().strip() if stdout else ""
        err = stderr.decode().strip() if stderr else ""

        if proc.returncode != 0:
            raise CLIError(
                err or f"tailscale {args[0]} failed (exit {proc.returncode})",
                returncode=proc.returncode or 1,
            )
        return out

    async def _run_json(self, *args: str, timeout: float = 30.0) -> dict[str, Any]:
        raw = await self._run(*args, timeout=timeout)
        return json.loads(raw)

    async def status(self) -> dict[str, Any]:
        return await self._run_json("status", "--json")

    async def ping(self, target: str, count: int = 1) -> str:
        return await self._run(
            "ping", "--c", str(count), target, timeout=count * 10 + 10
        )

    async def up(self, auth_key: str, hostname: str | None = None) -> str:
        args = ["up", f"--authkey={auth_key}"]
        if hostname:
            args.append(f"--hostname={hostname}")
        return await self._run(*args, timeout=60.0)

    async def down(self) -> str:
        return await self._run("down")

    async def set_exit_node(self, node: str) -> str:
        return await self._run("set", f"--exit-node={node}")

    async def clear_exit_node(self) -> str:
        return await self._run("set", "--exit-node=")

    async def ip(self) -> list[str]:
        raw = await self._run("ip")
        return [line.strip() for line in raw.splitlines() if line.strip()]

    async def version(self) -> str:
        return await self._run("version")

    async def file_send(self, filepath: str, target: str, timeout: float = 120.0) -> str:
        """Send a file to a peer: tailscale file cp <filepath> <target>:"""
        return await self._run("file", "cp", filepath, f"{target}:", timeout=timeout)

    async def file_get(self, directory: str, wait: bool = False, timeout: float = 120.0) -> str:
        """Receive waiting files: tailscale file get [--wait] <directory>"""
        args = ["file", "get"]
        if wait:
            args.append("--wait")
        args.append(directory)
        return await self._run(*args, timeout=timeout)

    # Convenience methods parsed from status()

    async def self_node(self) -> dict[str, Any]:
        s = await self.status()
        return s.get("Self", {})

    async def peers(self) -> list[dict[str, Any]]:
        s = await self.status()
        peer_map = s.get("Peer") or {}
        return list(peer_map.values())

    async def is_connected(self) -> bool:
        try:
            s = await self.status()
            return s.get("BackendState") == "Running"
        except Exception:
            return False


class TailscaleAPI:
    """Async REST API client for https://api.tailscale.com/api/v2."""

    def __init__(
        self,
        api_key: str | None = None,
        tailnet: str | None = None,
        base_url: str | None = None,
        timeout: float = 30.0,
    ) -> None:
        self._api_key = api_key or TAILSCALE_API_KEY
        self._tailnet = tailnet or TAILSCALE_TAILNET
        self._base_url = (base_url or TAILSCALE_API_BASE).rstrip("/")
        self._http = httpx.AsyncClient(timeout=timeout)

    @property
    def available(self) -> bool:
        return bool(self._api_key)

    async def _request(
        self,
        method: str,
        path: str,
        *,
        params: dict[str, Any] | None = None,
        json_body: dict[str, Any] | None = None,
    ) -> dict[str, Any]:
        if not self._api_key:
            raise ValueError("TAILSCALE_API_KEY is required for API operations")

        url = f"{self._base_url}{path}"
        headers = {"Authorization": f"Bearer {self._api_key}"}

        response = await self._http.request(
            method=method.upper(),
            url=url,
            params=params,
            json=json_body,
            headers=headers,
        )
        response.raise_for_status()
        if response.status_code == 204:
            return {}
        return response.json()

    async def close(self) -> None:
        try:
            await self._http.aclose()
        except Exception:
            pass

    async def list_devices(self) -> list[dict[str, Any]]:
        data = await self._request("GET", f"/tailnet/{self._tailnet}/devices")
        return data.get("devices", [])

    async def get_device(self, device_id: str) -> dict[str, Any]:
        return await self._request("GET", f"/device/{device_id}")

    async def authorize_device(self, device_id: str) -> dict[str, Any]:
        return await self._request(
            "POST",
            f"/device/{device_id}/authorized",
            json_body={"authorized": True},
        )

    async def delete_device(self, device_id: str) -> bool:
        await self._request("DELETE", f"/device/{device_id}")
        return True

    async def get_device_routes(self, device_id: str) -> dict[str, Any]:
        return await self._request("GET", f"/device/{device_id}/routes")

    async def set_device_routes(
        self, device_id: str, routes: list[str]
    ) -> dict[str, Any]:
        return await self._request(
            "POST",
            f"/device/{device_id}/routes",
            json_body={"routes": routes},
        )