Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Improve messaging and checks in NSFW protection #529

Merged
merged 5 commits into from
Oct 1, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
41 changes: 34 additions & 7 deletions src/protections/NsfwProtection.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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, '')
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

this could do with a comment for what it's fixing about the input (I'm not sure myself).

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
}
}
}
Expand Down
12 changes: 12 additions & 0 deletions test/integration/nsfwProtectionTest.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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")
});
});
Loading