From e44ee76e8d68b47d04c20cc828aa7ca32578814c Mon Sep 17 00:00:00 2001
From: Michael Thiesen <michael@understorytech.com>
Date: Thu, 17 Oct 2024 23:38:51 -0700
Subject: [PATCH] Enhanced DNS handling and added shell command execution
 support

---
 biome.json   |  3 ++-
 src/dns.ts   | 10 ++++++-
 src/index.ts |  3 ++-
 src/ssh.ts   | 75 +++++++++++++++++++++++++++++++++++++++++++++++++---
 4 files changed, 84 insertions(+), 7 deletions(-)

diff --git a/biome.json b/biome.json
index 61e7cf7..8f1c6fe 100644
--- a/biome.json
+++ b/biome.json
@@ -11,7 +11,8 @@
       },
       "suspicious": {
         "noExplicitAny": "off",
-        "noShadowRestrictedNames": "off"
+        "noShadowRestrictedNames": "off",
+        "noControlCharactersInRegex": "off"
       }
     }
   },
diff --git a/src/dns.ts b/src/dns.ts
index f62853c..e62a8e3 100644
--- a/src/dns.ts
+++ b/src/dns.ts
@@ -9,7 +9,15 @@ const ALT_DNS_SERVERS: [string, string] = ["1.1.1.1", "1.0.0.1"];
 export async function toggleDns(ssh: SNodeSSH) {
   if (timeoutId) clearTimeout(timeoutId);
 
-  const originalDnsServers = parseDnsServers(await ssh.execSafe("show"));
+  await ssh.shellSafe("echo $");
+  await ssh.shellSafe("echo $SHELL");
+  await ssh.shellSafe("echo $TERM");
+  await ssh.shellSafe("echo $-");
+  await ssh.shellSafe("shopt login_shell");
+
+  const originalDnsServers = parseDnsServers(
+    await ssh.shellSafe("show dns forwarding nameservers"),
+  );
   console.log(`Found existing DNS servers: ${originalDnsServers}`);
 
   console.log("Setting alternative DNS servers...");
diff --git a/src/index.ts b/src/index.ts
index 1146c25..15cef39 100644
--- a/src/index.ts
+++ b/src/index.ts
@@ -14,6 +14,7 @@ app.get("/toggle-dns", async (c) => {
     await toggleDns(ssh);
     return c.text("DNS servers toggled");
   } catch (err: any) {
+    console.warn(err.message);
     return c.text(`Error: ${err.message}`);
   } finally {
     if (ssh) ssh.dispose();
@@ -21,5 +22,5 @@ app.get("/toggle-dns", async (c) => {
 });
 
 serve({ fetch: app.fetch, port: 5888 }, (info) => {
-  console.log(`\nListening on http://localhost:${info.port}`);
+  console.log(`Listening on http://localhost:${info.port}`);
 });
diff --git a/src/ssh.ts b/src/ssh.ts
index 26029e0..4fa0791 100644
--- a/src/ssh.ts
+++ b/src/ssh.ts
@@ -3,6 +3,7 @@ import env from "./env";
 
 export type SNodeSSH = NodeSSH & {
   execSafe: (command: string) => Promise<string>;
+  shellSafe: (command: string) => Promise<string>;
 };
 
 export async function getSshConnection(): Promise<SNodeSSH> {
@@ -16,9 +17,17 @@ export async function getSshConnection(): Promise<SNodeSSH> {
   ssh.execSafe = async (command: string) => {
     console.log(`>>> ${command}`);
     const result = await ssh.execCommand(command);
-    if (result.stderr) console.warn(prefixCaret(result.stderr));
+    if (result.stderr) console.warn(prefixBar(result.stderr));
     if (result.stderr) throw new Error(`Error executing command: ${result.stderr}`);
-    if (result.stdout) console.log(prefixCaret(result.stdout));
+    if (result.stdout) console.log(prefixBar(result.stdout));
+    return result.stdout;
+  };
+  ssh.shellSafe = async (command: string) => {
+    console.log(`>>> ${command}`);
+    const result = await withShell(ssh, command);
+    if (result.stderr) console.warn(prefixBar(result.stderr));
+    if (result.stderr) throw new Error(`Error executing command: ${result.stderr}`);
+    if (result.stdout) console.log(prefixBar(result.stdout));
     return result.stdout;
   };
 
@@ -30,6 +39,64 @@ export async function getSshConnection(): Promise<SNodeSSH> {
   return ssh;
 }
 
-function prefixCaret(text: string): string {
-  return text.replace(/^/gm, "> ");
+async function withShell(
+  ssh: SNodeSSH,
+  command: string,
+): Promise<{ stdout: string; stderr: string }> {
+  let stdout = "";
+  let stderr = "";
+
+  await ssh.withShell(
+    async (channel) => {
+      return new Promise<void>((resolve, reject) => {
+        channel.on("data", (data: Buffer) => {
+          stdout += data.toString();
+        });
+
+        channel.stderr.on("data", (data: Buffer) => {
+          stderr += data.toString();
+        });
+
+        channel.on("close", () => {
+          resolve();
+        });
+
+        channel.on("error", (err: any) => {
+          reject(err);
+        });
+
+        channel.write(`${command}\nexit\n`);
+      });
+    },
+    {
+      modes: {
+        ECHO: 1, // Enable echoing
+        ICANON: 1, // Enable canonical input
+        ISIG: 1, // Enable signals like INTR, QUIT, SUSP
+        IEXTEN: 1, // Enable extensions
+        OPOST: 1, // Enable output processing
+        ONLCR: 1, // Translate newline to CR-NL
+        TTY_OP_ISPEED: 9600, // Input baud rate
+        TTY_OP_OSPEED: 9600, // Output baud rate
+      },
+    },
+  );
+
+  return { stdout: extractCommandOutput(stdout), stderr };
+}
+
+const shellOutputRegex = /(?<=@ubnt:~\$.+\n)([\s\S]+?)(?=\n.+@ubnt:~\$ exit)/gims;
+function extractCommandOutput(fullOutput: string): string {
+  const matches = fullOutput.match(shellOutputRegex);
+  if (!matches) return "";
+  const lines = matches.map(cleanLine);
+  return lines.join("\n");
+}
+
+function cleanLine(line: string): string {
+  return line.replace(/\x1B\[[0-9;]*[a-zA-Z]|\x1B\[[0-9;]*m/g, "").trim();
+}
+
+function prefixBar(text: string): string {
+  return text.replace(/^/gm, "| ");
 }