Stream uploads + downloads + allow for remote url uploads

- Switch to using got - Use encryption/decryption streams for speed & lesser memory consumption - Allow for stream based download & simultaneous upload of media
2021-01-13 22:48:28 +05:30
parent 500805236a
commit 0344d6336c
19 changed files with 501 additions and 146 deletions
@@ -1,16 +1,20 @@
 import * as Crypto from 'crypto'
+import { Readable, Transform } from 'stream'
 import HKDF from 'futoin-hkdf'
 import Jimp from 'jimp'
-import {promises as fs} from 'fs'
+import {createReadStream, createWriteStream, promises as fs, WriteStream} from 'fs'
 import { exec } from 'child_process'
-import {platform, release} from 'os'
+import {platform, release, tmpdir} from 'os'
 import HttpsProxyAgent from 'https-proxy-agent'
 import { URL } from 'url'
 import { Agent } from 'https'
 import Decoder from '../Binary/Decoder'
-import { MessageType, HKDFInfoKeys, MessageOptions, WAChat, WAMessageContent, BaileysError, WAMessageProto, TimedOutError, CancelledError, WAGenericMediaMessage, WAMessage, WAMessageKey } from './Constants'
+import { MessageType, HKDFInfoKeys, MessageOptions, WAChat, WAMessageContent, BaileysError, WAMessageProto, TimedOutError, CancelledError, WAGenericMediaMessage, WAMessage, WAMessageKey, DEFAULT_ORIGIN, WAMediaUpload } from './Constants'
 import KeyedDB from '@adiwajshing/keyed-db'
-import { Response } from 'node-fetch'
+import got, { Options, Response } from 'got'
+import { join } from 'path'
+import { IAudioMetadata } from 'music-metadata'
+

 const platformMap = {
    'aix': 'AIX',
@@ -222,9 +226,10 @@ const extractVideoThumb = async (
        })
    }) as Promise<void>

-export const compressImage = async (buffer: Buffer) => {
-    const jimp = await Jimp.read (buffer)
-    return jimp.resize(48, 48).getBufferAsync (Jimp.MIME_JPEG)
+export const compressImage = async (bufferOrFilePath: Buffer | string) => {
+    const jimp = await Jimp.read(bufferOrFilePath as any)
+    const result = await jimp.resize(48, 48).getBufferAsync(Jimp.MIME_JPEG)
+    return result
 }
 export const generateProfilePicture = async (buffer: Buffer) => {
    const jimp = await Jimp.read (buffer)
@@ -241,42 +246,131 @@ export const mediaMessageSHA256B64 = (message: WAMessageContent) => {
    const media = Object.values(message)[0] as WAGenericMediaMessage
    return media?.fileSha256 && Buffer.from(media.fileSha256).toString ('base64')
 }
-export async function getAudioDuration (buffer: Buffer) {
+export async function getAudioDuration (buffer: Buffer | string) {
    const musicMetadata = await import ('music-metadata')
-    const metadata = await musicMetadata.parseBuffer (buffer, null, {duration: true});
+    let metadata: IAudioMetadata
+    if(Buffer.isBuffer(buffer)) {
+        metadata = await musicMetadata.parseBuffer(buffer, null, { duration: true })
+    } else {
+        const rStream = createReadStream(buffer)
+        metadata = await musicMetadata.parseStream(rStream, null, { duration: true })
+        rStream.close()
+    }
    return metadata.format.duration;
 }
-
+export const toReadable = (buffer: Buffer) => {
+    const readable = new Readable({ read: () => {} })
+    readable.push(buffer)
+    readable.push(null)
+    return readable
+}
+export const getStream = async (item: WAMediaUpload) => {
+    if(Buffer.isBuffer(item)) return { stream: toReadable(item), type: 'buffer' }
+    if(item.url.toString().startsWith('http://') || item.url.toString().startsWith('https://')) {
+        return { stream: await getGotStream(item.url), type: 'remote' }
+    }
+    return { stream: createReadStream(item.url), type: 'file' }
+}
 /** generates a thumbnail for a given media, if required */
-export async function generateThumbnail(buffer: Buffer, mediaType: MessageType, info: MessageOptions) {
+export async function generateThumbnail(file: string, mediaType: MessageType, info: MessageOptions) {
    if ('thumbnail' in info) {
        // don't do anything if the thumbnail is already provided, or is null
        if (mediaType === MessageType.audio) {
            throw new Error('audio messages cannot have thumbnails')
        }
    } else if (mediaType === MessageType.image) {
-        const buff = await compressImage (buffer)
+        const buff = await compressImage(file)
        info.thumbnail = buff.toString('base64')
    } else if (mediaType === MessageType.video) {
-        const filename = './' + randomBytes(5).toString('hex') + '.mp4'
-        const imgFilename = filename + '.jpg'
-        await fs.writeFile(filename, buffer)
+        const imgFilename = join(tmpdir(), generateMessageID() + '.jpg')
        try {
-            await extractVideoThumb(filename, imgFilename, '00:00:00', { width: 48, height: 48 })
+            await extractVideoThumb(file, imgFilename, '00:00:00', { width: 48, height: 48 })
            const buff = await fs.readFile(imgFilename)
            info.thumbnail = buff.toString('base64')
            await fs.unlink(imgFilename)
        } catch (err) {
            console.log('could not generate video thumb: ' + err)
        }
-        await fs.unlink(filename)
+    }
+}
+export const getGotStream = async(url: string | URL, options: Options & { isStream?: true } = {}) => {
+    const fetched = got.stream(url, options)
+    await new Promise((resolve, reject) => {
+        fetched.once('response', ({statusCode: status}: Response) => {
+            if (status >= 400) {
+                reject(new BaileysError (
+                    'Invalid code (' + status + ') returned', 
+                    { status }
+                ))
+            } else {
+                resolve(undefined)
+            }
+        })
+    })
+    return fetched
+} 
+export const encryptedStream = async(media: WAMediaUpload, mediaType: MessageType) => {
+    const { stream, type } = await getStream(media)
+
+    const mediaKey = randomBytes(32)
+    const {cipherKey, iv, macKey} = getMediaKeys(mediaKey, mediaType)
+    // random name
+    const encBodyPath = join(tmpdir(), mediaType + generateMessageID() + '.enc')
+    const encWriteStream = createWriteStream(encBodyPath)
+    let bodyPath: string
+    let writeStream: WriteStream
+    if(type === 'file') {
+        bodyPath = (media as any).url
+    } else {
+        bodyPath = join(tmpdir(), mediaType + generateMessageID())
+        writeStream = createWriteStream(bodyPath)
+    }
+    
+    let fileLength = 0
+    const aes = Crypto.createCipheriv('aes-256-cbc', cipherKey, iv)
+    let hmac = Crypto.createHmac('sha256', macKey).update(iv)
+    let sha256Plain = Crypto.createHash('sha256')
+    let sha256Enc = Crypto.createHash('sha256')
+
+    const onChunk = (buff: Buffer) => {
+        sha256Enc = sha256Enc.update(buff)
+        hmac = hmac.update(buff)
+        encWriteStream.write(buff)
+    }
+    for await(const data of stream) {
+        fileLength += data.length
+        sha256Plain = sha256Plain.update(data)
+        writeStream && writeStream.write(data)
+        onChunk(aes.update(data))
+    }
+    onChunk(aes.final())
+
+    const mac = hmac.digest().slice(0, 10)
+    sha256Enc = sha256Enc.update(mac)
+    
+    const fileSha256 = sha256Plain.digest()
+    const fileEncSha256 = sha256Enc.digest()
+    
+    encWriteStream.write(mac)
+    encWriteStream.close()
+
+    writeStream && writeStream.close()
+
+    return {
+        mediaKey,
+        encBodyPath,
+        bodyPath,
+        mac,
+        fileEncSha256,
+        fileSha256,
+        fileLength
    }
 }
 /**
 * Decode a media message (video, image, document, audio) & return decrypted buffer
 * @param message the media message you want to decode
 */
-export async function decodeMediaMessageBuffer(message: WAMessageContent, fetchRequest: (host: string, method: string) => Promise<Response>) {
+export async function decryptMediaMessageBuffer(message: WAMessageContent): Promise<Readable> {
    /* 
        One can infer media type from the key in the message
        it is usually written as [mediaType]Message. Eg. imageMessage, audioMessage etc.
@@ -289,7 +383,11 @@ export async function decodeMediaMessageBuffer(message: WAMessageContent, fetchR
        throw new BaileysError('cannot decode text message', message)
    }
    if (type === MessageType.location || type === MessageType.liveLocation) {
-        return Buffer.from(message[type].jpegThumbnail)
+        const buffer = Buffer.from(message[type].jpegThumbnail)
+        const readable = new Readable({ read: () => {} })
+        readable.push(buffer)
+        readable.push(null)
+        return readable
    }
    let messageContent: WAGenericMediaMessage
    if (message.productMessage) {
@@ -299,41 +397,39 @@ export async function decodeMediaMessageBuffer(message: WAMessageContent, fetchR
    } else {
        messageContent = message[type]
    }
-    
    // download the message
-    const fetched = await fetchRequest(messageContent.url, 'GET')
-    if (fetched.status >= 400) {
-        throw new BaileysError ('Invalid code (' + fetched.status + ') returned. File has possibly been deleted from WA servers. Run `client.updateMediaMessage()` to refresh the url', {status: fetched.status})
-    }
-    const buffer = await fetched.buffer()
+    const fetched = await getGotStream(messageContent.url, {
+        headers: { Origin: DEFAULT_ORIGIN }
+    })
+    let remainingBytes = Buffer.from([])
+    const { cipherKey, iv } = getMediaKeys(messageContent.mediaKey, type)
+    const aes = Crypto.createDecipheriv("aes-256-cbc", cipherKey, iv)

-    const decryptedMedia = (type: MessageType) => {
-        // get the keys to decrypt the message
-        const mediaKeys = getMediaKeys(messageContent.mediaKey, type) //getMediaKeys(Buffer.from(messageContent.mediaKey, 'base64'), type)
-        // first part is actual file
-        const file = buffer.slice(0, buffer.length - 10)
-        // last 10 bytes is HMAC sign of file
-        const mac = buffer.slice(buffer.length - 10, buffer.length)
-        // sign IV+file & check for match with mac
-        const testBuff = Buffer.concat([mediaKeys.iv, file])
-        const sign = hmacSign(testBuff, mediaKeys.macKey).slice(0, 10)
-        // our sign should equal the mac
-        if (!sign.equals(mac)) throw new Error()
-        
-        return aesDecryptWithIV(file, mediaKeys.cipherKey, mediaKeys.iv) // decrypt media
-    }
-    const allTypes = [type, ...Object.keys(HKDFInfoKeys)]
-    for (let i = 0; i < allTypes.length;i++) {
-        try {
-            const decrypted = decryptedMedia (allTypes[i] as MessageType)
-            
-            if (i > 0) { console.log (`decryption of ${type} media with HKDF key of ${allTypes[i]}`) }
-            return decrypted
-        } catch {
-            if (i === 0) { console.log (`decryption of ${type} media with original HKDF key failed`) }
-        }
-    }
-    throw new BaileysError('Decryption failed, HMAC sign does not match', {status: 400})
+    const output = new Transform({
+        transform(chunk, _, callback) {
+            let data = Buffer.concat([remainingBytes, chunk])
+            const decryptLength =
+                Math.floor(data.length / 16) * 16
+            remainingBytes = data.slice(decryptLength)
+            data = data.slice(0, decryptLength)
+
+            try {
+                this.push(aes.update(data))
+                callback()
+            } catch(error) {
+                callback(error)
+            }  
+        },
+        final(callback) {
+            try {
+                this.push(aes.final())
+                callback()
+            } catch(error) {
+                callback(error)
+            }
+        },
+    })
+    return fetched.pipe(output, { end: true })
 }
 export function extensionForMediaMessage(message: WAMessageContent) {
    const getExtension = (mimetype: string) => mimetype.split(';')[0].split('/')[1]