From 9d8aa67f22d5c9e1fba6205b779f9acb41b1eb54 Mon Sep 17 00:00:00 2001
From: Adhiraj Singh <adhirajsingh1001@gmail.com>
Date: Thu, 2 Mar 2023 22:43:31 +0530
Subject: [PATCH] feat: process pollupdate automatically

---
 src/Socket/chats.ts          |  1 +
 src/Utils/process-message.ts | 67 +++++++++++++++++++++++++++++++++---
 2 files changed, 64 insertions(+), 4 deletions(-)

diff --git a/src/Socket/chats.ts b/src/Socket/chats.ts
index 40fe368..da53448 100644
--- a/src/Socket/chats.ts
+++ b/src/Socket/chats.ts
@@ -764,6 +764,7 @@ export const makeChatsSocket = (config: SocketConfig) => {
 					keyStore: authState.keys,
 					logger,
 					options: config.options,
+					getMessage: config.getMessage,
 				}
 			)
 		])
diff --git a/src/Utils/process-message.ts b/src/Utils/process-message.ts
index ad4bea8..294a63f 100644
--- a/src/Utils/process-message.ts
+++ b/src/Utils/process-message.ts
@@ -1,7 +1,7 @@
 import { AxiosRequestConfig } from 'axios'
 import type { Logger } from 'pino'
 import { proto } from '../../WAProto'
-import { AuthenticationCreds, BaileysEventEmitter, Chat, GroupMetadata, ParticipantAction, SignalKeyStoreWithTransaction, WAMessageStubType } from '../Types'
+import { AuthenticationCreds, BaileysEventEmitter, Chat, GroupMetadata, ParticipantAction, SignalKeyStoreWithTransaction, SocketConfig, WAMessageStubType } from '../Types'
 import { getContentType, normalizeMessageContent } from '../Utils/messages'
 import { areJidsSameUser, isJidBroadcast, isJidStatusBroadcast, jidNormalizedUser } from '../WABinary'
 import { aesDecryptGCM, hmacSign } from './crypto'
@@ -13,8 +13,9 @@ type ProcessMessageContext = {
 	creds: AuthenticationCreds
 	keyStore: SignalKeyStoreWithTransaction
 	ev: BaileysEventEmitter
+	getMessage: SocketConfig['getMessage']
 	logger?: Logger
-	options: AxiosRequestConfig<any>
+	options: AxiosRequestConfig<{}>
 }
 
 const REAL_MSG_STUB_TYPES = new Set([
@@ -36,7 +37,14 @@ export const cleanMessage = (message: proto.IWebMessageInfo, meId: string) => {
 	const content = normalizeMessageContent(message.message)
 	// if the message has a reaction, ensure fromMe & remoteJid are from our perspective
 	if(content?.reactionMessage) {
-		const msgKey = content.reactionMessage.key!
+		normaliseKey(content.reactionMessage.key!)
+	}
+
+	if(content?.pollUpdateMessage) {
+		normaliseKey(content.pollUpdateMessage.pollCreationMessageKey!)
+	}
+
+	function normaliseKey(msgKey: proto.IMessageKey) {
 		// if the reaction is from another user
 		// we've to correctly map the key to this user's perspective
 		if(!message.key.fromMe) {
@@ -69,6 +77,7 @@ export const isRealMessage = (message: proto.IWebMessageInfo, meId: string) => {
 	&& hasSomeContent
 	&& !normalizedContent?.protocolMessage
 	&& !normalizedContent?.reactionMessage
+	&& !normalizedContent?.pollUpdateMessage
 }
 
 export const shouldIncrementChatUnread = (message: proto.IWebMessageInfo) => (
@@ -136,6 +145,9 @@ export function decryptPollVote(
 
 	function toBinary(txt: string) {
 		return Buffer.from(txt)
+	}
+}
+
 const processMessage = async(
 	message: proto.IWebMessageInfo,
 	{
@@ -144,7 +156,8 @@ const processMessage = async(
 		creds,
 		keyStore,
 		logger,
-		options
+		options,
+		getMessage
 	}: ProcessMessageContext
 ) => {
 	const meId = creds.me!.id
@@ -321,6 +334,52 @@ const processMessage = async(
 			emitGroupUpdate({ inviteCode: code })
 			break
 		}
+	} else if(content?.pollUpdateMessage) {
+		const creationMsgKey = content.pollUpdateMessage.pollCreationMessageKey!
+		// we need to fetch the poll creation message to get the poll enc key
+		const pollMsg = await getMessage(creationMsgKey)
+		if(pollMsg) {
+			const meIdNormalised = jidNormalizedUser(meId)
+			const pollCreatorJid = getKeyAuthor(creationMsgKey, meIdNormalised)
+			const voterJid = getKeyAuthor(message.key!, meIdNormalised)
+			const pollEncKey = pollMsg.messageContextInfo?.messageSecret!
+
+			try {
+				const voteMsg = decryptPollVote(
+					content.pollUpdateMessage.vote!,
+					{
+						pollEncKey,
+						pollCreatorJid,
+						pollMsgId: creationMsgKey.id!,
+						voterJid,
+					}
+				)
+				ev.emit('messages.update', [
+					{
+						key: creationMsgKey,
+						update: {
+							pollUpdates: [
+								{
+									pollUpdateMessageKey: message.key,
+									vote: voteMsg,
+									senderTimestampMs: message.messageTimestamp,
+								}
+							]
+						}
+					}
+				])
+			} catch(err) {
+				logger?.warn(
+					{ err, creationMsgKey },
+					'failed to decrypt poll vote'
+				)
+			}
+		} else {
+			logger?.warn(
+				{ creationMsgKey },
+				'poll creation message not found, cannot decrypt update'
+			)
+		}
 	}
 
 	if(Object.keys(chat).length > 1) {