pikasTech-HWLAB/tools/unidesk-ssh.mjs

#!/usr/bin/env bun

const route = process.argv[2] ?? "";
const commandArgs = process.argv.slice(3);

if (route === "" || commandArgs.includes("--help") || commandArgs.includes("-h")) {
  printUsage(route === "");
}

const token = optionalEnv("UNIDESK_SSH_CLIENT_TOKEN");
if (!token) fail("UNIDESK_SSH_CLIENT_TOKEN is required; request tool=unidesk-ssh through AgentRun toolCredentials.");

const frontendUrl = frontendWebSocketUrl(requiredFrontendBaseUrl());
const parsedRoute = parseRoute(route);
const remoteCommand = buildRemoteCommand(parsedRoute, commandArgs);
const openTimeoutMs = positiveEnv("UNIDESK_SSH_OPEN_TIMEOUT_MS", 60000);
const runtimeTimeoutMs = positiveEnv("UNIDESK_SSH_RUNTIME_TIMEOUT_MS", 60000);
const ws = new WebSocket(frontendUrl, { headers: { authorization: `Bearer ${token}` } });

let exitCode = 255;
let settled = false;
let canSend = false;
let sessionReady = false;
const pending = [];
const pendingSessionMessages = [];

const openTimer = setTimeout(() => {
  if (sessionReady || settled) return;
  process.stderr.write("unidesk-ssh timed out waiting for provider session\n");
  closeSocket();
}, openTimeoutMs);

const runtimeTimer = setTimeout(() => {
  if (settled) return;
  exitCode = 124;
  process.stderr.write("unidesk-ssh runtime timeout; use short query plus poll semantics for long work\n");
  closeSocket();
}, runtimeTimeoutMs);

ws.addEventListener("open", () => {
  canSend = true;
  send({
    type: "ssh.open",
    providerId: parsedRoute.providerId,
    command: remoteCommand,
    cwd: parsedRoute.cwd ?? undefined,
    tty: false,
    stdinEotOnEnd: true,
    openTimeoutMs,
    runtimeTimeoutMs,
    cols: Number(process.stdout.columns) > 0 ? Number(process.stdout.columns) : 100,
    rows: Number(process.stdout.rows) > 0 ? Number(process.stdout.rows) : 30
  });
  flushPending();
});

ws.addEventListener("message", (event) => {
  let message;
  try {
    message = JSON.parse(webSocketDataText(event.data));
  } catch {
    process.stderr.write(`${webSocketDataText(event.data)}\n`);
    return;
  }

  if (message.type === "ssh.dispatched") return;
  if (message.type === "ssh.opened") {
    sessionReady = true;
    clearTimeout(openTimer);
    sendWhenSessionReady({ type: "ssh.eof" });
    flushSessionMessages();
    return;
  }
  if (message.type === "ssh.data") {
    const chunk = Buffer.from(String(message.data ?? ""), message.encoding === "base64" ? "base64" : "utf8");
    if (message.stream === "stderr") process.stderr.write(chunk);
    else process.stdout.write(chunk);
    return;
  }
  if (message.type === "ssh.error") {
    process.stderr.write(`${String(message.failureKind ?? "ssh-error")}: ${String(message.message ?? "ssh bridge error")}\n`);
    exitCode = 255;
    closeSocket();
    return;
  }
  if (message.type === "ssh.exit") {
    exitCode = Number.isInteger(message.exitCode) ? Number(message.exitCode) : 255;
    closeSocket();
  }
});

ws.addEventListener("close", () => finish(exitCode));
ws.addEventListener("error", () => {
  process.stderr.write("unidesk-ssh websocket error\n");
  finish(255);
});

function send(value) {
  const text = JSON.stringify(value);
  if (!canSend || ws.readyState !== WebSocket.OPEN) {
    pending.push(text);
    return;
  }
  ws.send(text);
}

function sendWhenSessionReady(value) {
  const text = JSON.stringify(value);
  if (!sessionReady || ws.readyState !== WebSocket.OPEN) {
    pendingSessionMessages.push(text);
    return;
  }
  ws.send(text);
}

function flushPending() {
  while (pending.length > 0 && ws.readyState === WebSocket.OPEN) ws.send(pending.shift());
}

function flushSessionMessages() {
  if (!sessionReady || ws.readyState !== WebSocket.OPEN) return;
  while (pendingSessionMessages.length > 0) ws.send(pendingSessionMessages.shift());
}

function closeSocket() {
  try {
    ws.close();
  } catch {
    finish(exitCode);
  }
}

function finish(code) {
  if (settled) return;
  settled = true;
  clearTimeout(openTimer);
  clearTimeout(runtimeTimer);
  process.exit(code);
}

function parseRoute(value) {
  const providerWorkspace = /^([^:]+):(\/.*)$/u.exec(value);
  if (providerWorkspace) return { providerId: providerWorkspace[1], cwd: providerWorkspace[2], plane: "host" };
  const parts = value.split(":");
  const providerId = parts[0] ?? "";
  if (!/^[A-Za-z0-9._-]+$/u.test(providerId)) fail(`invalid UniDesk provider route: ${value}`);
  const plane = parts[1] ?? "host";
  if (parts.length > 2) fail("unidesk-ssh supports provider, provider:/workspace, and provider:k3s routes only; use the full UniDesk CLI for nested routes.");
  if (plane !== "host" && plane !== "k3s") fail(`unsupported UniDesk route plane: ${plane}`);
  return { providerId, cwd: null, plane };
}

function buildRemoteCommand(parsed, args) {
  if (args.length === 0) fail("unidesk-ssh requires a command, for example: unidesk-ssh G14:/root/hwlab script -- 'pwd'");
  const [operation, ...rest] = args;
  let command;
  if (operation === "script" || operation === "shell" || operation === "sh") {
    const scriptArgs = rest[0] === "--" ? rest.slice(1) : rest;
    if (scriptArgs.length === 0) fail(`${operation} requires a script string`);
    command = `sh -lc ${shellQuote(scriptArgs.join(" "))}`;
  } else if (operation === "argv") {
    if (rest.length === 0) fail("argv requires at least one command argument");
    command = rest.map(shellQuote).join(" ");
  } else {
    command = args.map(shellQuote).join(" ");
  }
  if (parsed.plane === "k3s") return `export KUBECONFIG=/etc/rancher/k3s/k3s.yaml; ${command}`;
  return command;
}

function requiredFrontendBaseUrl() {
  const raw = optionalEnv("UNIDESK_FRONTEND_URL")
    ?? optionalEnv("UNIDESK_MAIN_SERVER_URL")
    ?? optionalEnv("UNIDESK_MAIN_SERVER_IP")
    ?? optionalEnv("UNIDESK_MAIN_SERVER_HOST");
  if (!raw) fail("UNIDESK_MAIN_SERVER_IP or UNIDESK_FRONTEND_URL is required for UniDesk SSH passthrough.");
  return /^https?:\/\//u.test(raw) ? raw : `http://${raw}`;
}

function frontendWebSocketUrl(base) {
  const url = new URL("/ws/ssh", base);
  url.protocol = url.protocol === "https:" ? "wss:" : "ws:";
  return url.toString();
}

function optionalEnv(name) {
  const value = process.env[name]?.trim() ?? "";
  return value.length > 0 ? value : null;
}

function positiveEnv(name, fallback) {
  const raw = optionalEnv(name);
  if (!raw) return fallback;
  const parsed = Number(raw);
  if (!Number.isFinite(parsed) || parsed <= 0) fail(`${name} must be a positive number`);
  return parsed;
}

function webSocketDataText(data) {
  if (typeof data === "string") return data;
  if (data instanceof ArrayBuffer) return Buffer.from(data).toString("utf8");
  if (ArrayBuffer.isView(data)) return Buffer.from(data.buffer, data.byteOffset, data.byteLength).toString("utf8");
  return String(data);
}

function shellQuote(value) {
  return `'${String(value).replace(/'/gu, `'"'"'`)}'`;
}

function fail(message) {
  process.stderr.write(`${message}\n`);
  process.exit(2);
}

function printUsage(failed) {
  process.stderr.write(`Usage: unidesk-ssh <provider[:/workspace]|provider:k3s> <script|argv|command> [args...]\n\nExamples:\n  unidesk-ssh G14:/root/hwlab script -- 'pwd && git status --short --branch'\n  unidesk-ssh G14 argv hostname\n  unidesk-ssh G14:k3s argv kubectl get ns agentrun-v01\n`);
  process.exit(failed ? 2 : 0);
}