diff --git a/main.py b/main.py index 1ed11cb..491a556 100644 --- a/main.py +++ b/main.py @@ -6,6 +6,8 @@ import configparser import re import io import mimetypes +import asyncio +import logging from collections import defaultdict from PythonSed import Sed @@ -13,7 +15,6 @@ from slixmpp import ClientXMPP from urllib.parse import urlparse, parse_qs, urlunparse from pantomime import normalize_mimetype -sed = Sed() sed_parse = re.compile("(? data_limit): - return + lenght = 0 outfile = io.BytesIO() for chunk in r.iter_content( chunk_size=512, decode_unicode=False, ): + lenght += 512 + if lenght >= data_limit: + return outfile.write(chunk) await self.embed_file(url, sender, mtype, ftype, outfile) - except Exception as e: - print(e) + except Exception: + ... async def embed_file(self, url, sender, mtype, ftype, outfile): ext = mimetypes.guess_extension(ftype) @@ -174,6 +176,8 @@ class AngelBot(ClientXMPP): message.send() async def parse_urls(self, msg, urls, sender, mtype): + if "nsfw" in msg["body"].lower(): + return for u in urls: if u in self.messages[sender]["links"]: continue @@ -183,29 +187,10 @@ class AngelBot(ClientXMPP): uri = urlparse(u) await self.parse_uri(uri, sender, mtype) - def muc_word(self, func): - name = func.__name__ - self.muc_commands[name] = func - return func - - def muc_command(self, func): - name = self.prefix + func.__name__ - self.muc_commands[name] = func - return func - - def word(self, func): - name = func.__name__ - self.commands[name] = func - return func - - def command(self, func): - name = self.prefix + func.__name__ - self.commands[name] = func - return func - - def s(self, msg, sender, mtype): + def sed_command(self, msg, sender, mtype): try: text = msg["body"] + print(f"{text = }") if not sed_cmd.match(text): self.messages[sender]["messages"].add(text) return @@ -214,16 +199,19 @@ class AngelBot(ClientXMPP): if len(sed_args) < 4: return + sed = Sed() + sed.load_string(text) for message in self.messages[sender]["messages"]: - if sed_args[1] in message: - sed.load_string(text) - fobj = io.StringIO(message) - res = sed.apply(fobj, None) - res = "\n".join(res) - self.messages[sender]["messages"].add(res) + message_io = io.StringIO(message) + res = sed.apply(message_io, None) + print(f"{ res = }") + out = "\n".join(res) + print(f"{ out = }") + if out.strip() != message.strip(): + self.messages[sender]["messages"].add(out) return self.send_message( mto=sender, - mbody=res, + mbody=out, mtype=mtype, ) @@ -231,26 +219,11 @@ class AngelBot(ClientXMPP): print(e) return - def __init__( - self, - jid, - password, - owner, - nick="angel", - prefix="!", - autojoin=None, - ): + def __init__(self, jid, password, nick="angel", autojoin=None): ClientXMPP.__init__(self, jid, password) - self.owner = owner self.jid = jid - self.prefix = prefix self.nick = nick self.autojoin = autojoin or [] - 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("disconnected", lambda _: self.connect()) - self.register_plugin("xep_0030") self.register_plugin("xep_0060") self.register_plugin("xep_0054") @@ -260,14 +233,25 @@ class AngelBot(ClientXMPP): self.register_plugin("xep_0153") self.register_plugin("xep_0363") + 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): self.send_presence() await self.get_roster() for channel in self.autojoin: - self.plugin["xep_0045"].join_muc(channel, self.nick) - await self.avatar() + try: + self.plugin["xep_0045"].join_muc(channel, self.nick) + except: + ... + # await self.update_info() + logging.info("Session started!") - async def avatar(self): + async def update_info(self): with open("angel.png", "rb") as avatar_file: avatar = avatar_file.read() @@ -275,38 +259,40 @@ class AngelBot(ClientXMPP): avatar_id = self.plugin["xep_0084"].generate_id(avatar) avatar_bytes = len(avatar) - info = {"id": avatar_id, "type": avatar_type, "bytes": avatar_bytes} + info = { + "id": avatar_id, + "type": avatar_type, + "bytes": avatar_bytes, + } - await self.plugin["xep_0153"].set_avatar( - jid=self.jid, - avatar=avatar, - mtype=avatar_type, + asyncio.gather(self.plugin["xep_0084"].publish_avatar(avatar)) + + asyncio.gather( + self.plugin["xep_0153"].set_avatar( + avatar=avatar, + mtype=avatar_type, + ) ) - await self.plugin["xep_0084"].publish_avatar(avatar) - await self.plugin["xep_0084"].publish_avatar_metadata([info]) + + asyncio.gather(self.plugin["xep_0084"].publish_avatar_metadata([info])) async def message(self, msg): if msg["type"] in ("chat", "normal"): mtype = "chat" sender = msg["from"].bare - message = msg["body"] - ctx = message.strip().split() + edit = "urn:xmpp:message-correct:0" in str(msg) + if edit: + return + try: if not msg["oob"]["url"]: if urls := self.get_urls(msg): await self.parse_urls(msg, urls, sender, mtype) except Exception: - pass + ... - edit = "urn:xmpp:message-correct:0" in str(msg) - - cm = ctx.pop(0) - if cm in self.commands: - if not edit: - self.commands[cm](msg, ctx, sender) - else: - self.s(msg, sender, mtype) + self.sed_command(msg, sender, mtype) async def muc_message(self, msg): if msg["type"] in ("groupchat", "normal"): @@ -315,56 +301,29 @@ class AngelBot(ClientXMPP): if msg["mucnick"] == self.nick: return - ctx = msg["body"].strip().split() + edit = "urn:xmpp:message-correct:0" in str(msg) + if edit: + return + try: if not msg["oob"]["url"]: if urls := self.get_urls(msg): await self.parse_urls(msg, urls, sender, mtype) except Exception: pass - edit = "urn:xmpp:message-correct:0" in str(msg) - cm = ctx.pop(0) - if cm in self.muc_commands: - if not edit: - self.muc_commands[cm](msg, ctx, sender) - else: - self.s(msg, sender, mtype) - @self.muc_word - def ping(msg, ctx, sender): - if ctx: - return - bot.send_message( - mto=sender, mbody=f"pong: {msg['mucnick']}", mtype="groupchat" - ) - - @self.word - def ping(msg, ctx, sender): - if ctx: - return - bot.send_message( - mto=sender, - mbody=f"pong: {msg.get_from().bare}", - mtype="chat", - ) - - @self.command - def join(msg, ctx, sender): - if sender == bot.owner: - for channel in ctx: - bot.plugin["xep_0045"].join_muc(channel, bot.nick) + self.sed_command(msg, sender, mtype) if __name__ == "__main__": config = configparser.ConfigParser() config.read("config.ini") - + logging.basicConfig(level=logging.DEBUG) jid = config["angel"]["jid"] password = config["angel"]["password"] - owner = config["angel"]["owner"] autojoin = config["angel"]["autojoin"].split() - bot = AngelBot(jid, password, owner, autojoin=autojoin) + bot = AngelBot(jid, password, autojoin=autojoin) bot.connect() - bot.process() + bot.process(forever=True)