From 3d8ae1ebb8a1567ece5501d04cf63ef4c38bf54f Mon Sep 17 00:00:00 2001 From: Donovan Daniels Date: Sat, 19 Oct 2024 13:43:43 -0500 Subject: [PATCH] overhaul auto source handling --- docker-compose.dev.yml | 4 + src/config/index.ts | 6 + src/config/private | 2 +- src/events/messageCreate.ts | 19 +- src/util/@types/femboyfans.d.ts | 105 +++++++++++ src/util/RequestProxy.ts | 5 +- src/util/Sauce.ts | 296 +++++++++++++++++++++----------- src/util/req/FemboyFans.ts | 61 +++++++ 8 files changed, 389 insertions(+), 109 deletions(-) create mode 100644 docker-compose.dev.yml create mode 100644 src/util/@types/femboyfans.d.ts create mode 100644 src/util/req/FemboyFans.ts diff --git a/docker-compose.dev.yml b/docker-compose.dev.yml new file mode 100644 index 0000000..56377b2 --- /dev/null +++ b/docker-compose.dev.yml @@ -0,0 +1,4 @@ +services: + maidboye: + image: maidboye + build: . diff --git a/src/config/index.ts b/src/config/index.ts index e8eb5a6..4e70090 100644 --- a/src/config/index.ts +++ b/src/config/index.ts @@ -137,6 +137,12 @@ export class Configuration extends PrivateConfiguration { static override get tempAuth() { return super.tempAuth; } + static override get femboyFansUser() { + return super.femboyFansUser; + } + static override get femboyFansKey() { + return super.femboyFansKey; + } /* db */ static get dbHost() { diff --git a/src/config/private b/src/config/private index a9350cd..711b8b1 160000 --- a/src/config/private +++ b/src/config/private @@ -1 +1 @@ -Subproject commit a9350cd0f09e873dc557af37a5eaa4290dce56be +Subproject commit 711b8b1c58fbb0b3aa96eb3444d345db94968b2e diff --git a/src/events/messageCreate.ts b/src/events/messageCreate.ts index cd0a1df..664d0cd 100644 --- a/src/events/messageCreate.ts +++ b/src/events/messageCreate.ts @@ -187,6 +187,9 @@ export default new ClientEvent("messageCreate", async function messageCreateEven if (sauce.post !== null && !sources.some(s => s.startsWith("https://e621.net/posts/"))) { sources.unshift(`https://e621.net/posts/${sauce.post.id}`); } + if (sauce.ffpost !== null && !sources.some(s => s.startsWith("https://femboy.fan/posts/"))) { + sources.unshift(`https://femboy.fan/posts/${sauce.ffpost.id}`); + } await this.rest.channels.createMessage(msg.channelID, { content: sources.join("\n"), messageReference: { @@ -214,6 +217,9 @@ export default new ClientEvent("messageCreate", async function messageCreateEven if (sauce.post !== null && !sources.some(s => s.startsWith("https://e621.net/posts/"))) { sources.unshift(`https://e621.net/posts/${sauce.post.id}`); } + if (sauce.ffpost !== null && !sources.some(s => s.startsWith("https://femboy.fan/posts/"))) { + sources.unshift(`https://femboy.fan/posts/${sauce.ffpost.id}`); + } if (sources.length !== 0) { await this.rest.channels.createMessage(msg.channelID, { content: sources.join("\n"), @@ -252,11 +258,14 @@ export default new ClientEvent("messageCreate", async function messageCreateEven } }); const md5 = createHash("md5").update(Buffer.from(await file.arrayBuffer())).digest("hex"); - const att = await directMD5(md5); - if (att !== null) { - const sources = Array.isArray(att.sourceOverride) ? att.sourceOverride : (att.sourceOverride === undefined ? [] : [att.sourceOverride]); - if (att.post !== null && !sources.some(s => s.startsWith("https://e621.net/posts/"))) { - sources.unshift(`https://e621.net/posts/${att.post.id}`); + const sauce = await directMD5(md5); + if (sauce !== null) { + const sources = Array.isArray(sauce.sourceOverride) ? sauce.sourceOverride : (sauce.sourceOverride === undefined ? [] : [sauce.sourceOverride]); + if (sauce.post !== null && !sources.some(s => s.startsWith("https://e621.net/posts/"))) { + sources.unshift(`https://e621.net/posts/${sauce.post.id}`); + } + if (sauce.ffpost !== null && !sources.some(s => s.startsWith("https://femboy.fan/posts/"))) { + sources.unshift(`https://femboy.fan/posts/${sauce.ffpost.id}`); } if (sources.length !== 0) { await this.rest.channels.createMessage(msg.channelID, { diff --git a/src/util/@types/femboyfans.d.ts b/src/util/@types/femboyfans.d.ts new file mode 100644 index 0000000..8d4c232 --- /dev/null +++ b/src/util/@types/femboyfans.d.ts @@ -0,0 +1,105 @@ +declare namespace FemboyFans { + export interface Post { + approver_id: number | null; + change_seq: number; + comment_count: number; + created_at: string; + crop: Crop; + description: string; + duration: null; + fav_count: number; + file: File; + flags: Flags; + framecount: null; + has_notes: boolean; + id: number; + is_favorited: boolean; + locked_tags: Array; + own_vote: number; + pools: Array; + preview: Crop; + qtags: Array; + rating: string; + relationships: Relationships; + sample: Crop; + score: Score; + sources: Array; + tags: Tags; + thumbnail_frame: null; + updated_at: string; + upload_url: string | null; + uploader_id: number; + views: Views; + } + + export interface Crop { + alternates?: Alternates; + has?: boolean; + height: number; + url: string; + width: number; + } + + export interface Alternates { + "480p"?: Alternate; + "720p"?: Alternate; + "original": Alternate; + } + + export interface Alternate { + height: number; + type: "video"; + urls: [webm: string | null, mp4: string | null]; + width: number; + } + + export interface File { + ext: string; + height: number; + md5: string; + size: number; + url: string; + width: number; + } + + export interface Flags { + deleted: boolean; + flagged: boolean; + note_locked: boolean; + pending: boolean; + rating_locked: boolean; + status_locked: boolean; + } + + export interface Relationships { + children: Array; + has_active_children: boolean; + has_children: boolean; + parent_id: number | null; + } + + export interface Score { + down: number; + total: number; + up: number; + } + + export interface Tags { + artist: Array; + character: Array; + copyright: Array; + gender: Array; + general: Array; + invalid: Array; + lore: Array; + meta: Array; + species: Array; + voice_actor: Array; + } + + export interface Views { + daily: number; + total: number; + } + +} diff --git a/src/util/RequestProxy.ts b/src/util/RequestProxy.ts index 7011ccb..b80b9b8 100644 --- a/src/util/RequestProxy.ts +++ b/src/util/RequestProxy.ts @@ -3,7 +3,10 @@ import { fetch } from "undici"; export default class RequestProxy { static DIRECT_WHITELIST = [ - /^https?:\/\/static\d+\.e(621|926)\.net\//, + /^https?:\/\/static\d+\.e(621|926)\.net/, + /^https?:\/\/static\.femboy\.fan/, + /^https?:\/\/v2\.yiff\.media/, + /^https?:\/\/yiff\.media\/V2/, /^https:\/\/media\.discordapp\.net/, /^https:\/\/cdn\.discordapp\.com/ ]; diff --git a/src/util/Sauce.ts b/src/util/Sauce.ts index 48f145b..1352011 100644 --- a/src/util/Sauce.ts +++ b/src/util/Sauce.ts @@ -2,14 +2,13 @@ import E621 from "./req/E621.js"; import SauceNAO from "./req/SauceNAO.js"; import RequestProxy from "./RequestProxy.js"; import Yiffy from "./req/Yiffy.js"; +import FemboyFans from "./req/FemboyFans.js"; import Config from "../config/index.js"; import type { Post } from "e621"; import type { JSONResponse } from "yiffy"; import { Strings } from "@uwu-codes/utils"; -import { fetch } from "undici"; import { fileTypeFromBuffer } from "file-type"; import { STATUS_CODES } from "node:http"; -import { Blob } from "node:buffer"; export class PreCheckError extends Error { override name = "PreCheckError"; @@ -21,91 +20,121 @@ export const autoMimeTypes = [ "image/png", "image/webp" ]; -export default async function Sauce(input: string, simularity = 80, skipCheck = false, skipSauceNao = false) { +export default async function Sauce(input: string, similarity = 80, skipCheck = false, skipSauceNao = false) { // I could include file extensions, but I couldn't be bothered since I only need the md5 const yrRegex = /(?:https?:\/\/)?yiff\.rocks\/([A-Z_-\da-z]+)/; // E621 - e621.net const e621Regex = /(?:https?:\/\/)?static\d\.(?:e621|e926)\.net\/data\/(?:sample\/)?(?:[\da-z]{2}\/){2}([\da-z]+)\.[\da-z]+/; // YiffyAPI V2 - v2.yiff.rest const yiffy2Regex = /(?:https?:\/\/)?(?:v2\.yiff\.media|yiff\.media\/V2)\/(?:.*\/)+([\da-z]+)\.[\da-z]+/; + // FemboyFans - static.femboy.fan + const femboyfansRegex = /(?:https?:\/\/)?static\.femboy\.fan\/(?:.*\/)+([\da-z]+)\.[\da-z]+/; let post: Post | null = null, + ffpost: FemboyFans.Post | null = null, method: typeof tried[number] | null = null, sourceOverride: string | Array | undefined, saucePercent = 0, snRateLimited = false; - const tried: Array<"e621" | "yiffy2" | "iqdb" | "saucenao"> = []; + const tried: Array<"e621" | "yiffy2" | "femboyfans" | "iqdb" | "ffiqdb" | "saucenao"> = []; - if (Strings.validateURL(input)) { - if (skipCheck === false) { - const head = await RequestProxy.head(input); - if (head.status !== 200 && head.status !== 204) { - throw new PreCheckError(`A pre-check failed when trying to fetch the image "${input}".\nA \`HEAD\` request returned a non 2XX response (${head.status} ${STATUS_CODES[head.status] || "UNKNOWN"})\n\nThis means we either can't access the file, the server is configured incorrectly, or the file does not exist.`); - } + if (!Strings.validateURL(input)) { + return null; + } + if (skipCheck === false) { + const head = await RequestProxy.head(input); + if (head.status !== 200 && head.status !== 204) { + throw new PreCheckError(`A pre-check failed when trying to fetch the image "${input}".\nA \`HEAD\` request returned a non 2XX response (${head.status} ${STATUS_CODES[head.status] || "UNKNOWN"})\n\nThis means we either can't access the file, the server is configured incorrectly, or the file does not exist.`); } + } + + let match: RegExpExecArray | null; - const yr = yrRegex.exec(input); - if (yr?.[1]) { - const res = await Yiffy.shortener.get(yr[1]); + outer: { + if ((match = yrRegex.exec(input))) { + const res = await Yiffy.shortener.get(match[1]); if (res !== null) { input = res.fullURL; } } - const e6 = e621Regex.exec(input); - const y2 = yiffy2Regex.exec(input); - if (e6?.[1]) { + + femboyfans: if ((match = femboyfansRegex.exec(input))) { + tried.push("femboyfans"); + ffpost = await FemboyFans.getPostByMD5(match[1]); + if (ffpost === null) { + break femboyfans; + } + + method = "femboyfans"; + break outer; + } + + e621: if ((match = e621Regex.exec(input))) { tried.push("e621"); - post = await E621.posts.getByMD5(e6[1]); - if (post !== null) { - method = "e621"; + post = await E621.posts.get(match[1]); + if (post === null) { + break e621; } + + method = "e621"; + break outer; } - if (y2?.[1] && !method) { + yiffy2: if ((match = yiffy2Regex.exec(input))) { tried.push("yiffy2"); - const d = await fetch(`https://v2.yiff.rest/images/${y2[1]}`, { + const yapi = await fetch(`https://v2.yiff.rest/images/${match[1]}`, { headers: { "User-Agent": Config.userAgent } }) .then(res => res.json() as Promise<{ data: JSONResponse; success: true; }>) .catch(() => null); - if (d !== null && d.success === true) { - const s = d.data.sources.find(so => so.includes("e621.net")); - const m = /https:\/\/e621\.net\/posts\/(\d+)/.exec((s || "")); - // dont - if (s && m?.[1]) { - post = await E621.posts.get(Number(m[1])); - if (post !== null) { - method = "e621"; - } - } else { - method = "yiffy2"; - sourceOverride = d.data.sources; - } + if (yapi === null || !yapi.success) { + break yiffy2; + } + + const ffSource = yapi.data.sources.find(s => s.startsWith("https://femboy.fan/posts/")); + const e6Source = yapi.data.sources.find(s => s.startsWith("https://e621.net/posts/")); + + if (ffSource !== undefined && (match = /https:\/\/femboy\.fan\/posts\/(\d+)/.exec(ffSource))) { + ffpost = await FemboyFans.getPost(Number(match[1])); + method = "femboyfans"; + break outer; } + if (e6Source !== undefined && (match = /https:\/\/e621\.net\/posts\/(\d+)/.exec(e6Source))) { + post = await E621.posts.get(Number(match[1])); + method = "e621"; + break outer; + } + method = "yiffy2"; + sourceOverride = yapi.data.sources; + break outer; } - // saucenao is fucky and their api sucks - if (!skipSauceNao && !method) { - const sa = await SauceNAO(input, [29, 40, 41, 42]).catch(() => null); - if (sa !== null && Array.isArray(sa) && sa.length !== 0) { - const top = sa.sort((a, b)=> a.header.similarity - b.header.similarity).find(v => v.header.similarity >= simularity); + saucenao: { + // saucenao is fucky and their api sucks + if (skipSauceNao) { + break saucenao; + } + + tried.push("saucenao"); + const sn = await SauceNAO(input, [29, 40, 41, 42]).catch(() => null); + if (sn !== null && Array.isArray(sn) && sn.length !== 0) { + const top = sn.sort((a, b)=> a.header.similarity - b.header.similarity).find(v => v.header.similarity >= similarity); if (top && top.data.ext_urls.length !== 0) { method = "saucenao"; saucePercent = top.header.similarity; sourceOverride = top.data.ext_urls; + break outer; } } - if (sa === "RateLimited") { + if (sn === "RateLimited") { snRateLimited = true; - } else { - tried.push("saucenao"); - } // so we don't tell the user we both couldn't try & tried saucenao + } } - iqdb: if (!method) { + iqdb: { const img = await RequestProxy.get(input); if (!img.ok) { break iqdb; @@ -113,93 +142,156 @@ export default async function Sauce(input: string, simularity = 80, skipCheck = const content = Buffer.from(await img.response.arrayBuffer()); const type = await fileTypeFromBuffer(content); - if (type && ["image/png", "image/jpeg"].includes(type.mime)) { - const result = await fetch(`https://e621.net/iqdb_queries.json?search[score_cutoff]=${simularity}`, { - method: "POST", - body: new Blob([content], { type: type.mime }) + if (!type || !["image/png", "image/jpeg"].includes(type.mime)) { + break iqdb; + } + + ffiqdb: { + tried.push("ffiqdb"); + const result = await FemboyFans.queryIQDB(content, similarity); + if (result === null) { + break ffiqdb; + } + + method = "ffiqdb"; + ffpost = await FemboyFans.getPost(result.post_id); + saucePercent = result.score; + break outer; + } + + e6iqdb: { + tried.push("iqdb"); + const body = new FormData(); + body.append("file", new Blob([content], { type: type.mime })); + const result = await fetch(`https://e621.net/iqdb_queries.json?search[score_cutoff]=${similarity}`, { + method: "POST", + body, + headers: { + "Authorization": `Basic ${Buffer.from(`${Config.e621User}:${Config.e621APIKey}`).toString("base64")}`, + "User-Agent": Config.userAgent + } }); if (result.status !== 200) { - break iqdb; + break e6iqdb; } - const res = (await result.json()) as Array<{ post_id: number; score: number; }>; + const results = (await result.json()) as Array<{ post_id: number; score: number; }>; + const res = results.sort((a, b) => b.score - a.score).at(0) ?? null; - if (res.length !== 0) { - method = "iqdb"; - post = await E621.posts.get(res[0].post_id); - saucePercent = res[0].score; + if (res === null) { + break e6iqdb; } + + method = "iqdb"; + post = await E621.posts.get(res.post_id); + saucePercent = res.score; + break outer; } } + } - if (post && post.flags.deleted && post.relationships.parent_id !== null) { - const parent = await E621.posts.get(post.relationships.parent_id); - if (parent !== null) { - post = parent; - } + if (post && post.flags.deleted && post.relationships.parent_id !== null) { + const parent = await E621.posts.get(post.relationships.parent_id); + if (parent !== null) { + post = parent; } + } - return { - method, - tried, - post, - saucePercent, - sourceOverride, - snRateLimited, - url: input - }; - } else { + if (ffpost && ffpost.flags.deleted && ffpost.relationships.parent_id !== null) { + const parent = await FemboyFans.getPost(ffpost.relationships.parent_id); + if (parent !== null) { + ffpost = parent; + } + } + + if (method === null) { return null; } + + return { + method, + tried, + post, + ffpost, + saucePercent, + sourceOverride, + snRateLimited, + url: input + }; } export async function directMD5(md5: string) { - let post = await E621.posts.getByMD5(md5), - method: "e621" | "yiffy2" | null = null, + let post: Post | null = null, ffpost: FemboyFans.Post | null = null, + method: "e621" | "yiffy2" | "femboyfans" | null = null, sourceOverride: string | Array | undefined, url: string | null = null; - if (post !== null) { - if (post.flags.deleted && post.relationships.parent_id !== null) { - const parent = await E621.posts.get(post.relationships.parent_id); - if (parent !== null) { - post = parent; + + outer: { + femboyfans: { + ffpost = await FemboyFans.getPostByMD5(md5); + if (ffpost === null) { + break femboyfans; + } + + if (ffpost.flags.deleted && ffpost.relationships.parent_id !== null) { + ffpost = await FemboyFans.getPost(ffpost.relationships.parent_id); } + + method = "femboyfans"; + break outer; } - url = post.file.url; - method = "e621"; - } - if (post === null) { - const yapi = await Yiffy.images.getImage(md5); - if (yapi !== null) { - const e = yapi.sources.find(source => source.startsWith("https://e621.net/posts/")); + e621: { + post = await E621.posts.getByMD5(md5); + if (post === null) { + break e621; + } + + if (post.flags.deleted && post.relationships.parent_id !== null) { + post = await E621.posts.get(post.relationships.parent_id); + } + + method = "e621"; + break outer; + } + + yiffy2: { + const yapi = await Yiffy.images.getImage(md5); + + if (yapi === null) { + break yiffy2; + } + + const ffSource = yapi.sources.find(s => s.startsWith("https://femboy.fan/posts/")); + const e6Source = yapi.sources.find(s => s.startsWith("https://e621.net/posts/")); + let match: RegExpExecArray | null; - if (e) { - if ((match = /https:\/\/e621\.net\/posts\/(\d+)/.exec(e))) { - post = await E621.posts.get(Number(match[1])); - if (post !== null) { - if (post.flags.deleted && post.relationships.parent_id !== null) { - const parent = await E621.posts.get(post.relationships.parent_id); - if (parent !== null) { - post = parent; - } - } - url = post.file.url; - method = "e621"; - } - } else { - url = yapi.url; - method = "yiffy2"; - sourceOverride = yapi.sources; - post = null; - } + if (ffSource !== undefined && (match = /https:\/\/femboy\.fan\/posts\/(\d+)/.exec(ffSource))) { + ffpost = await FemboyFans.getPost(Number(match[1])); + method = "femboyfans"; + break outer; + } + if (e6Source !== undefined && (match = /https:\/\/e621\.net\/posts\/(\d+)/.exec(e6Source))) { + post = await E621.posts.get(Number(match[1])); + method = "e621"; + break outer; } + + url = yapi.url; + method = "yiffy2"; + sourceOverride = yapi.sources; + break outer; } } + if (method === null) { + return null; + } + return { method, + ffpost, post, sourceOverride, url diff --git a/src/util/req/FemboyFans.ts b/src/util/req/FemboyFans.ts new file mode 100644 index 0000000..4e322a2 --- /dev/null +++ b/src/util/req/FemboyFans.ts @@ -0,0 +1,61 @@ +/// +import Config from "../../config/index.js"; +import Logger from "@uwu-codes/logger"; +import { fileTypeFromBuffer } from "file-type"; + +export default class FemboyFans { + static AUTH = `Basic ${Buffer.from(`${Config.femboyFansUser}:${Config.femboyFansKey}`).toString("base64")}`; + static URL = "https://femboy.fan"; + + static async getPost(id: number): Promise { + const result = await fetch(`${this.URL}/posts/${id}`, { + headers: { + "Authorization": this.AUTH, + "User-Agent": Config.userAgent + } + }); + if (result.status !== 200) { + Logger.getLogger("FemboyFans#getPost").error(`Unexpected ${result.status} ${result.statusText}:`); + Logger.getLogger("FemboyFans#getPost").error(await result.text()); + return null; + } + return (result.json() as Promise); + } + + static async getPostByMD5(md5: string): Promise { + const result = await fetch(`${this.URL}/posts.json?md5=${md5}`, { + headers: { + "Authorization": this.AUTH, + "User-Agent": Config.userAgent + } + }); + if (result.status !== 200) { + Logger.getLogger("FemboyFans#getPostByMD5").error(`Unexpected ${result.status} ${result.statusText}:`); + Logger.getLogger("FemboyFans#getPostByMD5").error(await result.text()); + return null; + } + return (result.json() as Promise>).then(([post]) => post); + } + + static async queryIQDB(img: Buffer, similarity = 60): Promise<{ post_id: number; score: number; } | null> { + const type = await fileTypeFromBuffer(img); + if (type === undefined) { + return null; + } + const body = new FormData(); + body.append("file", new Blob([img], { type: type.mime })); + const result = await fetch(`${this.URL}/posts/iqdb.json?search[score_cutoff]=${similarity}`, { + method: "POST", + body + }); + + if (result.status !== 200) { + Logger.getLogger("FemboyFans#queryIQDB").error(`Unexpected ${result.status} ${result.statusText}:`); + Logger.getLogger("FemboyFans#queryIQDB").error(await result.text()); + return null; + } + + const res = (await result.json()) as Array<{ post_id: number; score: number; }>; + return res.sort((a, b) => b.score - a.score).at(0) ?? null; + } +}