Add a NSFW protection

config/default.yaml

@@ -258,3 +258,8 @@ pollReports: false
 # Whether or not new reports, received either by webapi or polling,
 # should be printed to our managementRoom.
 displayReports: true
+
+# How sensitive the NsfwProtection should be, which determines if an image should be redacted. A number between 0 - .99,
+# with a lower number indicating greater sensitivity, possibly resulting in images being more aggressively flagged
+# and redacted as NSFW
+nsfwSensitivity: .6

package.json

@@ -46,6 +46,7 @@
   "dependencies": {
     "@sentry/node": "^7.17.2",
     "@sentry/tracing": "^7.17.2",
+    "@tensorflow/tfjs-node": "^4.21.0",
     "await-lock": "^2.2.2",
     "body-parser": "^1.20.1",
     "config": "^3.3.8",
@@ -56,6 +57,7 @@
     "js-yaml": "^4.1.0",
     "jsdom": "^16.6.0",
     "matrix-appservice-bridge": "8.1.2",
+    "nsfwjs": "^4.1.0",
     "parse-duration": "^1.0.2",
     "pg": "^8.8.0",
     "prom-client": "^14.1.0",

src/config.ts

@@ -174,6 +174,7 @@ export interface IConfig {
             enabled: boolean;
         }
     }
+    nsfwSensitivity: number
 
     /**
      * Config options only set at runtime. Try to avoid using the objects
@@ -248,7 +249,7 @@ const defaultConfig: IConfig = {
             enabled: false,
         },
     },
-
+    nsfwSensitivity: .6,
     // Needed to make the interface happy.
     RUNTIME: {
     },

src/protections/NsfwProtection.ts

@@ -0,0 +1,85 @@
+/*
+Copyright 2024 The Matrix.org Foundation C.I.C.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+import { Protection } from "./IProtection";
+import { Mjolnir } from "../Mjolnir";
+import * as nsfw from 'nsfwjs';
+import {LogLevel} from "matrix-bot-sdk";
+import { node } from '@tensorflow/tfjs-node';
+
+
+export class NsfwProtection extends Protection {
+    settings = {};
+    // @ts-ignore
+    private model: any;
+
+    constructor() {
+        super();
+    }
+
+    async initialize() {
+        this.model = await nsfw.load();
+    }
+
+    public get name(): string {
+        return 'NsfwProtection';
+    }
+
+    public get description(): string {
+        return "Scans all images sent into a protected room to determine if the image is " +
+            "NSFW. If it is, the image will automatically be redacted.";
+    }
+
+    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 formattedBody = content['formatted_body'] || '';
+            const isMedia = msgtype === 'm.image' || formattedBody.toLowerCase().includes('<img');
+
+            if (isMedia) {
+                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)
+
+                for (const prediction of predictions) {
+                    if (prediction["className"] === "Porn") {
+                        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}`);
+
+                            }
+                        }
+                    } 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()
+            }
+        }
+    }
+}

src/protections/ProtectionManager.ts

@@ -31,6 +31,7 @@ import { htmlEscape } from "../utils";
 import { ERROR_KIND_FATAL, ERROR_KIND_PERMISSION } from "../ErrorCache";
 import { RoomUpdateError } from "../models/RoomUpdateError";
 import { LocalAbuseReports } from "./LocalAbuseReports";
+import {NsfwProtection} from "./NsfwProtection";
 
 const PROTECTIONS: Protection[] = [
     new FirstMessageIsImage(),
@@ -42,6 +43,7 @@ const PROTECTIONS: Protection[] = [
     new DetectFederationLag(),
     new JoinWaveShortCircuit(),
     new LocalAbuseReports(),
+    new NsfwProtection()
 ];
 
 const ENABLED_PROTECTIONS_EVENT_TYPE = "org.matrix.mjolnir.enabled_protections";
@@ -100,6 +102,9 @@ export class ProtectionManager {
             protection.settings[key].setValue(value);
         }
         if (protection.enabled) {
+            if (protection.name === "NsfwProtection") {
+                (protection as NsfwProtection).initialize()
+            }
             for (let roomId of this.mjolnir.protectedRoomsTracker.getProtectedRooms()) {
                 await protection.startProtectingRoom(this.mjolnir, roomId);
             }

test/integration/fixtures.ts

@@ -1,6 +1,5 @@
 import { read as configRead } from "../../src/config";
 import { makeMjolnir, teardownManagementRoom } from "./mjolnirSetupUtils";
-import { register } from "prom-client";
 import dns from 'node:dns';
 
 // Necessary for CI: Node 17+ defaults to using ipv6 first, but Github Actions does not support ipv6

test/integration/nsfwProtectionTest.ts

@@ -0,0 +1,66 @@
+import {newTestUser} from "./clientHelper";
+
+import {MatrixClient} from "matrix-bot-sdk";
+import {getFirstReaction} from "./commands/commandUtils";
+import {strict as assert} from "assert";
+import { readFileSync } from 'fs';
+
+describe("Test: NSFW protection", function () {
+    let client: MatrixClient;
+    let room: string;
+    this.beforeEach(async function () {
+        client = await newTestUser(this.config.homeserverUrl, {name: {contains: "nsfw-protection"}});
+        await client.start();
+        const mjolnirId = await this.mjolnir.client.getUserId();
+        room = await client.createRoom({ invite: [mjolnirId] });
+        await client.joinRoom(room);
+        await client.joinRoom(this.config.managementRoom);
+        await client.setUserPowerLevel(mjolnirId, room, 100)
+    })
+    this.afterEach(async function () {
+        await client.stop();
+    })
+
+    function delay(ms: number) {
+        return new Promise(resolve => setTimeout(resolve, ms));
+    }
+
+
+    it("Nsfw protection doesn't redact sfw images", async function() {
+        this.timeout(20000);
+
+        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)
+
+        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");
+    });
+
+    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
+
+        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)
+
+        await delay(500)
+        let processedImage = await client.getEvent(room, imageMessage);
+        assert.equal(Object.keys(processedImage.content).length, 0, "This event should have been redacted");
+    });
+});

tsconfig.json

@@ -1,5 +1,6 @@
 {
     "compilerOptions": {
+        "skipLibCheck": true,
         "alwaysStrict": true,
         "esModuleInterop": true,
         "experimentalDecorators": true,