From 61a0fd5e3878614f2af61766d9761c70c0bcfda2 Mon Sep 17 00:00:00 2001
From: "H. Shay" <hillerys@element.io>
Date: Wed, 11 Sep 2024 15:05:11 -0700
Subject: [PATCH] requested changes

---
 src/protections/NsfwProtection.ts      | 22 ++++++----------------
 src/protections/ProtectionManager.ts   |  2 +-
 test/integration/nsfwProtectionTest.ts | 24 ++++++++++++------------
 3 files changed, 19 insertions(+), 29 deletions(-)

diff --git a/src/protections/NsfwProtection.ts b/src/protections/NsfwProtection.ts
index 61c1549..2237b0b 100644
--- a/src/protections/NsfwProtection.ts
+++ b/src/protections/NsfwProtection.ts
@@ -47,17 +47,16 @@ export class NsfwProtection extends Protection {
         if (event['type'] === 'm.room.message') {
             const content = event['content'] || {};
             const msgtype = content['msgtype'] || 'm.text';
-            const formattedBody = content['formatted_body'] || '';
-            const isMedia = msgtype === 'm.image' || formattedBody.toLowerCase().includes('<img');
+            const isMedia = msgtype === 'm.image';
 
             if (isMedia) {
-                const mxc = content["url"]
-                const image = await mjolnir.client.downloadContent(mxc)
+                const mxc = content["url"];
+                const image = await mjolnir.client.downloadContent(mxc);
                 const decodedImage = await node.decodeImage(image.data, 3);
-                const predictions = await this.model.classify(decodedImage)
+                const predictions = await this.model.classify(decodedImage);
 
                 for (const prediction of predictions) {
-                    if (prediction["className"] === "Porn") {
+                    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 {
@@ -67,18 +66,9 @@ export class NsfwProtection extends Protection {
 
                             }
                         }
-                    } else if (prediction["className"] === "Hentai") {
-                        if (prediction["probability"] > mjolnir.config.nsfwSensitivity) {
-                            await mjolnir.managementRoomOutput.logMessage(LogLevel.INFO, "NSFWProtection", `Redacting ${event["event_id"]} for inappropriate content.`);
-                            try {
-                                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}`);
-                            }
-                        }
                     }
                 }
-                decodedImage.dispose()
+                decodedImage.dispose();
             }
         }
     }
diff --git a/src/protections/ProtectionManager.ts b/src/protections/ProtectionManager.ts
index 9b460c9..9bd1e2c 100644
--- a/src/protections/ProtectionManager.ts
+++ b/src/protections/ProtectionManager.ts
@@ -103,7 +103,7 @@ export class ProtectionManager {
         }
         if (protection.enabled) {
             if (protection.name === "NsfwProtection") {
-                (protection as NsfwProtection).initialize()
+                (protection as NsfwProtection).initialize();
             }
             for (let roomId of this.mjolnir.protectedRoomsTracker.getProtectedRooms()) {
                 await protection.startProtectingRoom(this.mjolnir, roomId);
diff --git a/test/integration/nsfwProtectionTest.ts b/test/integration/nsfwProtectionTest.ts
index 7e83b94..f1aa8f1 100644
--- a/test/integration/nsfwProtectionTest.ts
+++ b/test/integration/nsfwProtectionTest.ts
@@ -15,7 +15,7 @@ describe("Test: NSFW protection", function () {
         room = await client.createRoom({ invite: [mjolnirId] });
         await client.joinRoom(room);
         await client.joinRoom(this.config.managementRoom);
-        await client.setUserPowerLevel(mjolnirId, room, 100)
+        await client.setUserPowerLevel(mjolnirId, room, 100);
     })
     this.afterEach(async function () {
         await client.stop();
@@ -34,12 +34,12 @@ describe("Test: NSFW protection", function () {
                 return await client.sendMessage(this.mjolnir.managementRoomId, { msgtype: 'm.text', body: `!mjolnir enable NsfwProtection` });
         });
 
-        const data = readFileSync('test_tree.jpg')
-        const mxc = await client.uploadContent(data, 'image/png')
-        let content = {"msgtype": "m.image", "body": "test.jpeg", "url": mxc}
-        let imageMessage = await client.sendMessage(room, content)
+        const data = readFileSync('test_tree.jpg');
+        const mxc = await client.uploadContent(data, 'image/png');
+        let content = {"msgtype": "m.image", "body": "test.jpeg", "url": mxc};
+        let imageMessage = await client.sendMessage(room, content);
 
-        await delay(500)
+        await delay(500);
         let processedImage = await client.getEvent(room, imageMessage);
         assert.equal(Object.keys(processedImage.content).length, 3, "This event should not have been redacted");
     });
@@ -47,19 +47,19 @@ describe("Test: NSFW protection", function () {
     it("Nsfw protection redacts nsfw images", async function() {
         this.timeout(20000);
         // dial the sensitivity on the protection way up so that all images are flagged as NSFW
-        this.mjolnir.config.nsfwSensitivity = 0.0
+        this.mjolnir.config.nsfwSensitivity = 0.0;
 
         await client.sendMessage(this.mjolnir.managementRoomId, { msgtype: 'm.text', body: `!mjolnir rooms add ${room}` });
         await getFirstReaction(client, this.mjolnir.managementRoomId, '✅', async () => {
                 return await client.sendMessage(this.mjolnir.managementRoomId, { msgtype: 'm.text', body: `!mjolnir enable NsfwProtection` });
         });
 
-        const data = readFileSync('test_tree.jpg')
-        const mxc = await client.uploadContent(data, 'image/png')
-        let content = {"msgtype": "m.image", "body": "test.jpeg", "url": mxc}
-        let imageMessage = await client.sendMessage(room, content)
+        const data = readFileSync('test_tree.jpg');
+        const mxc = await client.uploadContent(data, 'image/png');
+        let content = {"msgtype": "m.image", "body": "test.jpeg", "url": mxc};
+        let imageMessage = await client.sendMessage(room, content);
 
-        await delay(500)
+        await delay(500);
         let processedImage = await client.getEvent(room, imageMessage);
         assert.equal(Object.keys(processedImage.content).length, 0, "This event should have been redacted");
     });