feat: better command system

package.json

@@ -47,7 +47,8 @@
         "deploy": "node scripts/deploy-commands.js",
         "gen:schema": "node scripts/generate-config-schema.js",
         "clean": "rm -frv build tsconfig.tsbuildinfo; make clean",
-        "test": "vitest"
+        "test": "vitest",
+        "shell": "bun run src/main/typescript/shell.ts"
     },
     "_moduleAliases": {
         "@sudobot": "build/out",

src/framework/typescript/polyfills/FileSystem.ts

@@ -19,7 +19,7 @@
 
 import { existsSync } from "fs";
 import { readFile, writeFile } from "fs/promises";
-import type { Stringable } from "../types/Stringable";
+import type { StringLike } from "../types/StringLike";
 
 type ReadFileContentOptions<T extends boolean> = {
     json?: T;
@@ -81,7 +81,7 @@ export default class FileSystem {
      */
     public static async writeFileContents<J extends boolean = false>(
         path: string,
-        contents: J extends true ? object : Stringable,
+        contents: J extends true ? object : StringLike,
         json: J = false as J
     ): Promise<void> {
         if (process.versions.bun) {

src/framework/typescript/types/StringLike.ts

@@ -0,0 +1,20 @@
+/*
+ * This file is part of SudoBot.
+ *
+ * Copyright (C) 2021-2024 OSN Developers.
+ *
+ * SudoBot is free software; you can redistribute it and/or modify it
+ * under the terms of the GNU Affero General Public License as published by
+ * the Free Software Foundation, either version 3 of the License, or
+ * (at your option) any later version.
+ *
+ * SudoBot is distributed in the hope that it will be useful, but
+ * WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+ * GNU Affero General Public License for more details.
+ *
+ * You should have received a copy of the GNU Affero General Public License
+ * along with SudoBot. If not, see <https://www.gnu.org/licenses/>.
+ */
+
+export type StringLike = { toString: () => string } | string;

src/main/typescript/api/controllers/ShellCommandController.ts

@@ -48,7 +48,7 @@ class ShellCommandController extends Controller {
         const { command, args } = request.parsedBody as unknown as z.infer<
             typeof CommandPostSchema
         >;
-        const { output, error, code } = await this.shellService.simpleExecute(command, args);
+        const { output, error, code } = await this.shellService.simpleExecute(command, { args });
 
         return new Response({
             status: error ? 400 : 200,

src/main/typescript/services/ShellService.ts

@@ -1,12 +1,14 @@
 import Application from "@framework/app/Application";
 import { Name } from "@framework/services/Name";
 import { Service } from "@framework/services/Service";
+import { StringLike } from "@framework/types/StringLike";
 import { f } from "@framework/utils/string";
 import { env } from "@main/env/env";
+import ShellCommand from "@main/shell/core/ShellCommand";
 import { _meta as meta, version } from "@root/package.json";
 import chalk from "chalk";
 import { spawn } from "child_process";
-import { Awaitable } from "discord.js";
+import { Collection } from "discord.js";
 import { lstat, readdir } from "fs/promises";
 import { createServer } from "http";
 import path from "path";
@@ -17,6 +19,7 @@ import { WebSocket } from "ws";
 class ShellService extends Service {
     public readonly wss: InstanceType<typeof WebSocket.Server>;
     public readonly server = createServer();
+    private readonly commands = new Collection<string, ShellCommand>();
 
     public constructor(application: Application) {
         super(application);
@@ -26,7 +29,20 @@ class ShellService extends Service {
         });
     }
 
-    public override boot(): Awaitable<void> {
+    public override async boot(): Promise<void> {
+        const classes = await this.application.classLoader.loadClassesFromDirectory(
+            path.resolve(__dirname, "../shell/commands")
+        );
+
+        for (const ShellCommandClass of classes) {
+            const command = new ShellCommandClass(this.application) as ShellCommand;
+            this.commands.set(command.name, command);
+
+            for (const alias of command.aliases) {
+                this.commands.set(alias, command);
+            }
+        }
+
         this.wss.on("connection", ws => {
             this.application.logger.debug("Connection established");
 
@@ -147,8 +163,8 @@ class ShellService extends Service {
         ws.on("message", onMessage);
     }
 
-    public async simpleExecute(command: string, args: string[]) {
-        this.application.logger.event("Executing shell command: ", command, args);
+    public async simpleExecute(command: string, options?: ShellExecuteOptions) {
+        this.application.logger.event("Executing shell command: ", command, options);
 
         switch (command) {
             case "version":
@@ -160,12 +176,12 @@ class ShellService extends Service {
                 };
 
             case "cd":
-                if (!args[0]) {
+                if (!options?.args?.[0]) {
                     return { output: null, error: "cd: missing operand", code: 2 };
                 }
 
                 try {
-                    process.chdir(args[0]);
+                    process.chdir(options.args[0]);
                     return { output: null, error: null, code: 0 };
                 } catch (error) {
                     return {
@@ -214,8 +230,52 @@ class ShellService extends Service {
             }
         }
 
-        return { output: null, error: `${command}: command not found`, code: 127 };
+        let outputBuffer = "";
+
+        const shellCommand = this.commands.get(command);
+
+        if (!shellCommand) {
+            return { output: null, error: `${command}: command not found`, code: 127 };
+        }
+
+        const print = (...args: StringLike[]) => {
+            outputBuffer += args.join(" ");
+        };
+
+        const println = (...args: StringLike[]) => {
+            outputBuffer += args.join(" ") + "\n";
+        };
+
+        const ret = await shellCommand.execute({
+            elevatedPrivileges: options?.elevatedPrivileges ?? false,
+            args: options?.args ?? [],
+            print,
+            println
+        });
+
+        if (typeof ret === "string") {
+            return { output: ret, error: null, code: 0 };
+        }
+
+        if (
+            typeof ret === "object" &&
+            ret &&
+            ("output" in ret || "error" in ret || "code" in ret)
+        ) {
+            return {
+                output: (ret as Record<string, string>).output ?? null,
+                error: (ret as Record<string, string>).error ?? null,
+                code: (ret as Record<string, number>).code ?? 0
+            };
+        }
+
+        return { output: outputBuffer === "" ? null : outputBuffer, error: null, code: 0 };
     }
 }
 
+type ShellExecuteOptions = {
+    args?: string[];
+    elevatedPrivileges?: boolean;
+};
+
 export default ShellService;

src/main/typescript/shell.ts

@@ -0,0 +1,101 @@
+import Shell from "@main/shell/core/Shell";
+import axios, { AxiosError } from "axios";
+import chalk from "chalk";
+
+async function main() {
+    if (!process.env.SYSTEM_SHELL_KEY) {
+        throw new Error("Environment variable SYSTEM_SHELL_KEY is not defined");
+    }
+
+    const shell = new Shell();
+
+    await shell.awaitReady();
+
+    for await (const input of shell) {
+        if (input.startsWith("$")) {
+            if (input.slice(1).trim().length === 0) {
+                console.error(`${chalk.red.bold("sbsh:")} expected raw shell command after "$"`);
+                continue;
+            }
+
+            try {
+                shell.executeCommand(input.slice(1));
+
+                await new Promise<void>(resolve => {
+                    const handler = (message: MessageEvent) => {
+                        const { type, payload } = JSON.parse(message.data.toString("utf-8"));
+
+                        if (type === "exit" || type === "signal") {
+                            const code = +payload;
+                            shell.ws.removeEventListener("message", handler);
+                            shell.setExitCode(isNaN(code) ? 128 : code);
+                            resolve();
+                        }
+                    };
+
+                    shell.ws.addEventListener("message", handler);
+                });
+            } catch (error) {
+                if (error instanceof Error && error.name === "AbortError") {
+                    console.error(
+                        `${chalk.red.bold("sbsh:")} ${chalk.white.bold("SIGINT")} received, aborting command execution`
+                    );
+                    continue;
+                }
+
+                if (error instanceof AxiosError) {
+                    console.error(
+                        `${chalk.red.bold("sbsh:")} ${error.response?.data.shortError ?? error.response?.data.error ?? error?.message}`
+                    );
+
+                    shell.setExitCode(error.response?.data.code ?? 1);
+                    continue;
+                }
+
+                console.error(error);
+            }
+
+            continue;
+        }
+
+        const [command, ...args] = input.split(/\s+/);
+
+        if (shell.handleBuiltInCommands(command, args)) {
+            continue;
+        }
+
+        try {
+            const { data } = await axios.post(
+                `${process.env.SYSTEM_API_URL}/shell/command`,
+                {
+                    command,
+                    args
+                },
+                {
+                    headers: {
+                        Authorization: `Bearer ${process.env.SYSTEM_SHELL_KEY}`
+                    }
+                }
+            );
+
+            if (data.output !== null) {
+                console.log(data.output);
+            }
+
+            shell.setExitCode(0);
+        } catch (error) {
+            if (error instanceof AxiosError) {
+                console.error(
+                    `${chalk.red.bold("sbsh:")} ${error.response?.data.error ?? error?.message}`
+                );
+
+                shell.setExitCode(error.response?.data.code ?? 1);
+                continue;
+            }
+
+            console.error(error);
+        }
+    }
+}
+
+main().then();

src/main/typescript/shell/commands/AdminTestShellCommand.ts

@@ -0,0 +1,16 @@
+import ShellCommand from "@main/shell/core/ShellCommand";
+import type { ShellCommandContext } from "@main/shell/core/ShellCommandContext";
+
+class AdminTestShellCommand extends ShellCommand {
+    public override readonly name: string = "admintest";
+
+    public override async execute(context: ShellCommandContext): Promise<void> {
+        if (context.elevatedPrivileges) {
+            context.println("Looks good! You have elevated privileges.");
+        } else {
+            context.println("You do not have elevated privileges.");
+        }
+    }
+}
+
+export default AdminTestShellCommand;

src/main/typescript/shell/commands/RebootCommand.ts

@@ -0,0 +1,21 @@
+import ShellCommand from "@main/shell/core/ShellCommand";
+import type { ShellCommandContext } from "@main/shell/core/ShellCommandContext";
+
+class RebootCommand extends ShellCommand {
+    public override readonly name: string = "reboot";
+    public override readonly aliases: string[] = ["restart"];
+
+    public override async execute(context: ShellCommandContext): Promise<unknown> {
+        if (!context.elevatedPrivileges) {
+            return {
+                code: 1,
+                error: "reboot: Operation not permitted\nreboot: You may need elevated privileges to perform this action."
+            };
+        }
+
+        context.println("Rebooting in 5 seconds. You will lose connection to the shell.");
+        setTimeout(() => this.application.service("startupManager").requestRestart(), 5000);
+    }
+}
+
+export default RebootCommand;

src/main/typescript/shell/commands/SudoShellCommand.ts

@@ -0,0 +1,21 @@
+import ShellCommand from "@main/shell/core/ShellCommand";
+import type { ShellCommandContext } from "@main/shell/core/ShellCommandContext";
+
+class SudoShellCommand extends ShellCommand {
+    public override readonly name: string = "sudo";
+
+    public override async execute(context: ShellCommandContext): Promise<unknown> {
+        if (!context.args[0]) {
+            return { code: 1, error: "Usage: sudo <command>" };
+        }
+
+        const service = this.application.service("shellService");
+
+        return await service.simpleExecute(context.args[0], {
+            elevatedPrivileges: true,
+            args: context.args.slice(1)
+        });
+    }
+}
+
+export default SudoShellCommand;