hermes-sshx-link/__init__.py

"""
sshx-link plugin — creates a shell at sshx.io and returns the shareable link
============================================================================
- Auto arch detection (x86_64 / aarch64 / armv6 / armv7)
- Downloads sshx binary from S3 without curl|sh
- Starts `sshx` and parses https://sshx.io/s/... link from stdout
- No dashboard restart needed
"""
import subprocess
import threading
import logging
import shutil
import os
import platform
import stat
import urllib.request
import re
import tarfile
import tempfile
import time

logger = logging.getLogger(__name__)

_sshx_process = None
_sshx_link = None
_sshx_log = "/tmp/sshx_link.log"

SSHX_BIN = "/tmp/sshx"
SSHX_BIN_ALT = os.path.expanduser("~/.local/bin/sshx")
SSHX_TAR = "/tmp/sshx.tar.gz"

# Reuse ttyd fallback for offline mode (optional)
TTYD_BIN = "/tmp/ttyd"

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
    # default linux
    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"Detected arch={arch} suffix={suffix} -> sshx url={url}")
    return url


def _find_sshx_bin() -> str | None:
    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_installed() -> str | None:
    """Ensure sshx binary exists, download + extract if needed. Returns path or None."""
    existing = _find_sshx_bin()
    if existing:
        logger.info(f"sshx already installed at {existing}")
        return existing

    url = _get_sshx_url()
    dest = SSHX_BIN
    os.makedirs(os.path.dirname(SSHX_BIN_ALT), exist_ok=True)
    tmp_tar = SSHX_TAR
    try:
        logger.info(f"Downloading sshx from {url} -> {tmp_tar}")
        urllib.request.urlretrieve(url, tmp_tar)
        # extract tar.gz - contains single file 'sshx'
        logger.info(f"Extracting {tmp_tar}")
        with tarfile.open(tmp_tar, "r:gz") as tf:
            # find member named sshx, skip AppleDouble ._ files
            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]
            # extract to tmp dir then move
            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)
            # member.name may contain subdir
            if not os.path.exists(extracted):
                # try find sshx recursively
                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)
            # cache copy
            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 successfully to {dest}")
        return dest
    except Exception as e:
        logger.error(f"sshx install failed from {url}: {e}")
        return None


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


_sshx_last_output = ""
_sshx_last_error = ""

def _capture_link(proc, timeout=20) -> str | None:
    """Read stdout until link found or timeout. Returns link or None."""
    global _sshx_last_output
    link = None
    start = time.time()
    import select

    output = ""
    while time.time() - start < timeout:
        if proc.poll() is not None:
            # process died, collect remaining
            try:
                rem = proc.stdout.read() or ""
                output += rem
            except Exception:
                pass
            break
        # check if data available
        try:
            # select on stdout fd
            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
                # also check whole output so far (link may span?)
                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)
    _sshx_last_output = output[-2000:]
    if not link:
        logger.warning(f"sshx capture timeout, output tail: {output[-600:]}")
    return link


def _run_sshx(tool_ctx):
    """Start sshx and return link. Reuse if already running."""
    global _sshx_process, _sshx_link

    if _is_sshx_alive() and _sshx_link:
        tool_ctx.yield_result({"status": "already_running", "link": _sshx_link})
        return

    bin_path = _ensure_sshx_installed()
    if not bin_path:
        tool_ctx.yield_result({"status": "error", "message": "Failed to install sshx. Check network to s3.amazonaws.com/sshx"})
        return

    logger.info(f"Starting sshx via {bin_path} ...")
    try:
        # --quiet prints only the link (https://sshx.io/s/...), far more reliable to parse
        # fallback to no-flag if --quiet not supported (old version)
        _sshx_process = subprocess.Popen(
            [bin_path, "--quiet"],
            stdout=subprocess.PIPE,
            stderr=subprocess.STDOUT,
            text=True,
            bufsize=1,
            stdin=subprocess.DEVNULL,
            start_new_session=True,
        )
    except Exception as e:
        logger.error(f"Failed to start sshx: {e}")
        tool_ctx.yield_result({"status": "error", "message": f"Failed to start sshx: {e}"})
        return

    link = _capture_link(_sshx_process, timeout=20)
    if link:
        _sshx_link = link
        logger.info(f"sshx running at {_sshx_link} pid={_sshx_process.pid}")
        tool_ctx.yield_result({"status": "running", "link": _sshx_link, "pid": _sshx_process.pid})
    else:
        global _sshx_last_error
        if _sshx_process.poll() is not None:
            try:
                out = _sshx_process.stdout.read() or ""
            except Exception:
                out = _sshx_last_output
            _sshx_last_error = f"sshx exited quickly (code={_sshx_process.poll()}): {(_sshx_last_output or out)[:600]}"
            tool_ctx.yield_result({"status": "error", "message": _sshx_last_error, "output_tail": _sshx_last_output[-400:]})
            _sshx_process = None
        else:
            _sshx_last_error = f"sshx started (pid={_sshx_process.pid}) but link not parsed in 20s. tail: {_sshx_last_output[:400]}. Check outbound to https://sshx.io"
            tool_ctx.yield_result({"status": "error", "message": _sshx_last_error, "output_tail": _sshx_last_output[-400:], "pid": _sshx_process.pid})


def _stop_sshx(tool_ctx=None):
    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
    if tool_ctx:
        tool_ctx.yield_result({"status": "stopped"})


def register(ctx):
    # warmup download in background — no restart needed
    def _warmup():
        try:
            if not _find_sshx_bin():
                logger.info("Warming up sshx install...")
                _ensure_sshx_installed()
        except Exception as e:
            logger.debug(f"warmup failed: {e}")

    threading.Thread(target=_warmup, daemon=True).start()

    ctx.register_tool(
        name="sshx_start",
        toolset="sshx-link",
        schema={
            "description": "Create a shell at sshx.io and return the shareable link. Auto arch detection, downloads sshx binary if needed.",
            "parameters": {"type": "object", "properties": {}, "required": []},
        },
        handler=_run_sshx,
    )
    ctx.register_tool(
        name="sshx_stop",
        toolset="sshx-link",
        schema={
            "description": "Stop the running sshx session",
            "parameters": {"type": "object", "properties": {}, "required": []},
        },
        handler=_stop_sshx,
    )
    # alias for compat with old ttyd plugin
    ctx.register_tool(
        name="ttyd_start",
        toolset="sshx-link",
        schema={
            "description": "Alias of sshx_start (kept for compatibility)",
            "parameters": {"type": "object", "properties": {}, "required": []},
        },
        handler=_run_sshx,
    )
    logger.info("sshx-link plugin registered (sshx.io, auto arch, no restart)")