diff options
Diffstat (limited to 'velascobot2/bot.py')
| -rw-r--r-- | velascobot2/bot.py | 606 |
1 files changed, 606 insertions, 0 deletions
diff --git a/velascobot2/bot.py b/velascobot2/bot.py new file mode 100644 index 0000000..5914906 --- /dev/null +++ b/velascobot2/bot.py @@ -0,0 +1,606 @@ +import random +import logging +import time +import sys + +from telegram import Chat, ChatMember, Message, Update +import telegram +from telegram.ext import ( + Application, + CommandHandler, + ContextTypes, + MessageHandler, + filters, +) +from telegram.error import NetworkError + +from velascobot2.utils import ChatData, GraphicType, Response +from velascobot2.messenger import Messenger + +logger = logging.getLogger(__name__) + + +start_msg = ( + "Hello there! Ask me for /help to see an overview of the available" + " commands." +) + + +wake_msg = "Good morning. I just woke up" + + +help_msg = """I answer to the following commands: + +/start - I say hi. +/about - What I'm about. +/explain - I explain how I work. +/help - I send this message. +/count - I tell you how many messages from this chat I have processed. +/speak - Forces me to speak (if you are allowed to). +/period - Change the period of my messages. (Maximum of 100000) +/answer - Change my chance to answer at will. (Decimal between 0 and 1) +/restrict - Toggle restriction of configuration commands to admins only. +/silence - Toggle restriction on mentions by the bot. +/ping - Pong! +/where - Tell my configuration for this chat. +""" + + +about_msg = ( + "I am yet another Markov Bot experiment. I read everything you type to me" + " and then spit back nonsensical messages that look like yours.\n\nYou can" + " send /explain if you want further explanation." +) + + +explanation = ( + "I decompose every message I read in groups of 3 consecutive words, so for" + " each consecutive pair I save the word that can follow them. I then use" + " this to make my own messages. At first I will only repeat your messages" + " because for each 2 words I will have very few possible following" + " words.\n\nI also separate my vocabulary by chats, so anything I learn in" + " one chat I will only say in that chat. For privacy, you know. Also, I" + " save my vocabulary in the form of a json dictionary, so no logs are" + " kept.\n\nMy default period in private chats is one message of mine from" + " each 2 messages received, and in group chats it's 10 messages I read for" + " each message I send." +) + + +def static_reply(text: str, format=None): + async def reply(update: Update, context: ContextTypes.DEFAULT_TYPE): + assert update.message + await update.message.reply_text(text, parse_mode=format) + + return reply + + +async def stop(update: Update, context: ContextTypes.DEFAULT_TYPE): + assert update.message + chat = update.message.chat + chat_id = str(chat.id) + # TODO: remove chat from DB + logger.warning( + "I got blocked by {} [{}]".format(chat.effective_name, chat_id) + ) + + +# Legacy function from original velascobot, created before "effective_name" +# was implemented in the telegram library. Kept to ensure consistency +def get_chat_title(chat: Chat): + """Returns group chat name or falls back to user name on a private chat.""" + if chat.title is not None: + return chat.title + elif chat.first_name is not None: + if chat.last_name is not None: + return chat.first_name + " " + chat.last_name + else: + return chat.first_name + else: + return "" + + +async def ping(update: Update, _context: ContextTypes.DEFAULT_TYPE) -> None: + """Pongs back.""" + + assert update.message is not None + await update.message.reply_text("Pong!") + + +async def error_handler(update: object, context: ContextTypes.DEFAULT_TYPE): + logger.error(context.error) + sys.exit(1) + + +class Bot: + def __init__( + self, + messenger: Messenger, + bot_token: str, + motd_chats: list[str] = [], + motd_list: list[str] = [], + whitelist: list[str] | None = None, + testing: bool = False, + ) -> None: + self.messenger = messenger + self.motd_chats = motd_chats + self.motd_list = motd_list + self.whitelist = whitelist + self.testing = testing + self.countdown: dict[str, int] = {} + + if testing: + assert ( + whitelist is not None + ), "Testing mode enabled but no whitelist was given." + + # Create the Application and pass it your bot's token. + app = ( + Application.builder() + .token(bot_token) + .concurrent_updates(True) + .build() + ) + + # on different commands - answer in Telegram + app.add_handler(CommandHandler("start", static_reply(start_msg))) + app.add_handler(CommandHandler("about", static_reply(about_msg))) + app.add_handler(CommandHandler("explain", static_reply(explanation))) + app.add_handler(CommandHandler("help", static_reply(help_msg))) + app.add_handler(CommandHandler("stop", stop)) + app.add_handler(CommandHandler("ping", ping)) + app.add_handler(CommandHandler("count", self.get_processed_messages)) + app.add_handler(CommandHandler("speak", self.speak_command)) + app.add_handler(CommandHandler("period", self.period)) + app.add_handler(CommandHandler("answer", self.answer)) + app.add_handler(CommandHandler("restrict", self.restrict)) + app.add_handler(CommandHandler("silence", self.silence)) + app.add_handler( + CommandHandler("where", static_reply("Not implemented")) + ) + app.add_handler( + MessageHandler( + ( + filters.TEXT + | filters.Sticker.ALL + | filters.ANIMATION + | filters.VIDEO + ), + self.read, + ) + ) + + if self.testing: + app.add_error_handler(error_handler) + + self.app = app + + async def get_processed_messages( + self, update: Update, _context: ContextTypes.DEFAULT_TYPE + ) -> None: + assert update.message + + chat_id = str(update.message.chat.id) + chat_data = self.messenger.get_chat_data(chat_id) + if chat_data is None: + await update.message.reply_text( + "I've processed no messages so far." + ) + return + + num = ( + str(chat_data.processed_messages) + if chat_data.processed_messages > 0 + else "no" + ) + await update.message.reply_text( + "I've processed {} messages so far.".format(num) + ) + + async def period( + self, update: Update, _context: ContextTypes.DEFAULT_TYPE + ) -> None: + """Responds with the current message period or sets a new one.""" + assert update.message + chat = update.message.chat + + chat_id = str(chat.id) + chat_data = self.messenger.get_chat_data(chat_id) + assert chat_data + + assert update.message.text + words = update.message.text.split() + if len(words) <= 1: + await update.message.reply_text( + f"The current speech period is {chat_data.message_period}" + ) + return + + if chat_data.restricted: + assert update.message.from_user + user = await chat.get_member(update.message.from_user.id) + if not self.superuser(user): + await update.message.reply_text( + "You do not have permissions to do that." + ) + return + try: + period = int(words[1]) + self.messenger.update_message_period(chat_id, period) + await update.message.reply_text( + f"Period of speaking set to {period}." + ) + except Exception: + await update.message.reply_text( + "Format was confusing; period unchanged from" + f" {chat_data.message_period}." + ) + + async def answer( + self, update: Update, _context: ContextTypes.DEFAULT_TYPE + ) -> None: + """Responds with the current answer chance or sets a new one.""" + assert update.message + chat = update.message.chat + + chat_id = str(chat.id) + chat_data = self.messenger.get_chat_data(chat_id) + assert chat_data + + assert update.message.text + words = update.message.text.split() + if len(words) <= 1: + await update.message.reply_text( + f"The current answer chance is {chat_data.answer_chance}" + ) + return + + if chat_data.restricted: + assert update.message.from_user + user = await chat.get_member(update.message.from_user.id) + if not self.superuser(user): + await update.message.reply_text( + "You do not have permissions to do that." + ) + return + try: + chance = float(words[1]) + self.messenger.update_answer_chance(chat_id, chance) + await update.message.reply_text( + f"Chance of answering set to {chance}." + ) + except Exception: + await update.message.reply_text( + "Format was confusing; answer chance unchanged from" + f" {chat_data.answer_chance}." + ) + + async def restrict( + self, update: Update, _context: ContextTypes.DEFAULT_TYPE + ) -> None: + """Toggles the current restriction state.""" + + assert update.message + if "group" not in update.message.chat.type: + await update.message.reply_text("That only works in groups.") + return + chat = update.message.chat + chat_id = str(chat.id) + assert update.message.from_user + user = await chat.get_member(update.message.from_user.id) + chat_data = self.messenger.get_chat_data(chat_id) + assert chat_data + if chat_data.restricted: + if not self.superuser(user): + await update.message.reply_text( + "You do not have permissions to do that." + ) + return + self.messenger.update_set_restricted(chat_id, not chat_data.restricted) + allowed = ( + "let only admins" if not chat_data.restricted else "let everyone" + ) + await update.message.reply_text(f"I will {allowed} configure me now.") + + async def silence( + self, update: Update, _context: ContextTypes.DEFAULT_TYPE + ) -> None: + """Toggles the current silence state.""" + + assert update.message + if "group" not in update.message.chat.type: + await update.message.reply_text("That only works in groups.") + return + chat = update.message.chat + chat_id = str(chat.id) + assert update.message.from_user + user = await chat.get_member(update.message.from_user.id) + chat_data = self.messenger.get_chat_data(chat_id) + assert chat_data + if chat_data.silenced: + if not self.superuser(user): + await update.message.reply_text( + "You do not have permissions to do that." + ) + return + self.messenger.update_set_silenced(chat_id, not chat_data.silenced) + allowed = "avoid mentioning" if not chat_data.silenced else "mention" + await update.message.reply_text(f"I will {allowed} people now.") + + async def send( + self, + chat_id: int | str, + message_type: Response, + content: str | tuple[GraphicType, str], + **kwargs, + ): + bot = self.app.bot + + if message_type == Response.TEXT: + assert isinstance(content, str) + if content == "": + logger.warning("Generated empty text response.") + return + max_length = telegram.constants.MessageLimit.MAX_TEXT_LENGTH + if len(content) > max_length: + content = content[:max_length] + msg = "reply" if (kwargs.get("reply_to_message_id")) else "message" + logger.info( + "Sending a {} to {}: '{}'".format(msg, chat_id, content) + ) + await bot.send_message(chat_id, content, **kwargs) + else: + assert isinstance(content, tuple) + graph_type, graph_id = content + if graph_id == "": + logger.warning(f"Generated empty {graph_type} response.") + return + logger.info( + 'Sending {} "{}" to {}'.format(graph_type, graph_id, chat_id) + ) + + if graph_type == GraphicType.STICKER: + sent = await bot.send_sticker( + chat_id=chat_id, sticker=graph_id, **kwargs + ) + elif graph_type == GraphicType.ANIMATION: + sent = await bot.send_animation( + chat_id=chat_id, animation=graph_id, **kwargs + ) + elif graph_type == GraphicType.VIDEO: + sent = await bot.send_video( + chat_id=chat_id, video=graph_id, **kwargs + ) + else: + raise Exception("Sending unkown graphic type") + logger.debug(f"Sent graphic: {sent}") + + async def announce(self, announcement, check=(lambda _: True)) -> None: + """Sends an announcement to all chats that pass a check.""" + for chat in self.messenger.get_all_chats(): + if check(chat): + try: + await self.send(chat.id, Response.TEXT, announcement) + logger.info( + "Sending announcement to chat {}".format(chat.id) + ) + except Exception: + pass + + async def wake_up(self, exclusive: bool = False) -> None: + if not self.messenger.wake_up: + return + if len(self.motd_list) == 0: + return + + def valid_chat(chat: ChatData) -> bool: + if self.whitelist is not None and chat.id not in self.whitelist: + return False + if self.motd_chats and exclusive: + return chat.id in self.motd_chats + if not self.motd_chats: + return "group" in chat.type + else: + return "group" in chat.type or chat.id in self.motd_chats + + wake_text = random.choice(self.motd_list) + await self.announce(wake_text, valid_chat) + + async def should_reply(self, message: Message, chat: ChatData) -> bool: + """Returns True if the bot should reply to a message.""" + if self.messenger.is_mute(): + # Not if mute time hasn't finished + return False + if not self.testing and chat.restricted: + # If we're not in testing mode and the chat is restricted + assert message.from_user + user = await message.chat.get_member(message.from_user.id) + if not self.superuser(user): + # ...And the user has no permissions, should not reply + return False + + # otherwise (testing mode, or the chat is unrestricted, + # or the user is a superuser) + replied_message = message.reply_to_message + text = message.text.casefold() if message.text else "" + replies_to_me = ( + replied_message is not None + and replied_message.from_user + and replied_message.from_user.name == self.messenger.username + ) + # Only if it's a reply to a message of ours or the bot is mentioned in + # the message + return replies_to_me or self.messenger.mentioned(text) + + def superuser(self, member: ChatMember) -> bool: + """Returns True if the user is a group chat admin, creator or in the + bosses list.""" + logger.info( + "user {} ({}) requesting a restricted action".format( + str(member.user.id), member.user.name + ) + ) + return ( + (member.status == "creator") + or (member.status == "administrator") + or (member.user.id in self.messenger.bosses) + ) + + async def speak(self, chat_id: str, **kwargs) -> None: + """Send a new message.""" + + if self.whitelist is not None and chat_id not in self.whitelist: + # Don't, if there's a whitelist and this chat is not in it + return + if self.messenger.is_mute(): + # Don't, if mute time isn't over + return + try: + response, content = self.messenger.give_response(chat_id) + await self.send(chat_id, response, content, **kwargs) + while random.random() <= self.messenger.repeat: + del kwargs["reply_to_message_id"] + response, content = self.messenger.give_response(chat_id) + await self.send(chat_id, response, content, **kwargs) + # Consider any Network Error as a Telegram temporary ban, as I + # couldn't find out in the documentation how error 429 is handled by + # python-telegram-bot + except NetworkError as e: + logger.error("Sending a message caused network error:") + logger.exception(e) + logger.error( + "Going mute for {} seconds.".format(self.messenger.mute_time) + ) + self.messenger.set_mute_timer(int(time.perf_counter())) + except Exception as e: + logger.error("Sending a message caused exception:") + logger.exception(e) + + def extract_graphic(self, message: Message) -> tuple[GraphicType, str]: + logger.debug(f"Sticker: {message.sticker}") + logger.debug(f"Animation: {message.animation}") + logger.debug(f"Sticker: {message.video}") + if message.sticker is not None: + return GraphicType.STICKER, message.sticker.file_id + elif message.animation is not None: + return GraphicType.ANIMATION, message.animation.file_id + if message.video is not None: + return GraphicType.VIDEO, message.video.file_id + else: + raise Exception("Unknown graphic type") + + # Rolls the chance for answering in this specific chat, + # according to the answer probability + def is_answering(self, answer_chance: float) -> bool: + """Rolls the chance for answering to a message.""" + if answer_chance == 1: + return True + elif answer_chance == 0: + return False + rand = random.random() + return rand <= answer_chance + + async def read(self, update: Update, context: ContextTypes.DEFAULT_TYPE): + """Reads a non-command message.""" + # Ignore non-message updates + if update.message is None: + return + + message = update.message + chat = message.chat + chat_id = str(chat.id) + chat_name = get_chat_title(chat) + + self.messenger.ensure_chat_exists(chat_id, chat.type, chat_name) + chat_data = self.messenger.get_chat_data(chat_id) + assert chat_data + + if message.text is not None: + self.messenger.process_text(chat_id, message.text) + else: + graph_type, graph_id = self.extract_graphic(message) + self.messenger.process_graphic(chat_id, graph_type, graph_id) + + # Update the ChatData name if it has changed + if chat_name != chat_data.name: + self.messenger.update_chat_name(chat_id, chat_name) + + # Check if it's a "replyable" message & roll the chance to do so + should_reply = await self.should_reply(update.message, chat_data) + if should_reply and self.is_answering(chat_data.answer_chance): + await self.speak(chat_id, reply_to_message_id=message.message_id) + return + + # Decrease the countdown for the chat, and send a message if it + # reached 0 + if chat_id not in self.countdown: + self.countdown[chat_id] = chat_data.message_period + self.countdown[chat_id] -= 1 + if self.countdown[chat_id] < 0: + self.countdown[chat_id] = chat_data.message_period + await self.speak(chat_id) + + async def speak_command( + self, update: Update, context: ContextTypes.DEFAULT_TYPE + ): + """Handles the /speak command.""" + chat = update.message.chat + chat_id = str(chat.id) + chat_data = self.messenger.get_chat_data(chat_id) + + if chat_data is None: + return + + if not self.testing and chat_data.restricted: + user = update.message.chat.get_member(update.message.from_user.id) + if not self.superuser(user): + # update.message.reply_text("You do not have permissions to do that.") + return + + mid = str(update.message.message_id) + replied = update.message.reply_to_message + # Reply to the message that the command replies to, otherwise to the + # command itself + rid = replied.message_id if replied else mid + # words = update.message.text.split() + # if len(words) > 1: + # self.messenger.process_text(chat_id, ' '.join(words[1:])) + if chat_data.processed_messages > 0: + await self.speak(chat_id, reply_to_message_id=rid) + else: + empty_gen_warning = "I haven't learned a single word yet." + await self.send( + chat_id, + Response.TEXT, + content=empty_gen_warning, + reply_to_message_id=rid, + ) + + async def show_config( + self, update: Update, context: ContextTypes.DEFAULT_TYPE + ): + """Sends a summary of the chat configuration.""" + msg = update.message + assert msg + chat = msg.chat + chat_id = str(chat.id) + chat_data = self.messenger.get_chat_data(chat_id) + assert chat_data + if chat_data.restricted and chat_data.silenced: + permissions = "restricted and silenced" + elif chat_data.restricted: + permissions = "restricted but not silenced" + elif chat_data.silenced: + permissions = "not restricted but silenced" + else: + permissions = "neither restricted nor silenced" + + answer = ( + f"You're messaging in the chat of saved title" + f"__{chat_data.name}__, with id `{chat_data.id}`, message count " + f"{chat_data.processed_messages}, period " + f"{chat_data.message_period}, and answer probability " + f"{chat_data.answer_chance}.\n\nThis chat is {permissions}." + ) + + await msg.reply_markdown(answer) |