Provide the clamav socket, correctly upsert file hashes

pull/42/head
Skyler Grey 3 years ago
parent 0a4846c560
commit d1157318c0
Signed by: Minion3665
GPG Key ID: 1AFD10256B3C714D

@ -1090,18 +1090,18 @@ const callback = async (interaction: CommandInteraction): Promise<void> => {
break; break;
} }
case "malware": { case "malware": {
i.deferUpdate(); await i.deferUpdate();
config.malware = !config.malware; config.malware = !config.malware;
break; break;
} }
case "images": { case "images": {
i.deferUpdate(); await i.deferUpdate();
const next = await imageMenu(i, m, config.images); const next = await imageMenu(i, m, config.images);
config.images = next; config.images = next;
break; break;
} }
case "clean": { case "clean": {
i.deferUpdate(); await i.deferUpdate();
const next = await cleanMenu(i, m, config.clean); const next = await cleanMenu(i, m, config.clean);
config.clean = next; config.clean = next;
break; break;

@ -19,6 +19,7 @@ declare const config: {
}; };
baseUrl: string; baseUrl: string;
rapidApiKey: string; rapidApiKey: string;
clamavSocket: string;
}; };
export default config; export default config;

@ -11,6 +11,7 @@ import * as tf from "@tensorflow/tfjs-node";
import EmojiEmbed from "../utils/generateEmojiEmbed.js"; import EmojiEmbed from "../utils/generateEmojiEmbed.js";
import getEmojiByName from "../utils/getEmojiByName.js"; import getEmojiByName from "../utils/getEmojiByName.js";
import { ActionRowBuilder, ButtonBuilder, ButtonStyle } from "discord.js"; import { ActionRowBuilder, ButtonBuilder, ButtonStyle } from "discord.js";
import config from "../config/main.js";
interface NSFWSchema { interface NSFWSchema {
nsfw: boolean; nsfw: boolean;
@ -22,7 +23,11 @@ interface MalwareSchema {
} }
const nsfw_model = await nsfwjs.load(); const nsfw_model = await nsfwjs.load();
const clamscanner = await new ClamScan().init({}); const clamscanner = await new ClamScan().init({
clamdscan: {
socket: config.clamavSocket
}
});
export async function testNSFW(link: string): Promise<NSFWSchema> { export async function testNSFW(link: string): Promise<NSFWSchema> {
const [fileStream, hash] = await streamAttachment(link); const [fileStream, hash] = await streamAttachment(link);
@ -34,26 +39,29 @@ export async function testNSFW(link: string): Promise<NSFWSchema> {
const predictions = (await nsfw_model.classify(image, 1))[0]!; const predictions = (await nsfw_model.classify(image, 1))[0]!;
image.dispose(); image.dispose();
return { nsfw: predictions.className === "Hentai" || predictions.className === "Porn" }; const nsfw = predictions.className === "Hentai" || predictions.className === "Porn";
await client.database.scanCache.write(hash, "nsfw", nsfw);
return { nsfw };
} }
export async function testMalware(link: string): Promise<MalwareSchema> { export async function testMalware(link: string): Promise<MalwareSchema> {
const [fileName, hash] = await saveAttachment(link); const [fileName, hash] = await saveAttachment(link);
const alreadyHaveCheck = await client.database.scanCache.read(hash); const alreadyHaveCheck = await client.database.scanCache.read(hash);
if (alreadyHaveCheck?.malware) return { safe: alreadyHaveCheck.malware }; if (alreadyHaveCheck?.malware) return { safe: alreadyHaveCheck.malware };
let safe; let malware;
try { try {
safe = !(await clamscanner.scanFile(fileName)).isInfected; malware = (await clamscanner.scanFile(fileName)).isInfected;
} catch (e) { } catch (e) {
return { safe: true }; return { safe: true };
} }
client.database.scanCache.write(hash, "malware", safe); client.database.scanCache.write(hash, "malware", malware);
return { safe }; return { safe: !malware };
} }
export async function testLink(link: string): Promise<{ safe: boolean; tags: string[] }> { export async function testLink(link: string): Promise<{ safe: boolean; tags: string[] }> {
const alreadyHaveCheck = await client.database.scanCache.read(link); const alreadyHaveCheck = await client.database.scanCache.read(link);
if (alreadyHaveCheck?.bad_link) return { safe: alreadyHaveCheck.bad_link, tags: alreadyHaveCheck.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", { const scanned: { safe?: boolean; tags?: string[] } = await fetch("https://unscan.p.rapidapi.com/link", {
method: "POST", method: "POST",
headers: { headers: {

@ -591,7 +591,7 @@ interface ScanCacheSchema {
nsfw?: boolean; nsfw?: boolean;
malware?: boolean; malware?: boolean;
bad_link?: boolean; bad_link?: boolean;
tags: string[]; tags?: string[];
} }
export class ScanCache { export class ScanCache {
@ -608,7 +608,27 @@ export class ScanCache {
async write(hash: string, type: "nsfw" | "malware" | "bad_link", data: boolean, tags?: string[]) { async write(hash: string, type: "nsfw" | "malware" | "bad_link", data: boolean, tags?: string[]) {
await this.scanCache.updateOne( await this.scanCache.updateOne(
{ hash: hash }, { hash: hash },
{ hash: hash, [type]: data, tags: tags ?? [], addedAt: new Date() }, {
$set: (() => {
switch (type) {
case "nsfw": {
return { nsfw: data, addedAt: new Date() };
}
case "malware": {
return { malware: data, addedAt: new Date() };
}
case "bad_link": {
return { bad_link: data, tags: tags ?? [], addedAt: new Date() };
}
default: {
throw new Error("Invalid type");
}
}
})()
// No you can't just do { [type]: data }, yes it's a typescript error, no I don't know how to fix it
// cleanly, yes it would be marginally more elegant, no it's not essential, yes I'd be happy to review
// PRs that did improve this snippet
},
Object.assign({ upsert: true }, collectionOptions) Object.assign({ upsert: true }, collectionOptions)
); );
} }

Loading…
Cancel
Save