diff --git a/src/protections/NsfwProtection.ts b/src/protections/NsfwProtection.ts index 6be229ae..5269d5e7 100644 --- 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 { 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 = `
+ Redacted an image in ${room} +
${eventId}
${room}
+                                                  
` + 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 index 7b867fde..c86fd384 100644 --- 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 = `` + 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") }); }); \ No newline at end of file