create abstract ShellSession, extend Node and Local from that

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

src/main/shell-session/index.ts

+export * from "./shell-session";
+export * from "./local-shell-session";
+export * from "./node-shell-session";

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

 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;
+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 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();
+  protected async rawOpen() {
     const env = await this.getCachedShellEnv();
-    const shell = env.PTYSHELL;
+    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: this.cwd() || env.HOME,
+      cwd,
       env,
       name: "xterm-256color",
       rows: 30,
@@ -60,18 +37,24 @@ export class ShellSession extends EventEmitter {
     this.closeWebsocketOnProcessExit();
     this.exitProcessOnWebsocketClose();
 
-    appEventBus.emit({name: "shell", action: "open"});
+    appEventBus.emit({ name: this.EventName, action: "open" });
   }
 
-  protected cwd(): string {
-    if(!this.preferences || !this.preferences.terminalCWD || this.preferences.terminalCWD === "") {
-      return null;
-    }
+  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.preferences.terminalCWD;
+    return this.rawOpen();
   }
 
-  protected async getShellArgs(shell: string): Promise<Array<string>> {
+  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"}`];
@@ -86,75 +69,9 @@ export class ShellSession extends EventEmitter {
     }
   }
 
-  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);
-    }));
+    this.shellProcess.onData(this.sendResponse);
   }
 
   protected pipeStdin() {
@@ -169,9 +86,9 @@ export class ShellSession extends EventEmitter {
           this.shellProcess.write(message);
           break;
         case "4":
-          const resizeMsgObj = JSON.parse(message);
+          const { Width, Height } = JSON.parse(message);
 
-          this.shellProcess.resize(resizeMsgObj["Width"], resizeMsgObj["Height"]);
+          this.shellProcess.resize(Width, Height);
           break;
         case "9":
           this.emit("newToken", message);
@@ -188,12 +105,12 @@ export class ShellSession extends EventEmitter {
   protected closeWebsocketOnProcessExit() {
     this.shellProcess.onExit(({ exitCode }) => {
       this.running = false;
-      let timeout = 0;
+      const timeout = exitCode > 0 ? 15 * 1000 : 0;
 
       if (exitCode > 0) {
         this.sendResponse("Terminal will auto-close in 15 seconds ...");
-        timeout = 15*1000;
       }
+
       setTimeout(() => {
         this.exit();
       }, timeout);
@@ -207,7 +124,7 @@ export class ShellSession extends EventEmitter {
   }
 
   protected killShellProcess(){
-    if(this.running) {
+    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 {
@@ -221,6 +138,7 @@ export class ShellSession extends EventEmitter {
     }
   }
 
+  @autobind()
   protected sendResponse(msg: string) {
     this.websocket.send(`1${Buffer.from(msg).toString("base64")}`);
   }

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

 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 podId: string;
+  protected readonly EventName = "node-shell";
+  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, {
-      cols: 80,
-      cwd: this.cwd() || shellEnv["HOME"],
-      env: shellEnv,
-      name: "xterm-256color",
-      rows: 30,
-    });
-    this.running = true;
-    this.pipeStdout();
-    this.pipeStdin();
-    this.closeWebsocketOnProcessExit();
-    this.exitProcessOnWebsocketClose();
+  protected async getShell(): Promise<string> {
+    return this.kubectl.getPath();
+  }
 
-    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

+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;
+  }
+}