From 144327101dbca557911c3c8417bc4c21d7fa9786 Mon Sep 17 00:00:00 2001 From: Skyler Grey Date: Tue, 7 Mar 2023 23:40:50 +0000 Subject: [PATCH 1/2] Fix nickname and nsfw pfp scanning --- package.json | 2 +- src/premium/attachmentLogs.ts | 6 ++--- src/reflex/scanners.ts | 41 ++++++++++++++++++----------------- src/utils/database.ts | 10 ++++----- 4 files changed, 30 insertions(+), 29 deletions(-) diff --git a/package.json b/package.json index 6457877..9d3df42 100644 --- a/package.json +++ b/package.json @@ -21,7 +21,7 @@ "mongodb": "^4.7.0", "node-fetch": "^3.3.0", "node-tesseract-ocr": "^2.2.1", - "nsfwjs": "2.4.2", + "nsfwjs": "https://github.com/infinitered/nsfwjs", "seedrandom": "^3.0.5", "structured-clone": "^0.2.2", "systeminformation": "^5.17.3" diff --git a/src/premium/attachmentLogs.ts b/src/premium/attachmentLogs.ts index e719f22..ca60b2c 100644 --- a/src/premium/attachmentLogs.ts +++ b/src/premium/attachmentLogs.ts @@ -2,7 +2,7 @@ import { getCommandMentionByName } from "./../utils/getCommandDataByName.js"; import client from "../utils/client.js"; import keyValueList from "../utils/generateKeyValueList.js"; import singleNotify from "../utils/singleNotify.js"; -import { saveAttachment } from "../reflex/scanners.js"; +import { streamAttachment } from "../reflex/scanners.js"; import EmojiEmbed from "../utils/generateEmojiEmbed.js"; import addPlural from "../utils/plurals.js"; import type { GuildTextBasedChannel, Message } from "discord.js"; @@ -14,7 +14,7 @@ export default async function logAttachment(message: Message): Promise { - const [_fileName, hash] = await saveAttachment(link); + const [fileStream, hash] = await streamAttachment(link); const alreadyHaveCheck = await client.database.scanCache.read(hash); - if (alreadyHaveCheck) return { nsfw: alreadyHaveCheck.data }; + if (alreadyHaveCheck?.nsfw) return { nsfw: alreadyHaveCheck.nsfw }; - // const image = tf.node.decodePng() + const image = tf.node.decodeImage(new Uint8Array(fileStream), 3) as tf.Tensor3D; - // const result = await model.classify(image) + const predictions = (await nsfw_model.classify(image, 1))[0]!; + image.dispose(); - return { nsfw: false }; + return { nsfw: predictions.className === "Hentai" || predictions.className === "Porn" }; } export async function testMalware(link: string): Promise { - const [p, hash] = await saveAttachment(link); + const [p, hash] = await streamAttachment(link); const alreadyHaveCheck = await client.database.scanCache.read(hash); - if (alreadyHaveCheck) return { safe: alreadyHaveCheck.data }; + if (alreadyHaveCheck?.malware) return { safe: alreadyHaveCheck.malware }; + return { safe: true }; const data = new URLSearchParams(); - const f = createReadStream(p); + // const f = createReadStream(p); data.append("file", f.read(fs.statSync(p).size)); const result = await fetch("https://unscan.p.rapidapi.com/malware", { method: "POST", @@ -58,14 +60,14 @@ export async function testMalware(link: string): Promise { return { safe: true, errored: true }; }); if (!result.errored) { - client.database.scanCache.write(hash, result.safe); + client.database.scanCache.write(hash, "malware", result.safe); } return { safe: result.safe }; } export async function testLink(link: string): Promise<{ safe: boolean; tags: string[] }> { const alreadyHaveCheck = await client.database.scanCache.read(link); - if (alreadyHaveCheck) return { safe: alreadyHaveCheck.data, tags: [] }; + if (alreadyHaveCheck?.bad_link) return { safe: alreadyHaveCheck.bad_link, tags: alreadyHaveCheck.tags }; const scanned: { safe?: boolean; tags?: string[] } = await fetch("https://unscan.p.rapidapi.com/link", { method: "POST", headers: { @@ -79,19 +81,19 @@ export async function testLink(link: string): Promise<{ safe: boolean; tags: str console.error(err); return { safe: true, tags: [] }; }); - client.database.scanCache.write(link, scanned.safe ?? true, []); + client.database.scanCache.write(link, "bad_link", scanned.safe ?? true, scanned.tags ?? []); return { safe: scanned.safe ?? true, tags: scanned.tags ?? [] }; } -export async function saveAttachment(link: string): Promise<[string, string]> { +export async function streamAttachment(link: string): Promise<[ArrayBuffer, string]> { const image = await (await fetch(link)).arrayBuffer(); const fileName = generateFileName(link.split("/").pop()!.split(".").pop()!); const enc = new TextDecoder("utf-8"); writeFileSync(fileName, new DataView(image), "base64"); - return [fileName, createHash("sha512").update(enc.decode(image), "base64").digest("base64")]; + return [image, createHash("sha512").update(enc.decode(image), "base64").digest("base64")]; } const linkTypes = { @@ -218,11 +220,10 @@ export async function doMemberChecks(member: Discord.GuildMember, guild: Discord const avatarCheck = guildData.filters.images.NSFW && (await NSFWCheck(member.user.displayAvatarURL({ forceStatic: true }))); // Does the username contain an invite - const inviteCheck = - guildData.filters.invite.enabled && member.user.username.match(/discord\.gg\/[a-zA-Z0-9]+/gi) !== null; + const inviteCheck = guildData.filters.invite.enabled && /discord\.gg\/[a-zA-Z0-9]+/gi.test(member.user.username); // Does the nickname contain an invite const nicknameInviteCheck = - guildData.filters.invite.enabled && member.nickname?.match(/discord\.gg\/[a-zA-Z0-9]+/gi) !== null; + guildData.filters.invite.enabled && /discord\.gg\/[a-zA-Z0-9]+/gi.test(member.nickname ?? ""); if ( usernameCheck !== null || diff --git a/src/utils/database.ts b/src/utils/database.ts index 4f94712..a107d06 100644 --- a/src/utils/database.ts +++ b/src/utils/database.ts @@ -588,7 +588,9 @@ export class History { interface ScanCacheSchema { addedAt: Date; hash: string; - data: boolean; + nsfw?: boolean; + malware?: boolean; + bad_link?: boolean; tags: string[]; } @@ -600,14 +602,12 @@ export class ScanCache { } async read(hash: string) { - // console.log("ScanCache read"); return await this.scanCache.findOne({ hash: hash }); } - async write(hash: string, data: boolean, tags?: string[]) { - // console.log("ScanCache write"); + async write(hash: string, type: "nsfw" | "malware" | "bad_link", data: boolean, tags?: string[]) { await this.scanCache.insertOne( - { hash: hash, data: data, tags: tags ?? [], addedAt: new Date() }, + { hash: hash, [type]: data, tags: tags ?? [], addedAt: new Date() }, collectionOptions ); } From 3217998db27e7824182aea043a00474269391628 Mon Sep 17 00:00:00 2001 From: Skyler Grey Date: Tue, 7 Mar 2023 23:59:06 +0000 Subject: [PATCH 2/2] Downgrade nsfwjs and return saveAttachment as it is elsewhere used --- package.json | 4 +-- src/premium/attachmentLogs.ts | 6 ++-- src/reflex/scanners.ts | 64 +++++++++++++++++++---------------- 3 files changed, 40 insertions(+), 34 deletions(-) diff --git a/package.json b/package.json index 9d3df42..629c6dd 100644 --- a/package.json +++ b/package.json @@ -1,7 +1,7 @@ { "dependencies": { "@hokify/agenda": "^6.2.12", - "@tensorflow/tfjs": "^4.0.0", + "@tensorflow/tfjs": "^3.18.0", "@tensorflow/tfjs-node": "^4.2.0", "@total-typescript/ts-reset": "^0.3.7", "@tsconfig/node18-strictest-esm": "^1.0.0", @@ -21,7 +21,7 @@ "mongodb": "^4.7.0", "node-fetch": "^3.3.0", "node-tesseract-ocr": "^2.2.1", - "nsfwjs": "https://github.com/infinitered/nsfwjs", + "nsfwjs": "^2.4.2", "seedrandom": "^3.0.5", "structured-clone": "^0.2.2", "systeminformation": "^5.17.3" diff --git a/src/premium/attachmentLogs.ts b/src/premium/attachmentLogs.ts index ca60b2c..e719f22 100644 --- a/src/premium/attachmentLogs.ts +++ b/src/premium/attachmentLogs.ts @@ -2,7 +2,7 @@ import { getCommandMentionByName } from "./../utils/getCommandDataByName.js"; import client from "../utils/client.js"; import keyValueList from "../utils/generateKeyValueList.js"; import singleNotify from "../utils/singleNotify.js"; -import { streamAttachment } from "../reflex/scanners.js"; +import { saveAttachment } from "../reflex/scanners.js"; import EmojiEmbed from "../utils/generateEmojiEmbed.js"; import addPlural from "../utils/plurals.js"; import type { GuildTextBasedChannel, Message } from "discord.js"; @@ -14,7 +14,7 @@ export default async function logAttachment(message: Message): Promise { const alreadyHaveCheck = await client.database.scanCache.read(hash); if (alreadyHaveCheck?.nsfw) return { nsfw: alreadyHaveCheck.nsfw }; - const image = tf.node.decodeImage(new Uint8Array(fileStream), 3) as tf.Tensor3D; + const image = tf.tensor3d(new Uint8Array(fileStream)); const predictions = (await nsfw_model.classify(image, 1))[0]!; image.dispose(); @@ -37,32 +37,32 @@ export async function testNSFW(link: string): Promise { } export async function testMalware(link: string): Promise { - const [p, hash] = await streamAttachment(link); + const [_, hash] = await saveAttachment(link); const alreadyHaveCheck = await client.database.scanCache.read(hash); if (alreadyHaveCheck?.malware) return { safe: alreadyHaveCheck.malware }; return { safe: true }; - const data = new URLSearchParams(); - // const f = createReadStream(p); - data.append("file", f.read(fs.statSync(p).size)); - const result = await fetch("https://unscan.p.rapidapi.com/malware", { - method: "POST", - headers: { - "X-RapidAPI-Key": client.config.rapidApiKey, - "X-RapidAPI-Host": "unscan.p.rapidapi.com" - }, - body: data - }) - .then((response) => - response.status === 200 ? (response.json() as Promise) : { safe: true, errored: true } - ) - .catch((err) => { - console.error(err); - return { safe: true, errored: true }; - }); - if (!result.errored) { - client.database.scanCache.write(hash, "malware", result.safe); - } - return { safe: result.safe }; + // const data = new URLSearchParams(); + // // const f = createReadStream(p); + // data.append("file", f.read(fs.statSync(p).size)); + // const result = await fetch("https://unscan.p.rapidapi.com/malware", { + // method: "POST", + // headers: { + // "X-RapidAPI-Key": client.config.rapidApiKey, + // "X-RapidAPI-Host": "unscan.p.rapidapi.com" + // }, + // body: data + // }) + // .then((response) => + // response.status === 200 ? (response.json() as Promise) : { safe: true, errored: true } + // ) + // .catch((err) => { + // console.error(err); + // return { safe: true, errored: true }; + // }); + // if (!result.errored) { + // client.database.scanCache.write(hash, "malware", result.safe); + // } + // return { safe: result.safe }; } export async function testLink(link: string): Promise<{ safe: boolean; tags: string[] }> { @@ -89,11 +89,17 @@ export async function testLink(link: string): Promise<{ safe: boolean; tags: str } export async function streamAttachment(link: string): Promise<[ArrayBuffer, string]> { + const image = await (await fetch(link)).arrayBuffer(); + const enc = new TextDecoder("utf-8"); + return [image, createHash("sha512").update(enc.decode(image), "base64").digest("base64")]; +} + +export async function saveAttachment(link: string): Promise<[string, string]> { const image = await (await fetch(link)).arrayBuffer(); const fileName = generateFileName(link.split("/").pop()!.split(".").pop()!); const enc = new TextDecoder("utf-8"); writeFileSync(fileName, new DataView(image), "base64"); - return [image, createHash("sha512").update(enc.decode(image), "base64").digest("base64")]; + return [fileName, createHash("sha512").update(enc.decode(image), "base64").digest("base64")]; } const linkTypes = {