From 423fc69f9fd3015efd68aa1da74e78bd1cfe0ca9 Mon Sep 17 00:00:00 2001
From: Tarrgon <61888458+Tarrgon@users.noreply.github.com>
Date: Fri, 6 Jun 2025 12:05:32 -0400
Subject: [PATCH] Calculate md5 hashes of uploaded files

Discord does some weird stuff to file data, and we have to reverse it.
---
 src/events/handle-message.ts | 21 ++++++++----
 src/utils/file-utils.ts      | 63 ++++++++++++++++++++++++++++++++++++
 src/utils/index.ts           |  1 +
 3 files changed, 78 insertions(+), 7 deletions(-)
 create mode 100644 src/utils/file-utils.ts

diff --git a/src/events/handle-message.ts b/src/events/handle-message.ts
index 96119ab..efc4360 100644
--- a/src/events/handle-message.ts
+++ b/src/events/handle-message.ts
@@ -5,7 +5,7 @@ import { getE621Post, getE621PostByMd5, getPostUrl, PostAction, spoilerOrBlackli
 import { Database } from '../shared/Database';
 import { logDeletion, logEdit } from '../utils/message-logger';
 import { isEdited } from '../utils/message-utils';
-import { blipIDRegex, channelIgnoresLinks, channelIsInStaffCategory, channelIsSafe, commentIDRegex, forumTopicIDRegex, poolIDRegex, postIDRegex, recordIDRegex, searchLinkRegex, setIDRegex, takedownIDRegex, ticketIDRegex, userIDRegex, wikiLinkRegex } from '../utils';
+import { ALLOWED_MIMETYPES, blipIDRegex, calculateMD5FromURL, channelIgnoresLinks, channelIsInStaffCategory, channelIsSafe, commentIDRegex, forumTopicIDRegex, poolIDRegex, postIDRegex, recordIDRegex, searchLinkRegex, setIDRegex, takedownIDRegex, ticketIDRegex, userIDRegex, wikiLinkRegex } from '../utils';
 
 export type Message<InGuild extends boolean = boolean> = OmitPartialGroupDMChannel<DiscordMessage<InGuild>>;
 export type Partial = OmitPartialGroupDMChannel<PartialMessage>;
@@ -73,14 +73,21 @@ export async function handleMessageCreate(message: Message) {
     const match = md5Regex.exec(attachment.name);
     md5Regex.lastIndex = 0;
 
-    if (match) {
-      const post = await getE621PostByMd5(match[1]);
+    let md5;
 
-      if (post) {
-        if (await blacklistIfNecessary(message, [post])) return;
+    if (match) md5 = match[1];
+    else if (ALLOWED_MIMETYPES.includes(attachment.contentType!)) {
+      md5 = await calculateMD5FromURL(attachment.url);
+    }
 
-        responses.push(`<${getPostUrl(post)}>`);
-      }
+    if (!md5) continue;
+
+    const post = await getE621PostByMd5(md5);
+
+    if (post) {
+      if (await blacklistIfNecessary(message, [post])) return;
+
+      responses.push(`<${getPostUrl(post)}>`);
     }
   }
 
diff --git a/src/utils/file-utils.ts b/src/utils/file-utils.ts
new file mode 100644
index 0000000..7201701
--- /dev/null
+++ b/src/utils/file-utils.ts
@@ -0,0 +1,63 @@
+import crypto from 'crypto';
+
+const DISCORD_PNG_ADDITIONAL_BYTE_LENGTH = 26;
+const END_PNG_BYTES = 12;
+
+const DISCORD_JPG_START_OFFSET = 3;
+const DISCORD_JPG_REMOVE_BYTE_LENGTH = 23;
+const DISCORD_JPG_REINSERT = Buffer.from([0xE0, 0x00, 0x10, 0x4A, 0x46, 0x49, 0x46, 0x00, 0x01, 0x01, 0x00, 0x00, 0x01, 0x00, 0x01, 0x00, 0x00]);
+const DISCORD_JPG_REINSERT_BYTE_LENGTH = DISCORD_JPG_REINSERT.byteLength;
+
+export const ALLOWED_MIMETYPES = ['image/png', 'image/jpg', 'image/jpeg', 'image/gif', 'video/mp4', 'video/webm'];
+
+export function calculateMD5(data: Buffer): string {
+  return crypto.createHash('md5').update(data).digest('hex');
+}
+
+export async function downloadFile(url: string): Promise<Buffer | null> {
+  try {
+    const res = await fetch(url);
+
+    const mimeType = res.headers.get('Content-Type')!;
+
+    if (!ALLOWED_MIMETYPES.includes(mimeType)) return null;
+
+    const data = await res.arrayBuffer();
+
+    let finalData: Buffer;
+
+    if (mimeType == 'image/png') {
+      const startOffset = data.byteLength - DISCORD_PNG_ADDITIONAL_BYTE_LENGTH - END_PNG_BYTES;
+      const correctedData = Buffer.alloc(data.byteLength - DISCORD_PNG_ADDITIONAL_BYTE_LENGTH);
+      const buff = Buffer.from(data);
+      buff.copy(correctedData, 0, 0, startOffset);
+      buff.copy(correctedData, startOffset, startOffset + DISCORD_PNG_ADDITIONAL_BYTE_LENGTH);
+      finalData = correctedData;
+    } else if (mimeType == 'image/jpg' || mimeType == 'image/jpeg') {
+      const correctedData = Buffer.alloc(data.byteLength - DISCORD_JPG_REMOVE_BYTE_LENGTH + DISCORD_JPG_REINSERT_BYTE_LENGTH);
+      const buff = Buffer.from(data);
+      buff.copy(correctedData, 0, 0, DISCORD_JPG_START_OFFSET);
+      DISCORD_JPG_REINSERT.copy(correctedData, DISCORD_JPG_START_OFFSET);
+      buff.copy(correctedData, DISCORD_JPG_REMOVE_BYTE_LENGTH - DISCORD_JPG_START_OFFSET, DISCORD_JPG_START_OFFSET + DISCORD_JPG_REMOVE_BYTE_LENGTH);
+      finalData = correctedData;
+    } else {
+      finalData = Buffer.from(data);
+    }
+
+    return finalData;
+  } catch (e) {
+    console.error(e);
+    return null;
+  }
+}
+
+export async function calculateMD5FromURL(url: string): Promise<string | null> {
+  try {
+    const file = await downloadFile(url);
+    if (!file) return null;
+    return calculateMD5(file);
+  } catch (e) {
+    console.error(e);
+    return null;
+  }
+}
\ No newline at end of file
diff --git a/src/utils/index.ts b/src/utils/index.ts
index b981bdb..9ee0f48 100644
--- a/src/utils/index.ts
+++ b/src/utils/index.ts
@@ -5,6 +5,7 @@ export * from './ban-utils';
 export * from './channel-utils';
 export * from './commands';
 export * from './e621-utils';
+export * from './file-utils';
 export * from './interaction-utils';
 export * from './message-logger';
 export * from './message-utils';