326 lines
11 KiB
Python
326 lines
11 KiB
Python
"""
|
|
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
|
|
|
|
|
|
def _capture_link(proc, timeout=15) -> str | None:
|
|
"""Read stdout until link found or timeout. Returns link or None."""
|
|
link = None
|
|
start = time.time()
|
|
# use non-blocking read via thread? simpler: read line with timeout via poll
|
|
# We'll read via proc.stdout readline with timeout loop
|
|
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)
|
|
# also try reading logfile if we used file
|
|
if not link:
|
|
logger.warning(f"sshx capture timeout, output so far: {output[:500]}")
|
|
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=15)
|
|
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:
|
|
# process may still be alive but link not captured — check if died
|
|
if _sshx_process.poll() is not None:
|
|
try:
|
|
out = _sshx_process.stdout.read() or ""
|
|
except Exception:
|
|
out = ""
|
|
tool_ctx.yield_result({"status": "error", "message": f"sshx exited quickly (code={_sshx_process.poll()}): {out[:300]}"})
|
|
_sshx_process = None
|
|
else:
|
|
# still running but no link parsed — return with log hint
|
|
tool_ctx.yield_result({"status": "running", "link": None, "message": "sshx started but link not parsed yet, check logs", "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)")
|