From d1157318c09ce1401be40d0cee9800aa6a618d4d Mon Sep 17 00:00:00 2001
From: Skyler Grey <skyler3665@gmail.com>
Date: Wed, 8 Mar 2023 10:07:38 +0000
Subject: [PATCH] Provide the clamav socket, correctly upsert file hashes

---
 src/commands/settings/automod.ts |  6 +++---
 src/config/main.d.ts             |  1 +
 src/reflex/scanners.ts           | 22 +++++++++++++++-------
 src/utils/database.ts            | 24 ++++++++++++++++++++++--
 4 files changed, 41 insertions(+), 12 deletions(-)

diff --git a/src/commands/settings/automod.ts b/src/commands/settings/automod.ts
index db4a1c4..d8464f0 100644
--- a/src/commands/settings/automod.ts
+++ b/src/commands/settings/automod.ts
@@ -1090,18 +1090,18 @@ const callback = async (interaction: CommandInteraction): Promise<void> => {
                     break;
                 }
                 case "malware": {
-                    i.deferUpdate();
+                    await i.deferUpdate();
                     config.malware = !config.malware;
                     break;
                 }
                 case "images": {
-                    i.deferUpdate();
+                    await i.deferUpdate();
                     const next = await imageMenu(i, m, config.images);
                     config.images = next;
                     break;
                 }
                 case "clean": {
-                    i.deferUpdate();
+                    await i.deferUpdate();
                     const next = await cleanMenu(i, m, config.clean);
                     config.clean = next;
                     break;
diff --git a/src/config/main.d.ts b/src/config/main.d.ts
index 99c460d..419414a 100644
--- a/src/config/main.d.ts
+++ b/src/config/main.d.ts
@@ -19,6 +19,7 @@ declare const config: {
     };
     baseUrl: string;
     rapidApiKey: string;
+    clamavSocket: string;
 };
 
 export default config;
diff --git a/src/reflex/scanners.ts b/src/reflex/scanners.ts
index 49b203b..59bfa2a 100644
--- a/src/reflex/scanners.ts
+++ b/src/reflex/scanners.ts
@@ -11,6 +11,7 @@ import * as tf from "@tensorflow/tfjs-node";
 import EmojiEmbed from "../utils/generateEmojiEmbed.js";
 import getEmojiByName from "../utils/getEmojiByName.js";
 import { ActionRowBuilder, ButtonBuilder, ButtonStyle } from "discord.js";
+import config from "../config/main.js";
 
 interface NSFWSchema {
     nsfw: boolean;
@@ -22,7 +23,11 @@ interface MalwareSchema {
 }
 
 const nsfw_model = await nsfwjs.load();
-const clamscanner = await new ClamScan().init({});
+const clamscanner = await new ClamScan().init({
+    clamdscan: {
+        socket: config.clamavSocket
+    }
+});
 
 export async function testNSFW(link: string): Promise<NSFWSchema> {
     const [fileStream, hash] = await streamAttachment(link);
@@ -34,26 +39,29 @@ export async function testNSFW(link: string): Promise<NSFWSchema> {
     const predictions = (await nsfw_model.classify(image, 1))[0]!;
     image.dispose();
 
-    return { nsfw: predictions.className === "Hentai" || predictions.className === "Porn" };
+    const nsfw = predictions.className === "Hentai" || predictions.className === "Porn";
+    await client.database.scanCache.write(hash, "nsfw", nsfw);
+
+    return { nsfw };
 }
 
 export async function testMalware(link: string): Promise<MalwareSchema> {
     const [fileName, hash] = await saveAttachment(link);
     const alreadyHaveCheck = await client.database.scanCache.read(hash);
     if (alreadyHaveCheck?.malware) return { safe: alreadyHaveCheck.malware };
-    let safe;
+    let malware;
     try {
-        safe = !(await clamscanner.scanFile(fileName)).isInfected;
+        malware = (await clamscanner.scanFile(fileName)).isInfected;
     } catch (e) {
         return { safe: true };
     }
-    client.database.scanCache.write(hash, "malware", safe);
-    return { safe };
+    client.database.scanCache.write(hash, "malware", malware);
+    return { safe: !malware };
 }
 
 export async function testLink(link: string): Promise<{ safe: boolean; tags: string[] }> {
     const alreadyHaveCheck = await client.database.scanCache.read(link);
-    if (alreadyHaveCheck?.bad_link) return { safe: alreadyHaveCheck.bad_link, tags: alreadyHaveCheck.tags };
+    if (alreadyHaveCheck?.bad_link) return { safe: alreadyHaveCheck.bad_link, tags: alreadyHaveCheck.tags ?? [] };
     const scanned: { safe?: boolean; tags?: string[] } = await fetch("https://unscan.p.rapidapi.com/link", {
         method: "POST",
         headers: {
diff --git a/src/utils/database.ts b/src/utils/database.ts
index c5697f3..3fc2333 100644
--- a/src/utils/database.ts
+++ b/src/utils/database.ts
@@ -591,7 +591,7 @@ interface ScanCacheSchema {
     nsfw?: boolean;
     malware?: boolean;
     bad_link?: boolean;
-    tags: string[];
+    tags?: string[];
 }
 
 export class ScanCache {
@@ -608,7 +608,27 @@ export class ScanCache {
     async write(hash: string, type: "nsfw" | "malware" | "bad_link", data: boolean, tags?: string[]) {
         await this.scanCache.updateOne(
             { hash: hash },
-            { hash: hash, [type]: data, tags: tags ?? [], addedAt: new Date() },
+            {
+                $set: (() => {
+                    switch (type) {
+                        case "nsfw": {
+                            return { nsfw: data, addedAt: new Date() };
+                        }
+                        case "malware": {
+                            return { malware: data, addedAt: new Date() };
+                        }
+                        case "bad_link": {
+                            return { bad_link: data, tags: tags ?? [], addedAt: new Date() };
+                        }
+                        default: {
+                            throw new Error("Invalid type");
+                        }
+                    }
+                })()
+                // No you can't just do { [type]: data }, yes it's a typescript error, no I don't know how to fix it
+                // cleanly, yes it would be marginally more elegant, no it's not essential, yes I'd be happy to review
+                // PRs that did improve this snippet
+            },
             Object.assign({ upsert: true }, collectionOptions)
         );
     }