class WhatsAppAgent(object): def __init__(self): self.driver = WhatsAPIDriver(client=WhatsAppAgentConstants.BROWSER) logger.info("Waiting for successful QR Code scan...") self.driver.wait_for_login() logger.info("Logged in!") def get_all_chats(self): return self.driver.get_all_chats() def get_all_msgs_in_chat(self, chat=None): if chat is None: return [] return self.driver.get_all_messages_in_chat(chat) def get_msg_text(self, message=None): if message is None: return "" return message.safe_content def get_unread_msgs(self): for contact in self.driver.get_unread(): for message in contact.messages: if isinstance(message, Message): whatsAppMsg = WhatsAppAgentMessage( chatId=contact.chat.id, chatName=contact.chat.name, senderId="", senderName=message.sender.get_safe_name(), msgType="", msgData=message.safe_content, msgTime=message.timestamp) logger.debug("New message: %s", whatsAppMsg.readable_format()) def run(self): while True: time.sleep(WhatsAppAgentConstants.NEW_MSG_POLL) self.get_unread_msgs()
class WhatsappBot: def __init__(self, auto_run=False, auto_long_run=False, headless=False): self._chromedriver = os.environ.get( 'CHROMEDRIVE_PATH', os.path.join(os.path.dirname(__file__), "chromedriver")) self._profile_path = os.path.join(os.path.dirname(__file__), "chromeprofile") self._headless = headless self._driver = WhatsAPIDriver( username="******", client="chrome", profile=self._profile_path, executable_path=self._chromedriver, headless=self._headless, chrome_options=[ "user-agent=Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/83.0.4103.116 Safari/537.36", ], heroku=self._headless) self._bot_actions = BotActions(self._driver) self._thread_pool = ThreadPoolExecutor(max_workers=2) self._threaded_users = [] if auto_long_run: self.keep_running() elif auto_run: self.run() def run(self) -> None: """ • Check if there's some user with status == 6 (if yes, generate sticker pack) • Listen to new messages Here we can either keep listening to new messages (long_run == True) or check just once. """ # Check for new messages self.check_for_unread_messages() # Check for users waiting the pack user_in_last_stage = User.find_users_in_last_stage() for user in user_in_last_stage: if user[0] not in self._threaded_users: self._thread_pool.submit(self.create_sticker_pack, user) self._threaded_users.append(user[0]) logging.info(f"{user[0]} added to queue.") else: logging.info(f"{user[0]} already in queue.") def keep_running(self) -> None: """ Keeps running with self.run() :return: """ while True: if not self._driver.is_logged_in(): self._driver.screenshot('scrsht.png') self._driver.wait_for_login() try: self.run() sleep(2) except TypeError as err: logging.critical(err) logging.critical("---RESTARTING---") def create_sticker_pack(self, user_info: tuple) -> bool: """Create a sticker pack using StickerSet class.""" wa_chat_id = user_info[0] package_title = user_info[1] tg_chat_id = user_info[2] logging.info(f"Running for {wa_chat_id}") # Get stickers messages stickers = self.list_user_unread_stickers(wa_chat_id) # Create sticker set sticker_set = StickerSet(tg_chat_id) name = sticker_set.create_new_sticker_set( package_title, stickers[0].save_media_buffer(True)) if not name: logging.error(f"Can't create {wa_chat_id} pack: name = {name}") return False # Populate sticker set for sticker in stickers[1:]: stts = sticker_set.add_sticker_to_set( tg_chat_id, name, sticker.save_media_buffer(True)) logging.info(f"Added a sticker to {name}: {stts}") # Send confirmation self._bot_actions.confirmation(wa_chat_id, name) logging.info(f"Finished {wa_chat_id}") # Remove user from threading if wa_chat_id in self._threaded_users: self._threaded_users.remove(wa_chat_id) return True def check_for_unread_messages(self) -> None: """ Check for unread messages and call actions :return: """ unread = self._driver.get_unread() for msg_group in unread: print(f'Message from <{msg_group.chat.id}>.') for message in msg_group.messages: self.process_incoming_message(message) def process_incoming_message(self, message: Message) -> None: # Message properties: https://gist.github.com/hellmrf/6e06fc374bb43de0868fbb57c223aecd if message.type == 'chat': print( f"[{message.chat_id} {message.timestamp}]: {message.content}") user_is_threaded = message.chat_id in self._threaded_users self.treat_message(message.chat_id, message.content, queued=user_is_threaded) elif User.get_stage(message.chat_id) == 0: self.treat_message(message.chat_id, "Hello") def list_user_unread_stickers(self, chat_id: str) -> List[Message]: messages: List[Message] = self._driver.get_all_messages_in_chat( chat_id) stickers = [ message for message in messages if message.type == 'sticker' ] return stickers @staticmethod def upload_sticker_from_message(tg_user_id: int, sticker_message: Message) -> str: sticker = sticker_message.save_media_buffer(True) return StickerSet.upload_sticker(tg_user_id, sticker) def treat_message(self, chat_id: str, message: str, queued=False) -> bool: return self._bot_actions.answer(chat_id, message, queued=queued) @staticmethod def convert_sticker_to_png_base64(sticker: BytesIO) -> str: """ Converts a sticker file (webp) to base64 string (png) :param sticker: the sticker to be converted :return: the base64 string """ file = BytesIO() img = Image.open(sticker) img.save(file, 'png') base64_content = base64.b64encode(file.getvalue()).decode() base64_string = 'data:image/png;base64,' + base64_content return base64_string # TODO: remove that @staticmethod def temp_save_to_txt(base64_string: str, suffix="") -> None: with open( f"/home/helitonmrf/Projects/WhatsGram_Stickers/test/sticker{suffix}.html", 'w') as fl: fl.write(f"<img src='{base64_string}' />")