hermes-sshx-link/dashboard/plugin_api.py

import asyncio
import subprocess
import re
import shutil
import logging
import threading
from typing import Optional
from fastapi import APIRouter
from fastapi.responses import JSONResponse

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

_sshx_process = None
_sshx_link = None

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


def _ensure_sshx() -> bool:
    if shutil.which("sshx"):
        return True
    try:
        subprocess.run(
            ["sh", "-c", "curl -sSf https://sshx.io/get | sh"],
            capture_output=True, timeout=60,
        )
        return shutil.which("sshx") is not None
    except Exception as e:
        logger.error("sshx install failed: %s", e)
        return False


def _start_reader(proc, future):
    global _sshx_link

    def _read():
        try:
            found = False
            for line in proc.stdout:
                line = line.strip()
                if not found:
                    m = LINK_RE.search(line)
                    if m:
                        _sshx_link = m.group(0)
                        found = True
                        if not future.done():
                            future.set_result(_sshx_link)
            if not found and not future.done():
                future.set_exception(RuntimeError("sshx exited without link"))
        except Exception as e:
            if not future.done():
                future.set_exception(e)

    t = threading.Thread(target=_read, daemon=True)
    t.start()
    return t


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

    if _sshx_process and _sshx_process.poll() is None and _sshx_link:
        return {"status": "running", "link": _sshx_link}

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

    if _sshx_process and _sshx_process.poll() is None:
        _sshx_process.terminate()
        try:
            _sshx_process.wait(timeout=5)
        except subprocess.TimeoutExpired:
            _sshx_process.kill()

    _sshx_process = subprocess.Popen(
        ["sshx", "--quiet"],
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT,
        text=True,
        bufsize=1,
    )

    _sshx_link = None
    loop = asyncio.get_event_loop()
    link_future = loop.create_future()
    _start_reader(_sshx_process, link_future)

    try:
        link = await asyncio.wait_for(link_future, timeout=30)
        return {"status": "running", "link": link}
    except asyncio.TimeoutError:
        return JSONResponse(status_code=504, content={"status": "error", "message": "sshx did not return a link in time"})
    except Exception as e:
        return JSONResponse(status_code=500, content={"status": "error", "message": str(e)})


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


@router.post("/stop")
async def stop_sshx():
    global _sshx_process, _sshx_link
    if _sshx_process and _sshx_process.poll() is None:
        _sshx_process.terminate()
        try:
            _sshx_process.wait(timeout=5)
        except subprocess.TimeoutExpired:
            _sshx_process.kill()
        _sshx_process = None
        _sshx_link = None
        return {"status": "stopped"}
    return {"status": "already_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,
        )
        return {"status": "restarting"}
    except Exception as e:
        return JSONResponse(status_code=500, content={"status": "error", "message": str(e)})