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

@@ -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,8 +73,16 @@ 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 (match) md5 = match[1];
+    else if (ALLOWED_MIMETYPES.includes(attachment.contentType!)) {
+      md5 = await calculateMD5FromURL(attachment.url);
+    }
+
+    if (!md5) continue;
+
+    const post = await getE621PostByMd5(md5);
 
     if (post) {
       if (await blacklistIfNecessary(message, [post])) return;
@@ -82,7 +90,6 @@ export async function handleMessageCreate(message: Message) {
       responses.push(`<${getPostUrl(post)}>`);
     }
   }
-  }
 
   if (responses.length > 0) {
     await message.reply(responses.join('\n'));

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;
+  }
+}

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';