hermes-sshx-link/dashboard/plugin_api.py

import asyncio
import subprocess
import os
import logging
import time
import platform
import stat
import urllib.request
import re
import tarfile
import tempfile
import shutil
import select
from typing import Optional
from fastapi import APIRouter, Request
from fastapi.responses import JSONResponse

logger = logging.getLogger(__name__)
router = APIRouter()

_sshx_link = None
_sshx_process: Optional[subprocess.Popen] = None
_sshx_installing = False
_sshx_bin = "/tmp/sshx"
_sshx_bin_alt = os.path.expanduser("~/.local/bin/sshx")
_sshx_tar = "/tmp/sshx.tar.gz"

LINK_RE = re.compile(r"https://sshx\.io/s/[A-Za-z0-9\-_]+(?:#[^\s\"']*)?")


def _detect_sshx_arch():
    machine = platform.machine().lower()
    system = platform.system().lower()
    if system == "darwin":
        suffix = "-apple-darwin"
        if machine in ("aarch64", "arm64", "armv8l", "armv8b"):
            arch = "aarch64"
        elif machine in ("x86_64", "x64", "amd64"):
            arch = "x86_64"
        else:
            arch = "aarch64" if "arm" in machine else "x86_64"
        return arch, suffix
    suffix = "-unknown-linux-musl"
    if machine in ("aarch64", "aarch64_be", "arm64", "armv8b", "armv8l"):
        arch = "aarch64"
    elif machine in ("x86_64", "x64", "amd64"):
        arch = "x86_64"
    elif machine == "armv6l":
        arch = "arm"
        suffix += "eabihf"
    elif machine == "armv7l":
        arch = "armv7"
        suffix += "eabihf"
    else:
        arch = "x86_64"
    return arch, suffix


def _get_sshx_url() -> str:
    arch, suffix = _detect_sshx_arch()
    url = f"https://s3.amazonaws.com/sshx/sshx-{arch}{suffix}.tar.gz"
    logger.info(f"sshx arch={arch} suffix={suffix} url={url}")
    return url


def _find_sshx_bin() -> Optional[str]:
    for p in [_sshx_bin, _sshx_bin_alt, shutil.which("sshx")]:
        if p and os.path.exists(p) and os.access(p, os.X_OK):
            return p
    return None


def _ensure_sshx() -> Optional[str]:
    existing = _find_sshx_bin()
    if existing:
        return existing
    global _sshx_installing
    _sshx_installing = True
    try:
        url = _get_sshx_url()
        dest = _sshx_bin
        os.makedirs(os.path.dirname(_sshx_bin_alt), exist_ok=True)
        tmp_tar = _sshx_tar
        logger.info(f"Downloading sshx {url} -> {tmp_tar}")
        urllib.request.urlretrieve(url, tmp_tar)
        logger.info(f"Extracting {tmp_tar}")
        with tarfile.open(tmp_tar, "r:gz") as tf:
            member = None
            for m in tf.getmembers():
                base = os.path.basename(m.name)
                if base.startswith("._"):
                    continue
                if base == "sshx" and m.isfile():
                    member = m
                    break
            if not member:
                for m in tf.getmembers():
                    base = os.path.basename(m.name)
                    if base.startswith("._"):
                        continue
                    if m.isfile():
                        member = m
                        break
            if not member:
                member = tf.getmembers()[0]
            tmpdir = tempfile.mkdtemp()
            try:
                tf.extract(member, path=tmpdir, filter='fully_trusted')
            except TypeError:
                tf.extract(member, path=tmpdir)
            extracted = os.path.join(tmpdir, member.name)
            if not os.path.exists(extracted):
                for root, _, files in os.walk(tmpdir):
                    if "sshx" in files:
                        extracted = os.path.join(root, "sshx")
                        break
            shutil.copy2(extracted, dest)
            os.chmod(dest, os.stat(dest).st_mode | stat.S_IEXEC)
            try:
                if dest != _sshx_bin_alt:
                    shutil.copy2(dest, _sshx_bin_alt)
                    os.chmod(_sshx_bin_alt, os.stat(_sshx_bin_alt).st_mode | stat.S_IEXEC)
            except Exception:
                pass
            shutil.rmtree(tmpdir, ignore_errors=True)
        try:
            os.remove(tmp_tar)
        except Exception:
            pass
        logger.info(f"sshx installed to {dest}")
        return dest
    except Exception as e:
        logger.error(f"sshx install failed: {e}")
        return None
    finally:
        _sshx_installing = False


def _is_alive() -> bool:
    global _sshx_process
    if not _sshx_process:
        return False
    return _sshx_process.poll() is None


def _kill_sshx():
    global _sshx_process, _sshx_link
    if _sshx_process and _sshx_process.poll() is None:
        try:
            os.killpg(os.getpgid(_sshx_process.pid), 15)
        except Exception:
            try:
                _sshx_process.terminate()
            except Exception:
                pass
        time.sleep(0.5)
        if _sshx_process.poll() is None:
            try:
                os.killpg(os.getpgid(_sshx_process.pid), 9)
            except Exception:
                try:
                    _sshx_process.kill()
                except Exception:
                    pass
    _sshx_process = None
    _sshx_link = None


def _capture_link_sync(proc, timeout=15) -> Optional[str]:
    link = None
    start = time.time()
    output = ""
    while time.time() - start < timeout:
        if proc.poll() is not None:
            try:
                rem = proc.stdout.read() or ""
                output += rem
            except Exception:
                pass
            break
        try:
            r, _, _ = select.select([proc.stdout], [], [], 0.5)
            if r:
                line = proc.stdout.readline()
                if not line:
                    time.sleep(0.1)
                    continue
                output += line
                logger.info(f"sshx stdout: {line.strip()}")
                m = LINK_RE.search(line)
                if m:
                    link = m.group(0)
                    break
                m2 = LINK_RE.search(output)
                if m2:
                    link = m2.group(0)
                    break
        except Exception as e:
            logger.debug(f"capture error: {e}")
            time.sleep(0.2)
    if not link:
        m = LINK_RE.search(output)
        if m:
            link = m.group(0)
    if not link:
        logger.warning(f"sshx link not found after {timeout}s, output: {output[:600]}")
    return link


def _start_sshx_background() -> bool:
    global _sshx_process, _sshx_link
    _kill_sshx()
    bin_path = _ensure_sshx()
    if not bin_path:
        return False
    logger.info(f"Starting sshx {bin_path}")
    try:
        proc = subprocess.Popen(
            [bin_path, "--quiet"],
            stdout=subprocess.PIPE,
            stderr=subprocess.STDOUT,
            text=True,
            bufsize=1,
            stdin=subprocess.DEVNULL,
            start_new_session=True,
        )
        _sshx_process = proc
        link = _capture_link_sync(proc, timeout=15)
        if link:
            _sshx_link = link
            logger.info(f"sshx running link={_sshx_link} pid={proc.pid}")
            return True
        # if link not found but process still alive, keep it and return false to let caller handle
        if proc.poll() is None:
            logger.warning("sshx started but link not parsed yet")
            return False
        logger.error(f"sshx exited quickly code={proc.poll()}")
        _sshx_process = None
        return False
    except Exception as e:
        logger.error(f"Failed to start sshx: {e}")
        return False


@router.get("/start")
async def start_sshx(request: Request):
    global _sshx_link, _sshx_process
    if _is_alive() and _sshx_link:
        return {"status": "running", "link": _sshx_link, "pid": _sshx_process.pid if _sshx_process else None}

    bin_path = _find_sshx_bin()
    if not bin_path:
        loop = asyncio.get_event_loop()
        bin_path = await loop.run_in_executor(None, _ensure_sshx)
        if not bin_path:
            return JSONResponse(status_code=500, content={"status": "error", "message": "Failed to install sshx (arch auto-detect). Check network to s3.amazonaws.com"})

    loop = asyncio.get_event_loop()
    ok = await loop.run_in_executor(None, _start_sshx_background)

    if ok and _sshx_link:
        return {"status": "running", "link": _sshx_link, "pid": _sshx_process.pid if _sshx_process else None}
    # if process alive but link not yet, try one more check
    if _is_alive():
        return JSONResponse(status_code=202, content={"status": "starting", "message": "sshx started but link not ready yet, retry /status", "pid": _sshx_process.pid if _sshx_process else None})
    return JSONResponse(status_code=504, content={"status": "error", "message": "sshx did not start in time, check logs"})


@router.post("/start")
async def start_sshx_post(request: Request):
    return await start_sshx(request)


@router.get("/status")
async def sshx_status():
    global _sshx_link, _sshx_installing, _sshx_process
    if _is_alive() and _sshx_link:
        return {"status": "running", "link": _sshx_link, "pid": _sshx_process.pid}
    if _sshx_installing:
        return {"status": "installing", "link": None}
    if _is_alive() and not _sshx_link:
        return {"status": "starting", "link": None, "pid": _sshx_process.pid}
    return {"status": "stopped", "link": None}


@router.post("/stop")
async def stop_sshx():
    _kill_sshx()
    return {"status": "stopped"}


@router.get("/stop")
async def stop_sshx_get():
    _kill_sshx()
    return {"status": "stopped"}


# compat: old ttyd endpoint still works but now maps to sshx
@router.get("/restart-dashboard")
async def restart_dashboard():
    return {"status": "deprecated", "message": "No restart needed. This endpoint is no-op."}


@router.post("/restart-dashboard")
async def restart_dashboard_post():
    return {"status": "deprecated", "message": "No restart needed. This endpoint is no-op."}