From a98484c38bc92055e3c660a86687ee0425bc5c5b Mon Sep 17 00:00:00 2001
From: Adhiraj Singh <adhirajsingh1001@gmail.com>
Date: Thu, 21 Apr 2022 00:52:17 +0530
Subject: [PATCH] refactor: split "downloadContentFromMessage" to
 "downloadEncryptedContent""

---
 src/Types/Message.ts        |  6 ++++++
 src/Utils/messages-media.ts | 24 +++++++++++++++++++-----
 2 files changed, 25 insertions(+), 5 deletions(-)

diff --git a/src/Types/Message.ts b/src/Types/Message.ts
index a83afc7..bc55d89 100644
--- a/src/Types/Message.ts
+++ b/src/Types/Message.ts
@@ -178,3 +178,9 @@ export type WAMessageUpdate = { update: Partial<WAMessage>, key: proto.IMessageK
 export type WAMessageCursor = { before: WAMessageKey | undefined } | { after: WAMessageKey | undefined }
 
 export type MessageUserReceiptUpdate = { key: proto.IMessageKey, receipt: MessageUserReceipt }
+
+export type MediaDecryptionKeyInfo = {
+    iv: Buffer
+    cipherKey: Buffer
+    macKey?: Buffer
+}
\ No newline at end of file
diff --git a/src/Utils/messages-media.ts b/src/Utils/messages-media.ts
index 48bf4cc..03b7ccb 100644
--- a/src/Utils/messages-media.ts
+++ b/src/Utils/messages-media.ts
@@ -11,7 +11,7 @@ import type { Logger } from 'pino'
 import { Readable, Transform } from 'stream'
 import { URL } from 'url'
 import { DEFAULT_ORIGIN, MEDIA_PATH_MAP } from '../Defaults'
-import { CommonSocketConfig, DownloadableMessage, MediaConnInfo, MediaType, MessageType, WAGenericMediaMessage, WAMediaUpload, WAMediaUploadFunction, WAMessageContent, WAProto } from '../Types'
+import { CommonSocketConfig, DownloadableMessage, MediaConnInfo, MediaDecryptionKeyInfo, MediaType, MessageType, WAGenericMediaMessage, WAMediaUpload, WAMediaUploadFunction, WAMessageContent, WAProto } from '../Types'
 import { hkdf } from './crypto'
 import { generateMessageID } from './generics'
 
@@ -60,7 +60,7 @@ export const hkdfInfoKey = (type: MediaType) => {
 }
 
 /** generates all the keys required to encrypt/decrypt & sign a media message */
-export function getMediaKeys(buffer, mediaType: MediaType) {
+export function getMediaKeys(buffer: Uint8Array | string, mediaType: MediaType): MediaDecryptionKeyInfo {
 	if(typeof buffer === 'string') {
 		buffer = Buffer.from(buffer.replace('data:;base64,', ''), 'base64')
 	}
@@ -186,6 +186,7 @@ export const toBuffer = async(stream: Readable) => {
 		buff = Buffer.concat([ buff, chunk ])
 	}
 
+	stream.destroy()
 	return buff
 }
 
@@ -343,12 +344,26 @@ export type MediaDownloadOptions = {
     endByte?: number
 }
 
-export const downloadContentFromMessage = async(
+export const downloadContentFromMessage = (
 	{ mediaKey, directPath, url }: DownloadableMessage,
 	type: MediaType,
-	{ startByte, endByte }: MediaDownloadOptions = { }
+	opts: MediaDownloadOptions = { }
 ) => {
 	const downloadUrl = url || `https://${DEF_HOST}${directPath}`
+	const keys = getMediaKeys(mediaKey, type)
+
+	return downloadEncryptedContent(downloadUrl, keys, opts)
+}
+
+/**
+ * Decrypts and downloads an AES256-CBC encrypted file given the keys.
+ * Assumes the SHA256 of the plaintext is appended to the end of the ciphertext
+ * */
+export const downloadEncryptedContent = async(
+	downloadUrl: string,
+	{ cipherKey, iv }: MediaDecryptionKeyInfo,
+	{ startByte, endByte }: MediaDownloadOptions = { }
+) => {
 	let bytesFetched = 0
 	let startChunk = 0
 	let firstBlockIsIV = false
@@ -386,7 +401,6 @@ export const downloadContentFromMessage = async(
 	)
 
 	let remainingBytes = Buffer.from([])
-	const { cipherKey, iv } = getMediaKeys(mediaKey, type)
 
 	let aes: Crypto.Decipher