ForcePilot/backend/package/yuxi/channel/gateway/tailscale_auth.py

import asyncio
import json
import logging
import time
from dataclasses import dataclass

logger = logging.getLogger(__name__)

WHOIS_CACHE_TTL_MS = 60_000
WHOIS_ERROR_TTL_MS = 5_000
WHOIS_TIMEOUT_S = 5.0


@dataclass
class TailscaleWhoisIdentity:
    login: str
    name: str | None = None


@dataclass
class _CacheEntry:
    value: TailscaleWhoisIdentity | None
    expires_at: float


_whois_cache: dict[str, _CacheEntry] = {}


def _read_cached_whois(ip: str, now_ms: float) -> TailscaleWhoisIdentity | None | bool:
    """Returns None=cache miss, False=cached null, or the identity."""
    entry = _whois_cache.get(ip)
    if entry is None:
        return None
    if entry.expires_at <= now_ms:
        del _whois_cache[ip]
        return None
    return entry.value if entry.value is not None else False  # type: ignore[return-value]


def _write_cached_whois(ip: str, value: TailscaleWhoisIdentity | None, ttl_ms: float) -> None:
    _whois_cache[ip] = _CacheEntry(value=value, expires_at=time.time() * 1000 + ttl_ms)


def _parse_whois_identity(payload: dict) -> TailscaleWhoisIdentity | None:
    user_profile = payload.get("UserProfile") or payload.get("userProfile") or payload.get("User")
    if not isinstance(user_profile, dict):
        return None
    login = (
        user_profile.get("LoginName")
        or user_profile.get("Login")
        or user_profile.get("loginName")
        or user_profile.get("login")
    )
    if not login or not isinstance(login, str):
        return None
    name = (
        user_profile.get("DisplayName")
        or user_profile.get("displayName")
        or user_profile.get("Name")
        or user_profile.get("name")
    )
    return TailscaleWhoisIdentity(
        login=login.lower().strip(),
        name=name.strip() if isinstance(name, str) and name.strip() else None,
    )


def _parse_possibly_noisy_json(stdout: str) -> dict:
    trimmed = stdout.strip()
    start = trimmed.find("{")
    end = trimmed.rfind("}")
    if start >= 0 and end > start:
        return json.loads(trimmed[start : end + 1])
    return json.loads(trimmed)


async def _get_tailscale_binary() -> str | None:
    try:
        proc = await asyncio.create_subprocess_exec(
            "which",
            "tailscale",
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.DEVNULL,
        )
        stdout, _ = await proc.communicate()
        if proc.returncode == 0 and stdout:
            return stdout.decode().strip()
    except FileNotFoundError:
        logger.debug("tailscale binary not found via 'which'")
    except Exception:
        logger.debug("tailscale binary lookup failed via 'which'", exc_info=True)

    import os

    mac_path = "/Applications/Tailscale.app/Contents/MacOS/Tailscale"
    if os.path.exists(mac_path) and os.access(mac_path, os.X_OK):
        return mac_path
    return None


async def read_tailscale_whois_identity(
    ip: str,
    timeout_s: float = WHOIS_TIMEOUT_S,
    cache_ttl_ms: float = WHOIS_CACHE_TTL_MS,
    error_ttl_ms: float = WHOIS_ERROR_TTL_MS,
) -> TailscaleWhoisIdentity | None:
    normalized = ip.strip()
    if not normalized:
        return None

    now_ms = time.time() * 1000
    cached = _read_cached_whois(normalized, now_ms)
    if cached is not None:
        return cached if cached is not False else None

    tailscale_bin = await _get_tailscale_binary()
    if not tailscale_bin:
        logger.debug("tailscale binary not found, whois unavailable")
        _write_cached_whois(normalized, None, error_ttl_ms)
        return None

    proc: asyncio.subprocess.Process | None = None
    try:
        proc = await asyncio.create_subprocess_exec(
            tailscale_bin,
            "whois",
            "--json",
            normalized,
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE,
        )
        stdout, stderr = await asyncio.wait_for(
            proc.communicate(),
            timeout=timeout_s,
        )
        if proc.returncode != 0:
            logger.debug("tailscale whois failed for %s: %s", normalized, stderr.decode().strip())
            _write_cached_whois(normalized, None, error_ttl_ms)
            return None
        payload = _parse_possibly_noisy_json(stdout.decode())
        identity = _parse_whois_identity(payload)
        _write_cached_whois(normalized, identity, cache_ttl_ms)
        return identity
    except TimeoutError:
        logger.debug("tailscale whois timeout for %s", normalized)
        if proc is not None:
            try:
                proc.kill()
                await proc.wait()
            except Exception:
                pass
        _write_cached_whois(normalized, None, error_ttl_ms)
        return None
    except Exception:
        logger.debug("tailscale whois error for %s", normalized, exc_info=True)
        if proc is not None:
            try:
                proc.kill()
                await proc.wait()
            except Exception:
                pass
        _write_cached_whois(normalized, None, error_ttl_ms)
        return None


def get_tailscale_user_from_headers(headers: dict | None) -> tuple[str, str] | None:
    if not headers:
        return None
    login = headers.get("tailscale-user-login")
    if not login:
        return None
    name = headers.get("tailscale-user-name") or login
    return (login.lower().strip(), name.strip() if isinstance(name, str) else login)