create abstract ShellSession, extend Node and Local from that

Signed-off-by: Sebastian Malton <sebastian@malton.name>

src/main/shell-session.ts

@@ -1,227 +0,0 @@
-import * as pty from "node-pty";
-import * as WebSocket from "ws";
-import { EventEmitter } from "events";
-import path from "path";
-import shellEnv from "shell-env";
-import { app } from "electron";
-import { Kubectl } from "./kubectl";
-import { Cluster } from "./cluster";
-import { ClusterPreferences } from "../common/cluster-store";
-import { helmCli } from "./helm/helm-cli";
-import { isWindows } from "../common/vars";
-import { appEventBus } from "../common/event-bus";
-import { userStore } from "../common/user-store";
-
-export class ShellSession extends EventEmitter {
-  static shellEnvs: Map<string, any> = new Map();
-
-  protected websocket: WebSocket;
-  protected shellProcess: pty.IPty;
-  protected kubeconfigPath: string;
-  protected nodeShellPod: string;
-  protected kubectl: Kubectl;
-  protected kubectlBinDir: string;
-  protected kubectlPathDir: string;
-  protected helmBinDir: string;
-  protected preferences: ClusterPreferences;
-  protected running = false;
-  protected clusterId: string;
-
-  constructor(socket: WebSocket, cluster: Cluster) {
-    super();
-    this.websocket = socket;
-    this.kubeconfigPath =  cluster.getProxyKubeconfigPath();
-    this.kubectl = new Kubectl(cluster.version);
-    this.preferences = cluster.preferences || {};
-    this.clusterId = cluster.id;
-  }
-
-  public async open() {
-    this.kubectlBinDir = await this.kubectl.binDir();
-    const pathFromPreferences = userStore.preferences.kubectlBinariesPath || this.kubectl.getBundledPath();
-
-    this.kubectlPathDir = userStore.preferences.downloadKubectlBinaries ? this.kubectlBinDir : path.dirname(pathFromPreferences);
-    this.helmBinDir = helmCli.getBinaryDir();
-    const env = await this.getCachedShellEnv();
-    const shell = env.PTYSHELL;
-    const args = await this.getShellArgs(shell);
-
-    this.shellProcess = pty.spawn(shell, args, {
-      cols: 80,
-      cwd: this.cwd() || env.HOME,
-      env,
-      name: "xterm-256color",
-      rows: 30,
-    });
-    this.running = true;
-
-    this.pipeStdout();
-    this.pipeStdin();
-    this.closeWebsocketOnProcessExit();
-    this.exitProcessOnWebsocketClose();
-
-    appEventBus.emit({name: "shell", action: "open"});
-  }
-
-  protected cwd(): string {
-    if(!this.preferences || !this.preferences.terminalCWD || this.preferences.terminalCWD === "") {
-      return null;
-    }
-
-    return this.preferences.terminalCWD;
-  }
-
-  protected async getShellArgs(shell: string): Promise<Array<string>> {
-    switch(path.basename(shell)) {
-      case "powershell.exe":
-        return ["-NoExit", "-command", `& {Set-Location $Env:USERPROFILE; $Env:PATH="${this.helmBinDir};${this.kubectlPathDir};$Env:PATH"}`];
-      case "bash":
-        return ["--init-file", path.join(this.kubectlBinDir, ".bash_set_path")];
-      case "fish":
-        return ["--login", "--init-command", `export PATH="${this.helmBinDir}:${this.kubectlPathDir}:$PATH"; export KUBECONFIG="${this.kubeconfigPath}"`];
-      case "zsh":
-        return ["--login"];
-      default:
-        return [];
-    }
-  }
-
-  protected async getCachedShellEnv() {
-    let env = ShellSession.shellEnvs.get(this.clusterId);
-
-    if (!env) {
-      env = await this.getShellEnv();
-      ShellSession.shellEnvs.set(this.clusterId, env);
-    } else {
-      // refresh env in the background
-      this.getShellEnv().then((shellEnv: any) => {
-        ShellSession.shellEnvs.set(this.clusterId, shellEnv);
-      });
-    }
-
-    return env;
-  }
-
-  protected async getShellEnv() {
-    const env = JSON.parse(JSON.stringify(await shellEnv()));
-    const pathStr = [this.kubectlBinDir, this.helmBinDir, process.env.PATH].join(path.delimiter);
-
-    if(isWindows) {
-      env["SystemRoot"] = process.env.SystemRoot;
-      env["PTYSHELL"] = process.env.SHELL || "powershell.exe";
-      env["PATH"] = pathStr;
-      env["LENS_SESSION"] = "true";
-      const lensWslEnv = "KUBECONFIG/up:LENS_SESSION/u";
-
-      if (process.env.WSLENV != undefined) {
-        env["WSLENV"] = `${process.env["WSLENV"]}:${lensWslEnv}`;
-      } else {
-        env["WSLENV"] = lensWslEnv;
-      }
-    } else if(typeof(process.env.SHELL) != "undefined") {
-      env["PTYSHELL"] = process.env.SHELL;
-      env["PATH"] = pathStr;
-    } else {
-      env["PTYSHELL"] = ""; // blank runs the system default shell
-    }
-
-    if(path.basename(env["PTYSHELL"]) === "zsh") {
-      env["OLD_ZDOTDIR"] = env.ZDOTDIR || env.HOME;
-      env["ZDOTDIR"] = this.kubectlBinDir;
-      env["DISABLE_AUTO_UPDATE"] = "true";
-    }
-
-    env["PTYPID"] = process.pid.toString();
-    env["KUBECONFIG"] = this.kubeconfigPath;
-    env["TERM_PROGRAM"] = app.getName();
-    env["TERM_PROGRAM_VERSION"] = app.getVersion();
-
-    if (this.preferences.httpsProxy) {
-      env["HTTPS_PROXY"] = this.preferences.httpsProxy;
-    }
-    const no_proxy = ["localhost", "127.0.0.1", env["NO_PROXY"]];
-
-    env["NO_PROXY"] = no_proxy.filter(address => !!address).join();
-
-    if (env.DEBUG) { // do not pass debug option to bash
-      delete env["DEBUG"];
-    }
-
-    return(env);
-  }
-
-  protected pipeStdout() {
-    // send shell output to websocket
-    this.shellProcess.onData(((data: string) => {
-      this.sendResponse(data);
-    }));
-  }
-
-  protected pipeStdin() {
-    // write websocket messages to shellProcess
-    this.websocket.on("message", (data: string) => {
-      if (!this.running) { return; }
-
-      const message = Buffer.from(data.slice(1, data.length), "base64").toString();
-
-      switch (data[0]) {
-        case "0":
-          this.shellProcess.write(message);
-          break;
-        case "4":
-          const resizeMsgObj = JSON.parse(message);
-
-          this.shellProcess.resize(resizeMsgObj["Width"], resizeMsgObj["Height"]);
-          break;
-        case "9":
-          this.emit("newToken", message);
-          break;
-      }
-    });
-  }
-
-  protected exit(code = 1000) {
-    if (this.websocket.readyState == this.websocket.OPEN) this.websocket.close(code);
-    this.emit("exit");
-  }
-
-  protected closeWebsocketOnProcessExit() {
-    this.shellProcess.onExit(({ exitCode }) => {
-      this.running = false;
-      let timeout = 0;
-
-      if (exitCode > 0) {
-        this.sendResponse("Terminal will auto-close in 15 seconds ...");
-        timeout = 15*1000;
-      }
-      setTimeout(() => {
-        this.exit();
-      }, timeout);
-    });
-  }
-
-  protected exitProcessOnWebsocketClose() {
-    this.websocket.on("close", () => {
-      this.killShellProcess();
-    });
-  }
-
-  protected killShellProcess(){
-    if(this.running) {
-      // On Windows we need to kill the shell process by pid, since Lens won't respond after a while if using `this.shellProcess.kill()`
-      if (isWindows) {
-        try {
-          process.kill(this.shellProcess.pid);
-        } catch(e) {
-          return;
-        }
-      } else {
-        this.shellProcess.kill();
-      }
-    }
-  }
-
-  protected sendResponse(msg: string) {
-    this.websocket.send(`1${Buffer.from(msg).toString("base64")}`);
-  }
-}

src/main/shell-session/index.ts

@@ -0,0 +1,3 @@
+export * from "./shell-session";
+export * from "./local-shell-session";
+export * from "./node-shell-session";

src/main/shell-session/local-shell-session.ts

@@ -0,0 +1,145 @@
+import * as pty from "node-pty";
+import path from "path";
+import { helmCli } from "../helm/helm-cli";
+import { isWindows } from "../../common/vars";
+import { appEventBus } from "../../common/event-bus";
+import { userStore } from "../../common/user-store";
+import { autobind } from "../../common/utils";
+import { ShellSession } from "./shell-session";
+
+type EnvVarMap = Record<string, string>;
+
+export class LocalShellSession extends ShellSession {
+  protected readonly EventName: string = "shell";
+
+  protected kubectlBinDir: string;
+  protected kubectlPathDir: string;
+  protected helmBinDir: string;
+  protected running = false;
+
+  protected async rawOpen() {
+    const env = await this.getCachedShellEnv();
+    const shell = await this.getShell(env);
+    const args = await this.getShellArgs(shell);
+    const cwd = this.cwd(env);
+
+    this.shellProcess = pty.spawn(shell, args, {
+      cols: 80,
+      cwd,
+      env,
+      name: "xterm-256color",
+      rows: 30,
+    });
+    this.running = true;
+
+    this.pipeStdout();
+    this.pipeStdin();
+    this.closeWebsocketOnProcessExit();
+    this.exitProcessOnWebsocketClose();
+
+    appEventBus.emit({ name: this.EventName, action: "open" });
+  }
+
+  public async open() {
+    this.kubectlBinDir = await this.kubectl.binDir();
+    const pathFromPreferences = userStore.preferences.kubectlBinariesPath || this.kubectl.getBundledPath();
+
+    this.kubectlPathDir = userStore.preferences.downloadKubectlBinaries ? this.kubectlBinDir : path.dirname(pathFromPreferences);
+    this.helmBinDir = helmCli.getBinaryDir();
+
+    return this.rawOpen();
+  }
+
+  protected async getShell(env: EnvVarMap): Promise<string> {
+    return env.PTYSHELL;
+  }
+
+  protected async getShellArgs(shell: string): Promise<string[]> {
+    switch(path.basename(shell)) {
+      case "powershell.exe":
+        return ["-NoExit", "-command", `& {Set-Location $Env:USERPROFILE; $Env:PATH="${this.helmBinDir};${this.kubectlPathDir};$Env:PATH"}`];
+      case "bash":
+        return ["--init-file", path.join(this.kubectlBinDir, ".bash_set_path")];
+      case "fish":
+        return ["--login", "--init-command", `export PATH="${this.helmBinDir}:${this.kubectlPathDir}:$PATH"; export KUBECONFIG="${this.kubeconfigPath}"`];
+      case "zsh":
+        return ["--login"];
+      default:
+        return [];
+    }
+  }
+
+  protected pipeStdout() {
+    // send shell output to websocket
+    this.shellProcess.onData(this.sendResponse);
+  }
+
+  protected pipeStdin() {
+    // write websocket messages to shellProcess
+    this.websocket.on("message", (data: string) => {
+      if (!this.running) { return; }
+
+      const message = Buffer.from(data.slice(1, data.length), "base64").toString();
+
+      switch (data[0]) {
+        case "0":
+          this.shellProcess.write(message);
+          break;
+        case "4":
+          const { Width, Height } = JSON.parse(message);
+
+          this.shellProcess.resize(Width, Height);
+          break;
+        case "9":
+          this.emit("newToken", message);
+          break;
+      }
+    });
+  }
+
+  protected exit(code = 1000) {
+    if (this.websocket.readyState == this.websocket.OPEN) this.websocket.close(code);
+    this.emit("exit");
+  }
+
+  protected closeWebsocketOnProcessExit() {
+    this.shellProcess.onExit(({ exitCode }) => {
+      this.running = false;
+      const timeout = exitCode > 0 ? 15 * 1000 : 0;
+
+      if (exitCode > 0) {
+        this.sendResponse("Terminal will auto-close in 15 seconds ...");
+      }
+
+      setTimeout(() => {
+        this.exit();
+      }, timeout);
+    });
+  }
+
+  protected exitProcessOnWebsocketClose() {
+    this.websocket.on("close", () => {
+      this.killShellProcess();
+    });
+  }
+
+  protected killShellProcess(){
+    if (this.running) {
+      // On Windows we need to kill the shell process by pid, since Lens won't respond after a while if using `this.shellProcess.kill()`
+      if (isWindows) {
+        try {
+          process.kill(this.shellProcess.pid);
+        } catch(e) {
+          return;
+        }
+      } else {
+        this.shellProcess.kill();
+      }
+    }
+  }
+
+  @autobind()
+  protected sendResponse(msg: string) {
+    this.websocket.send(`1${Buffer.from(msg).toString("base64")}`);
+  }
+}

src/main/shell-session/node-shell-session.ts

@@ -1,52 +1,38 @@
 import * as WebSocket from "ws";
-import * as pty from "node-pty";
 import { ShellSession } from "./shell-session";
 import { v4 as uuid } from "uuid";
 import * as k8s from "@kubernetes/client-node";
 import { KubeConfig } from "@kubernetes/client-node";
 import { Cluster } from "./cluster";
 import logger from "./logger";
-import { appEventBus } from "../common/event-bus";
+import { LocalShellSession } from "./local-shell-session";
 
 export class NodeShellSession extends ShellSession {
-  protected nodeName: string;
+  protected readonly EventName = "node-shell";
-  protected podId: string;
+  protected podId = `node-shell-${uuid()}`;
   protected kc: KubeConfig;
 
-  constructor(socket: WebSocket, cluster: Cluster, nodeName: string) {
+  constructor(socket: WebSocket, cluster: Cluster, protected nodeName: string) {
     super(socket, cluster);
-    this.nodeName = nodeName;
-    this.podId = `node-shell-${uuid()}`;
     this.kc = cluster.getProxyKubeconfig();
   }
 
   public async open() {
-    const shell = await this.kubectl.getPath();
-    let args = [];
-
     if (this.createNodeShellPod(this.podId, this.nodeName)) {
       await this.waitForRunningPod(this.podId).catch(() => {
         this.exit(1001);
       });
     }
-    args = ["exec", "-i", "-t", "-n", "kube-system", this.podId, "--", "sh", "-c", "((clear && bash) || (clear && ash) || (clear && sh))"];
 
-    const shellEnv = await this.getCachedShellEnv();
+    return this.rawOpen();
+  }
 
-    this.shellProcess = pty.spawn(shell, args, {
+  protected async getShell(): Promise<string> {
-      cols: 80,
+    return this.kubectl.getPath();
-      cwd: this.cwd() || shellEnv["HOME"],
+  }
-      env: shellEnv,
-      name: "xterm-256color",
-      rows: 30,
-    });
-    this.running = true;
-    this.pipeStdout();
-    this.pipeStdin();
-    this.closeWebsocketOnProcessExit();
-    this.exitProcessOnWebsocketClose();
 
-    appEventBus.emit({name: "node-shell", action: "open"});
+  protected async getShellArgs(): Promise<string[]> {
+    return ["exec", "-i", "-t", "-n", "kube-system", this.podId, "--", "sh", "-c", "((clear && bash) || (clear && ash) || (clear && sh))"];
   }
 
   protected exit(code = 1000) {
@@ -146,7 +132,7 @@ export async function openShell(socket: WebSocket, cluster: Cluster, nodeName?:
   if (nodeName) {
     shell = new NodeShellSession(socket, cluster, nodeName);
   } else {
-    shell = new ShellSession(socket, cluster);
+    shell = new LocalShellSession(socket, cluster);
   }
   shell.open();
 

src/main/shell-session/shell-session.ts

@@ -0,0 +1,95 @@
+import { ClusterId, ClusterPreferences } from "../../common/cluster-store";
+import { Cluster } from "../cluster";
+import * as pty from "node-pty";
+import * as WebSocket from "ws";
+import { EventEmitter } from "events";
+import { Kubectl } from "../kubectl";
+import { isWindows } from "../../common/vars";
+
+export type EnvVarMap = Record<string, string>;
+
+/**
+ * Joins all the non-empty elements of `parts` using `sep` between each element
+ * @param parts the potential elements for the new ENV var multi-element value
+ * @param sep The separator to join the elements together
+ */
+function joinEnvParts(parts: (string | null | undefined)[], sep: string): string {
+  return parts.filter(Boolean).join(sep);
+}
+
+export abstract class ShellSession extends EventEmitter {
+  protected abstract readonly EventName: string;
+  protected static readonly ShellEnvs = new Map<ClusterId, EnvVarMap>();
+  protected running = false;
+  protected shellProcess: pty.IPty;
+  protected readonly kubeconfigPath: string;
+  protected readonly kubectl: Kubectl;
+
+  protected get clusterId(): ClusterId {
+    return this.cluster.id;
+  }
+
+  protected get preferences(): ClusterPreferences {
+    return this.cluster.preferences || {};
+  }
+
+  protected cwd(env: EnvVarMap): string {
+    return this.preferences?.terminalCWD || env.HOME;
+  }
+
+  constructor(protected websocket: WebSocket, protected cluster: Cluster) {
+    super();
+    this.kubeconfigPath = cluster.getProxyKubeconfigPath();
+    this.kubectl = new Kubectl(cluster.version);
+  }
+
+  protected async getCachedShellEnv(): Promise<EnvVarMap> {
+    if (!ShellSession.ShellEnvs.has(this.clusterId)) {
+      ShellSession.ShellEnvs.set(this.clusterId, await this.getShellEnv());
+    } else {
+      // refresh env in the background
+      this.getShellEnv()
+        .then(shellEnv => {
+          ShellSession.ShellEnvs.set(this.clusterId, shellEnv);
+        });
+    }
+
+    return ShellSession.ShellEnvs.get(this.clusterId);
+  }
+
+  protected async getShellEnv(): Promise<EnvVarMap> {
+    const env = JSON.parse(JSON.stringify(await shellEnv()));
+    const pathStr = joinEnvParts([this.kubectlBinDir, this.helmBinDir, process.env.PATH], path.delimiter);
+
+    if (isWindows) {
+      env["SystemRoot"] = process.env.SystemRoot;
+      env["PTYSHELL"] = process.env.SHELL || "powershell.exe";
+      env["PATH"] = pathStr;
+      env["LENS_SESSION"] = "true";
+      env["WSLENV"] = joinEnvParts([env["WSLENV"], "KUBECONFIG/up:LENS_SESSION/u"], ":");
+    } else {
+      env["PTYSHELL"] = process.env.SHELL ?? ""; // blank runs the system default shell
+      env["PATH"] = pathStr;
+    }
+
+    if (path.basename(env["PTYSHELL"]) === "zsh") {
+      env["OLD_ZDOTDIR"] = env.ZDOTDIR || env.HOME;
+      env["ZDOTDIR"] = this.kubectlBinDir;
+      env["DISABLE_AUTO_UPDATE"] = "true";
+    }
+
+    env["PTYPID"] = process.pid.toString();
+    env["KUBECONFIG"] = this.kubeconfigPath;
+    env["TERM_PROGRAM"] = app.getName();
+    env["TERM_PROGRAM_VERSION"] = app.getVersion();
+
+    if (this.preferences.httpsProxy) {
+      env["HTTPS_PROXY"] = this.preferences.httpsProxy;
+    }
+
+    env["WSLENV"] = joinEnvParts(["localhost", "127.0.0.1", env["NO_PROXY"]], ",");
+    delete env["DEBUG"];
+
+    return env;
+  }
+}