From d8dbc351bb53ec9d6042e8449f60ea1d3c7eabdf Mon Sep 17 00:00:00 2001
From: Elijah <elijah@computernewb.com>
Date: Fri, 8 Dec 2023 20:02:44 -0500
Subject: [PATCH] Add File and Discord logging

---
 config.example.json |  2 ++
 package.json        |  1 +
 src/IConfig.ts      |  2 ++
 src/VM.ts           | 23 +++++++++++++----------
 src/index.ts        | 29 ++++++++++++++++++++++++++++-
 5 files changed, 46 insertions(+), 11 deletions(-)

diff --git a/config.example.json b/config.example.json
index 1af74b6..077da9e 100644
--- a/config.example.json
+++ b/config.example.json
@@ -4,6 +4,8 @@
     "MaxFileSize": 104857600,
     "BlockedMD5": [],
     "RateLimit": 10,
+    "LogDir": "/var/log/agent",
+    "DiscordWebhook": "https://discordapp.com/api/webhooks/1234567890/abcdefghijklmnopqrstuvwxyz",
     "VMs": [
         {
             "ID": "vm1",
diff --git a/package.json b/package.json
index 77f0e40..2fea0fa 100644
--- a/package.json
+++ b/package.json
@@ -17,6 +17,7 @@
   "dependencies": {
     "@types/md5": "^2.3.5",
     "async-mutex": "^0.4.0",
+    "discord.js": "^14.14.1",
     "fastify": "^4.24.3",
     "md5": "^2.3.0",
     "msgpack-lite": "^0.1.26"
diff --git a/src/IConfig.ts b/src/IConfig.ts
index c7985dc..f740c78 100644
--- a/src/IConfig.ts
+++ b/src/IConfig.ts
@@ -4,6 +4,8 @@ export default interface IConfig {
     MaxFileSize : number;
     BlockedMD5: string[];
     RateLimit : number;
+    LogDir? : string;
+    DiscordWebhook? : string;
     VMs : {
         ID : string;
         SocketPath : string;
diff --git a/src/VM.ts b/src/VM.ts
index 77e8579..74b46b0 100644
--- a/src/VM.ts
+++ b/src/VM.ts
@@ -27,20 +27,23 @@ export default class VM {
                 FileData: data
             };
             var payload = msgpack.encode(msg);
-            var header = Buffer.alloc(4);
-            header.writeUInt32LE(payload.length);
-            // Make sure we dont write two messages at the same time
-            var release = await this.#writeLock.acquire();
-            await this.#writeData(header);
-            // shit gets fucked if i dont do this
-            await sleep(100);
-            await this.#writeData(payload);
-            // All done
-            release();
+            await this.#sendMessage(payload);
             res();
         });
     }
 
+    #sendMessage(data : Buffer) : Promise<void> {
+        return new Promise(async (res, rej) => {
+            var header = Buffer.alloc(4);
+            header.writeUInt32LE(data.length);
+            await this.#writeLock.runExclusive(async () => {
+                await this.#writeData(header);
+                await this.#writeData(data);
+                res();
+            });
+        });
+    }
+
     #writeData(data : Buffer) : Promise<void> {
         return new Promise((resolve, reject) => {
             this.#socket.write(data, (err) => {
diff --git a/src/index.ts b/src/index.ts
index 97d706c..f1f8a0d 100644
--- a/src/index.ts
+++ b/src/index.ts
@@ -5,6 +5,8 @@ import VM from './VM.js';
 import md5 from 'md5';
 import log from './log.js';
 import RateLimit from './Ratelimit.js';
+import * as fs from 'fs';
+import { EmbedBuilder, WebhookClient } from 'discord.js';
 
 log("INFO", "CollabVM Agent Server Starting up...");
 // Load the config file
@@ -17,6 +19,13 @@ try {
     process.exit(1);
 }
 
+var discordenabled : boolean = false;
+var discord : WebhookClient | null = null;
+if (config.DiscordWebhook) {
+    discordenabled = true;
+    discord = new WebhookClient({ url: config.DiscordWebhook });
+}
+
 var VMs = new Map<string, VM>();
 config.VMs.forEach((v) => {
     VMs.set(v.ID, new VM(v.SocketPath));
@@ -88,6 +97,14 @@ app.put("/:vm/:filename", async (req, res) => {
     }
     await VMs.get(vm)!.UploadFile(filename, filedata);
     log("INFO", `${vm}: ${req.ip} uploaded "${filename}" with MD5 ${hash}`);
+    if (config.LogDir)
+        await appendFile(`${config.LogDir}/${vm}.log`, `${req.ip} uploaded "${filename}" with MD5 ${hash}\n`);
+    if (discordenabled) {
+        var embed = new EmbedBuilder()
+        .addFields({"name": "VM", "value": vm, "inline": true}, {"name": "Filename", "value": filename, "inline": true}, {"name": "IP", "value": req.ip, "inline": true}, {"name": "MD5", "value": hash, "inline": true})
+        .setTitle("File Uploaded");
+        discord!.send({ embeds: [embed] });
+    }
     return { success: true, result: "File uploaded" };
 });
 
@@ -95,4 +112,14 @@ log("INFO", "Starting HTTP server on port " + config.ListenPort);
 app.listen({
     port: config.ListenPort,
     host: "127.0.0.1"
-});
\ No newline at end of file
+});
+
+function appendFile(path : string, data : string | Uint8Array) {
+    return new Promise<void>((res, rej) => {
+        fs.appendFile(path, data, (err) => {
+            if (err) {
+                rej(err);
+            } else res();
+        });
+    });
+}
\ No newline at end of file