Improve messaging and checks in NSFW protection (#529)

* improve messages and event checks * tests * streamline extracting mxcs * requested changes * add comment
matrix-org · Oct 1, 2024 · 74d2c9b · 74d2c9b
1 parent 37a2168
commit 74d2c9b
Show file tree

Hide file tree

Showing 2 changed files with 46 additions and 7 deletions.
diff --git a/src/protections/NsfwProtection.ts b/src/protections/NsfwProtection.ts
@@ -45,25 +45,52 @@ export class NsfwProtection extends Protection {
 
     public async handleEvent(mjolnir: Mjolnir, roomId: string, event: any): Promise<any> {
         if (event['type'] === 'm.room.message') {
-            const content = event['content'] || {};
-            const msgtype = content['msgtype'] || 'm.text';
-            const isMedia = msgtype === 'm.image';
+            let content = JSON.stringify(event['content']);
+            // clean up brackets and "" to simplify regex
+            content = content.replace(/"|{|}/g, '')
+            if (!content.toLowerCase().includes("mxc")) {
+                return;
+             }
+            // try and grab a human-readable alias for more helpful management room output
+            const maybeAlias = await mjolnir.client.getPublishedAlias(roomId)
+            const room = maybeAlias ? maybeAlias : roomId
 
-            if (isMedia) {
-                const mxc = content["url"];
+            const mxcs = content.match(/(mxc?:\/\/[^\s]+)/gim);
+            if (!mxcs) {
+                //something's gone wrong with the regex
+                await mjolnir.managementRoomOutput.logMessage(LogLevel.ERROR, "NSFWProtection", `Unable to find any mxcs in  ${event["event_id"]} in ${room}`);
+                return;
+            }
+
+            // @ts-ignore - see null check immediately above
+            for (const mxc of mxcs) {
                 const image = await mjolnir.client.downloadContent(mxc);
                 const decodedImage = await node.decodeImage(image.data, 3);
                 const predictions = await this.model.classify(decodedImage);
 
+
                 for (const prediction of predictions) {
                     if (["Hentai", "Porn"].includes(prediction["className"])) {
                         if (prediction["probability"] > mjolnir.config.nsfwSensitivity) {
-                            await mjolnir.managementRoomOutput.logMessage(LogLevel.INFO, "NSFWProtection", `Redacting ${event["event_id"]} for inappropriate content.`);
                             try {
                                 await mjolnir.client.redactEvent(roomId, event["event_id"]);
                             } catch (err) {
-                                await mjolnir.managementRoomOutput.logMessage(LogLevel.ERROR, "NSFWProtection", `There was an error redacting ${event["event_id"]}: ${err}`);
+                                await mjolnir.managementRoomOutput.logMessage(LogLevel.ERROR, "NSFWProtection", `There was an error redacting ${event["event_id"]} in ${room}: ${err}`);
                             }
+                            let eventId = event["event_id"]
+                            let body = `Redacted an image in ${room} ${eventId}`
+                            let formatted_body = `<details>
+                                                  <summary>Redacted an image in ${room}</summary>
+                                                  <pre>${eventId}</pre>  <pre></pre>${room}</pre>
+                                                  </details>`
+                            const msg = {
+                                msgtype: "m.notice",
+                                body: body,
+                                format: "org.matrix.custom.html",
+                                formatted_body: formatted_body
+                            };
+                            await mjolnir.client.sendMessage(mjolnir.managementRoomId, msg);
+                            break
                         }
                     }
                 }

diff --git a/test/integration/nsfwProtectionTest.ts b/test/integration/nsfwProtectionTest.ts
@@ -59,8 +59,20 @@ describe("Test: NSFW protection", function () {
         let content = {"msgtype": "m.image", "body": "test.jpeg", "url": mxc};
         let imageMessage = await client.sendMessage(room, content);
 
+        let formatted_body = `<img src=${mxc} />`
+        let htmlContent = {
+                                msgtype: "m.image",
+                                body: formatted_body,
+                                format: "org.matrix.custom.html",
+                                formatted_body: formatted_body
+                            };
+        let htmlMessage = await client.sendMessage(room, htmlContent)
+
         await delay(500);
         let processedImage = await client.getEvent(room, imageMessage);
         assert.equal(Object.keys(processedImage.content).length, 0, "This event should have been redacted");
+
+        let processedHtml = await client.getEvent(room, htmlMessage)
+        assert.equal(Object.keys(processedHtml.content).length, 0, "This html image event should have been redacted")
     });
 });