hermes-sshx-link/dashboard/plugin_api.py

import asyncio
import subprocess
import re
import os
import shutil
import logging
import time
from typing import Optional
from fastapi import APIRouter, BackgroundTasks
from fastapi.responses import JSONResponse

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

_sshx_link = None
_sshx_pid = None
_sshx_installing = False
_sshx_log = "/tmp/sshx_link.log"
_sshx_pid_file = "/tmp/sshx_link.pid"

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

import threading as _threading


def _auto_install():
    global _sshx_installing
    if shutil.which("sshx"):
        logger.info("sshx already installed at %s", shutil.which("sshx"))
        return
    _sshx_installing = True
    logger.info("Auto-installing sshx...")
    try:
        subprocess.run(
            ["sh", "-c", "curl -sSf https://sshx.io/get | sh"],
            capture_output=True, timeout=60,
        )
        if shutil.which("sshx"):
            logger.info("sshx installed successfully")
            _trigger_dashboard_restart()
        else:
            logger.error("sshx install failed")
    except Exception as e:
        logger.error("sshx auto-install error: %s", e)
    finally:
        _sshx_installing = False


@router.on_event("startup")
def _startup_auto_install():
    _threading.Thread(target=_auto_install, daemon=True).start()


def _ensure_sshx() -> tuple[bool, bool]:
    """Ensure sshx is installed. Returns (success, needs_restart)."""
    if shutil.which("sshx"):
        return True, False
    try:
        subprocess.run(
            ["sh", "-c", "curl -sSf https://sshx.io/get | sh"],
            capture_output=True, timeout=60,
        )
        if shutil.which("sshx"):
            return True, True  # newly installed, needs restart
        return False, False
    except Exception as e:
        logger.error("sshx install failed: %s", e)
        return False, False


def _trigger_dashboard_restart():
    try:
        subprocess.Popen(
            ["nohup", "bash", "-c",
             "sleep 3 && sudo systemctl restart hermes-dashboard"],
            stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL,
            stdin=subprocess.DEVNULL, start_new_session=True,
        )
    except Exception as e:
        logger.error("Failed to trigger dashboard restart: %s", e)


def _kill_sshx():
    global _sshx_pid, _sshx_link
    if _sshx_pid:
        try:
            os.kill(_sshx_pid, 15)
        except (ProcessLookupError, PermissionError):
            pass
    _sshx_pid = None
    _sshx_link = None
    for f in [_sshx_log, _sshx_pid_file]:
        if os.path.exists(f):
            os.remove(f)


def _is_alive() -> bool:
    global _sshx_pid
    if not _sshx_pid:
        return False
    try:
        os.kill(_sshx_pid, 0)
        return True
    except (ProcessLookupError, PermissionError):
        return False


def _read_link_from_log() -> Optional[str]:
    if not os.path.exists(_sshx_log):
        return None
    try:
        with open(_sshx_log) as f:
            content = f.read()
        m = LINK_RE.search(content)
        if m:
            return m.group(0)
    except Exception:
        pass
    return None


def _start_sshx_background():
    global _sshx_pid, _sshx_link
    _kill_sshx()
    _sshx_link = None

    log_fh = open(_sshx_log, "w")
    proc = subprocess.Popen(
        ["sshx", "--quiet"],
        stdout=log_fh,
        stderr=subprocess.DEVNULL,
        stdin=subprocess.DEVNULL,
        start_new_session=True,
    )
    _sshx_pid = proc.pid
    with open(_sshx_pid_file, "w") as pf:
        pf.write(str(proc.pid))

    for _ in range(30):
        time.sleep(1)
        link = _read_link_from_log()
        if link:
            _sshx_link = link
            return True
        if proc.poll() is not None:
            return False
    return False


@router.get("/start")
async def start_sshx(background_tasks: BackgroundTasks):
    global _sshx_link, _sshx_pid

    if _is_alive() and _sshx_link:
        return {"status": "running", "link": _sshx_link}

    ok, needs_restart = _ensure_sshx()
    if not ok:
        return JSONResponse(status_code=500, content={"status": "error", "message": "Failed to install sshx"})

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

    # Defer dashboard restart until AFTER the response is sent
    if needs_restart:
        background_tasks.add_task(_trigger_dashboard_restart)

    if result:
        return {"status": "running", "link": _sshx_link}
    return JSONResponse(status_code=504, content={"status": "error", "message": "sshx did not return a link in time"})


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


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


@router.post("/restart-dashboard")
async def restart_dashboard():
    try:
        subprocess.Popen(
            ["nohup", "bash", "-c", "sleep 1 && sudo systemctl restart hermes-dashboard"],
            stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL,
            start_new_session=True,
        )
        return {"status": "restarting"}
    except Exception as e:
        return JSONResponse(status_code=500, content={"status": "error", "message": str(e)})