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;
}
case "malware": {
i.deferUpdate();
await i.deferUpdate();
config.malware = !config.malware;
break;
}
case "images": {
i.deferUpdate();
await i.deferUpdate();
const next = await imageMenu(i, m, config.images);
config.images = next;
break;
}
case "clean": {
i.deferUpdate();
await i.deferUpdate();
const next = await cleanMenu(i, m, config.clean);
config.clean = next;
break;

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

@ -11,6 +11,7 @@ import * as tf from "@tensorflow/tfjs-node";
import EmojiEmbed from "../utils/generateEmojiEmbed.js";
import getEmojiByName from "../utils/getEmojiByName.js";
import { ActionRowBuilder, ButtonBuilder, ButtonStyle } from "discord.js";
import config from "../config/main.js";
interface NSFWSchema {
nsfw: boolean;
@ -22,7 +23,11 @@ interface MalwareSchema {
}
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> {
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]!;
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> {
const [fileName, hash] = await saveAttachment(link);
const alreadyHaveCheck = await client.database.scanCache.read(hash);
if (alreadyHaveCheck?.malware) return { safe: alreadyHaveCheck.malware };
let safe;
let malware;
try {
safe = !(await clamscanner.scanFile(fileName)).isInfected;
malware = (await clamscanner.scanFile(fileName)).isInfected;
} catch (e) {
return { safe: true };
}
client.database.scanCache.write(hash, "malware", safe);
return { safe };
client.database.scanCache.write(hash, "malware", malware);
return { safe: !malware };
}
export async function testLink(link: string): Promise<{ safe: boolean; tags: string[] }> {
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", {
method: "POST",
headers: {

@ -591,7 +591,7 @@ interface ScanCacheSchema {
nsfw?: boolean;
malware?: boolean;
bad_link?: boolean;
tags: string[];
tags?: string[];
}
export class ScanCache {
@ -608,7 +608,27 @@ export class ScanCache {
async write(hash: string, type: "nsfw" | "malware" | "bad_link", data: boolean, tags?: string[]) {
await this.scanCache.updateOne(
{ 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)
);
}

Loading…
Cancel
Save