add sf vm ssh (#107)

Author: coffinsfcompute

deno.lock

@@ -8,6 +8,7 @@
     "test": "deno test --allow-all"
   },
   "dependencies": {
+    "@alphahydrae/exec": "^1.1.0",
     "@commander-js/extra-typings": "^12.1.0",
     "@inkjs/ui": "^2.0.0",
     "@inquirer/prompts": "^5.1.2",
@@ -36,6 +37,7 @@
     "prettier": "^3.5.3",
     "react": "^18.3.1",
     "semver": "^7.6.3",
+    "shescape": "^2.1.1",
     "tiny-invariant": "^1.3.3",
     "tweetnacl": "^1.0.3",
     "tweetnacl-util": "^0.15.1",

package.json

@@ -45,9 +45,10 @@ const apiPaths: Record<string, Path<IdParams | never>> = {
 
   vms_instances_list: "/v0/vms/instances",
   vms_logs_list: "/v0/vms/logs",
+  vms_replace: "/v0/vms/replace",
   vms_script_post: "/v0/vms/script",
   vms_script_get: "/v0/vms/script",
-  vms_replace: "/v0/vms/replace",
+  vms_ssh_get: "/v0/vms/ssh",
 };
 
 export async function getWebAppUrl<T extends TokenParams | never>(

src/helpers/urls.ts

@@ -0,0 +1,135 @@
+import type { Command } from "@commander-js/extra-typings";
+import console from "node:console";
+import process from "node:process";
+import { Shescape } from "shescape";
+import { execvp } from "@alphahydrae/exec";
+import {
+  logAndQuit,
+  logSessionTokenExpiredAndQuit,
+} from "../helpers/errors.ts";
+import { getApiUrl } from "../helpers/urls.ts";
+import { getAuthToken } from "../helpers/config.ts";
+
+type SshHostKey = {
+  keyType: string;
+  base64EncodedKey: string;
+};
+
+export function registerSsh(program: Command) {
+  program
+    .command("ssh")
+    .option("-q, --quiet", "Quiet mode", false)
+    .option(
+      "--use-host-keys [value]",
+      "Use API provided SSH server host keys if known",
+      true,
+    )
+    .argument(
+      "<destination>",
+      "USERNAME@VM_ID The (optional) username, and VM id to SSH into.",
+    )
+    .allowExcessArguments(false)
+    .action(async (destination, options) => {
+      const splitDestination = destination.split("@");
+      let vmId: string;
+      let sshUsername: string | undefined;
+      if (splitDestination.length == 1) {
+        sshUsername = undefined;
+        vmId = splitDestination[0];
+      } else if (splitDestination.length == 2) {
+        sshUsername = splitDestination[0];
+        vmId = splitDestination[1];
+      } else {
+        logAndQuit(`Invalid SSH destination string: ${destination}`);
+      }
+
+      const baseUrl = await getApiUrl("vms_ssh_get");
+      const params = new URLSearchParams();
+      params.append("vm_id", vmId);
+      const url = `${baseUrl}?${params.toString()}`;
+      const response = await fetch(url, {
+        method: "GET",
+        headers: {
+          Authorization: `Bearer ${await getAuthToken()}`,
+        },
+      });
+
+      if (!response.ok) {
+        if (response.status === 401) {
+          logSessionTokenExpiredAndQuit();
+        }
+
+        logAndQuit(
+          `Failed to retrieve ssh information: ${response.statusText}`,
+        );
+      }
+
+      const data = (await response.json()) as {
+        ssh_hostname: string;
+        ssh_port: number;
+        ssh_host_keys: SshHostKey[] | undefined;
+      };
+      const sshHostname = data.ssh_hostname;
+      const sshPort = data.ssh_port;
+      const sshHostKeys = data.ssh_host_keys || [];
+
+      let sshDestination = sshHostname;
+      if (sshUsername !== undefined) {
+        sshDestination = `${sshUsername}@${sshDestination}`;
+      }
+      if (sshPort !== undefined) {
+        sshDestination = `${sshDestination}:${sshPort}`;
+      }
+      sshDestination = `ssh://${sshDestination}`;
+
+      let cmd = ["ssh"];
+
+      if (sshHostKeys.length > 0 && options.useHostKeys) {
+        let knownHostsCommand = ["/usr/bin/env", "printf", "%s %s %s\\n"];
+        for (const sshHostKey of sshHostKeys) {
+          knownHostsCommand = knownHostsCommand.concat([
+            sshHostname,
+            sshHostKey.keyType,
+            sshHostKey.base64EncodedKey,
+          ]);
+        }
+        // Escape all characters for proper pass through
+        for (const i in knownHostsCommand) {
+          knownHostsCommand[i] = knownHostsCommand[i].replaceAll("%", "%%");
+          knownHostsCommand[i] = knownHostsCommand[i].replaceAll('"', '\\"');
+          knownHostsCommand[i] = '"' + knownHostsCommand[i] + '"';
+        }
+        const knownHostsCommand_str = knownHostsCommand.join(" ");
+        cmd = cmd.concat(["-o", `KnownHostsCommand=${knownHostsCommand_str}`]);
+      }
+
+      cmd = cmd.concat([sshDestination]);
+
+      let shescape: undefined | Shescape = undefined;
+      let shell: undefined | string = undefined;
+      if (process.env.SHELL !== undefined) {
+        try {
+          shescape = new Shescape({
+            flagProtection: false,
+            shell: process.env.SHELL,
+          });
+          shell = process.env.SHELL;
+        } catch {
+          // shescape will stay undefined
+        }
+      }
+      if (shescape === undefined) {
+        shescape = new Shescape({
+          flagProtection: false,
+          shell: "/bin/sh",
+        });
+        shell = "/bin/sh";
+      }
+      const shell_cmd = shescape.quoteAll(cmd).join(" ");
+      if (!options.quiet) {
+        console.log(`Executing (${shell} style output): ${shell_cmd}`);
+      }
+
+      execvp(cmd[0], cmd.slice(1));
+    });
+}

src/lib/ssh.ts

@@ -14,6 +14,7 @@ import {
 } from "../helpers/errors.ts";
 import { getApiUrl } from "../helpers/urls.ts";
 import { isFeatureEnabled } from "./posthog.ts";
+import { registerSsh } from "./ssh.ts";
 
 type VMInstance = {
   id: string;
@@ -34,6 +35,8 @@ export async function registerVM(program: Command) {
     .aliases(["v", "vms"])
     .description("Manage virtual machines");
 
+  registerSsh(vm);
+
   vm.command("list")
     .description("List all virtual machines")
     .action(async () => {