import asyncio from collections import defaultdict from slixmpp import ClientXMPP from slixmpp.stanza import Message class Lifo(list): """Limited size LIFO array to store messages and urls.""" def __init__(self, size): """Initialize the LIFO array.""" super().__init__() self.size = size def add(self, item): """Add an item to the LIFO array.""" self.insert(0, item) if len(self) > self.size: self.pop() def create_messages_dict(): return defaultdict( lambda: { 'messages': Lifo(100), 'links': Lifo(10), 'previews': Lifo(10), } ) class RegexCmd: """Regex command decorator.""" def __init__(self, bot, pattern, block=False, matcher=None): """Initialize the decorator.""" self.pattern = pattern self.bot = bot self.block = block self.matcher = matcher def __call__(self, func): """Call the decorator.""" self.bot.regex_cmds.append(self) self.func = func return self class AngelBot(ClientXMPP): """AngelBot class.""" def __init__(self, jid, password, nick='angel', autojoin=None): """Initialize the bot.""" super().__init__(jid, password) self.jid = jid self.nick = nick self.autojoin = autojoin or [] self.messages = create_messages_dict() self.register_plugins() self.add_handlers() def reply(self, msg, body): """Reply to a message.""" self.save_message_history(msg) self.raw_reply(msg, body) def raw_reply(self, msg, body): """Reply to a message without saving history.""" self.send_message( mto=msg['from'].bare, mbody=body, mtype=msg['type'], ) def save_message_history(self, msg): """Save the history of messages.""" sender = msg['from'].bare self.messages[sender]['messages'].add(msg['body']) def get_message_history(self, msg): """Get the messages from the sender.""" sender = msg['from'].bare return self.messages[sender]['messages'] def save_link_history(self, msg, url): """Save the history of links.""" sender = msg['from'].bare self.messages[sender]['links'].add(url) def get_link_history(self, msg): """Get the links from the sender.""" sender = msg['from'].bare return self.messages[sender]['links'] def save_preview_history(self, msg, preview): """Save the history of previews.""" sender = msg['from'].bare self.messages[sender]['previews'].add(preview) def get_preview_history(self, msg): """Get the previews from the sender.""" sender = msg['from'].bare return self.messages[sender]['previews'] regex_cmds = [] async def embed_file(self, sender, mtype, ftype, fname, outfile): """Embed a file and send the result to the sender.""" furl = await self.plugin['xep_0363'].upload_file( fname, content_type=ftype, input_file=outfile ) self.messages[sender]['links'].add(furl) message = self.make_message(sender) message['body'] = furl message['type'] = mtype message['oob']['url'] = furl message.send() def register_plugins(self): self.register_plugin('xep_0030') self.register_plugin('xep_0060') self.register_plugin('xep_0054') self.register_plugin('xep_0045') self.register_plugin('xep_0066') self.register_plugin('xep_0084') self.register_plugin('xep_0153') self.register_plugin('xep_0363') def add_handlers(self): self.add_event_handler('session_start', self.session_start) self.add_event_handler('message', self.message) self.add_event_handler('groupchat_message', self.muc_message) # self.add_event_handler("vcard_avatar_update", self.debug_event) # self.add_event_handler("stream_error", self.debug_event) self.add_event_handler('disconnected', lambda _: self.connect()) async def session_start(self, event): """Start the bot.""" self.send_presence() await self.get_roster() await self.update_info() for channel in self.autojoin: try: self.plugin['xep_0045'].join_muc(channel, self.nick) except Exception as e: print(e) async def update_info(self): """Update the bot info.""" with open('angel.png', 'rb') as avatar_file: avatar = avatar_file.read() avatar_type = 'image/png' avatar_id = self.plugin['xep_0084'].generate_id(avatar) avatar_bytes = len(avatar) asyncio.gather(self.plugin['xep_0084'].publish_avatar(avatar)) asyncio.gather( self.plugin['xep_0153'].set_avatar( avatar=avatar, mtype=avatar_type, ) ) info = { 'id': avatar_id, 'type': avatar_type, 'bytes': avatar_bytes, } asyncio.gather(self.plugin['xep_0084'].publish_avatar_metadata([info])) vcard = self.plugin['xep_0054'].make_vcard() vcard['URL'] = 'https://wiki.kalli.st/Angel' vcard['DESC'] = 'Angel is a bot that can do link previews and embeds.' vcard['NICKNAME'] = 'Angel' vcard['FN'] = 'Angel' asyncio.gather(self.plugin['xep_0054'].publish_vcard(vcard)) async def message(self, msg): """Process a message.""" if msg['type'] in ('chat', 'normal'): edit = 'urn:xmpp:message-correct:0' in str(msg) if edit: return mtype = msg['type'] sender = msg['from'].bare self.process_commands(msg, sender, mtype) async def muc_message(self, msg): """Process a groupchat message.""" if msg['type'] in ('groupchat', 'normal'): edit = 'urn:xmpp:message-correct:0' in str(msg) if edit: return if msg['mucnick'] == self.nick: return mtype = msg['type'] sender = msg['from'].bare self.process_commands(msg, sender, mtype) def process_commands(self, msg, sender, mtype): """Process commands.""" for cmd in self.regex_cmds: if cmd.pattern.search(msg['body']): ctx = CommandContext(self, msg) if cmd.matcher and not cmd.matcher(ctx): continue cmd.func(ctx) if cmd.block: return self.messages[sender]['messages'].add(msg['body']) class CommandContext: """Command context.""" def __init__(self, bot: AngelBot, msg: Message): """Initialize the command context.""" self.bot = bot self.msg = msg def reply(self, body): """Get the reply function.""" return self.bot.reply(self.msg, body) def raw_reply(self, body): """Get the raw reply function.""" return self.bot.raw_reply(self.msg, body) @property def sender(self): """Get the sender of the message.""" return self.msg['from'].bare @property def mtype(self): """Get the message type.""" return self.msg['type'] @property def body(self): """Get the message body.""" return self.msg['body'] @property def message_history(self): """Get the message history.""" return self.bot.get_message_history(self.msg) @property def link_history(self): """Get the link history.""" return self.bot.get_link_history(self.msg) @property def preview_history(self): """Get the preview history.""" return self.bot.get_preview_history(self.msg) def save_link_history(self, url): """Save the link history.""" self.bot.save_link_history(self.msg, url) def save_message_history(self): """Save the message history.""" self.bot.save_message_history(self.msg) def save_preview_history(self, preview): """Save the preview history.""" self.bot.save_preview_history(self.msg, preview) @property def is_oob(self): """Check if the message is OOB.""" return bool(self.msg['oob']['url']) def embed_file(self, ftype, fname, outfile): """Embed a file and send the result to the sender.""" asyncio.gather( self.bot.embed_file(self.sender, self.mtype, ftype, fname, outfile) )