diff --git a/src/ollama_slixmpp_omemo_bot/main.py b/src/ollama_slixmpp_omemo_bot/main.py index 6b4e193..f942fd3 100644 --- a/src/ollama_slixmpp_omemo_bot/main.py +++ b/src/ollama_slixmpp_omemo_bot/main.py @@ -5,8 +5,6 @@ import logging from getpass import getpass from argparse import ArgumentParser -import slixmpp_omemo - from ollama_bot import OllamaBot log = logging.getLogger(__name__) @@ -33,34 +31,16 @@ if __name__ == "__main__": ) parser.add_argument("-j", "--jid", dest="jid", help="JID to use") parser.add_argument("-p", "--password", dest="password", help="password to use") - DATA_DIR = os.path.join( - os.path.dirname(os.path.abspath(__file__)), - "omemo", - ) - parser.add_argument( - "--data-dir", dest="data_dir", help="data directory", default=DATA_DIR - ) args = parser.parse_args() logging.basicConfig(level=args.loglevel, format="%(levelname)-8s %(message)s") if args.jid is None: args.jid = input("JID: ") if args.password is None: args.password = getpass("Password: ") - os.makedirs(args.data_dir, exist_ok=True) xmpp = OllamaBot(args.jid, args.password) xmpp.register_plugin("xep_0030") # Service Discovery xmpp.register_plugin("xep_0199") # XMPP Ping xmpp.register_plugin("xep_0380") # Explicit Message Encryption - try: - xmpp.register_plugin( - "xep_0384", - { - "data_dir": args.data_dir, - }, - module=slixmpp_omemo, - ) - except slixmpp_omemo.PluginCouldNotLoad: - log.exception("And error occured when loading the omemo plugin.") - sys.exit(1) + xmpp.connect() xmpp.process() diff --git a/src/ollama_slixmpp_omemo_bot/ollama_bot.py b/src/ollama_slixmpp_omemo_bot/ollama_bot.py index 4ae0815..18d4472 100644 --- a/src/ollama_slixmpp_omemo_bot/ollama_bot.py +++ b/src/ollama_slixmpp_omemo_bot/ollama_bot.py @@ -10,14 +10,6 @@ from slixmpp.types import JidStr, MessageTypes from slixmpp.xmlstream.handler import CoroutineCallback from slixmpp.xmlstream.handler.coroutine_callback import CoroutineFunction from slixmpp.xmlstream.matcher import MatchXPath -from slixmpp_omemo import ( - EncryptionPrepareException, - MissingOwnKey, - NoAvailableSession, - UndecidedException, - UntrustedException, -) -from omemo.exceptions import MissingBundleException class LEVELS(Enum): @@ -123,64 +115,20 @@ class OllamaBot(ClientXMPP): if not msg["body"]: return None log = open("log.txt", "a", 1) - if not self["xep_0384"].is_encrypted(msg): + body: Optional[bytes] = await self["xep_0384"].decrypt_message( + encrypted, mfrom, allow_untrusted + ) + if msg['body'] is not None: log.write(f"{mfrom}: {msg['body']}\n") - ollama_server_response: Optional[str] = self.message_to_ollama_server( - msg["body"] - ) - await self.plain_reply(mto, mtype, f"{ollama_server_response or ''}") - - return None - try: - encrypted = msg["omemo_encrypted"] - body: Optional[bytes] = await self["xep_0384"].decrypt_message( - encrypted, mfrom, allow_untrusted - ) - if body is not None: - decoded: str = body.decode("utf8") - log.write(f"{mfrom}: {decoded}\n") - if self.is_command(decoded): - await self.handle_command(mto, mtype, decoded) - elif self.debug_level == LEVELS.DEBUG: - ollama_server_response: Optional[str] = ( - self.message_to_ollama_server(decoded) - ) - await self.plain_reply( - mto, mtype, f"{ollama_server_response or ''}" - ) - except MissingOwnKey: - await self.plain_reply( - mto, - mtype, - "Error: Message not encrypted for me.", - ) - except NoAvailableSession: - await self.plain_reply( - mto, - mtype, - "Error: Message uses an encrypted session I don't know about.", - ) - except (UndecidedException, UntrustedException) as exn: - await self.plain_reply( - mto, - mtype, - ( - f"WARNING: Your device '{exn.device}' is not in my trusted devices." - f"Allowing untrusted..." - ), - ) - await self.message_handler(msg, allow_untrusted=True) - except EncryptionPrepareException: - await self.plain_reply( - mto, mtype, "Error: I was not able to decrypt the message." - ) - except Exception as exn: - await self.plain_reply( - mto, - mtype, - "Error: Exception occured while attempting decryption.\n%r" % exn, - ) - raise + if self.is_command(decoded): + await self.handle_command(mto, mtype, decoded) + elif self.debug_level == LEVELS.DEBUG: + ollama_server_response: Optional[str] = ( + self.message_to_ollama_server(decoded) + ) + await self.plain_reply( + mto, mtype, f"{ollama_server_response or ''}" + ) return None async def plain_reply(self, mto: JID, mtype: Optional[MessageTypes], body): @@ -190,49 +138,6 @@ class OllamaBot(ClientXMPP): msg["body"] = body return msg.send() - async def encrypted_reply(self, mto: JID, mtype: Optional[MessageTypes], body): - msg = self.make_message(mto=mto, mtype=mtype) - msg["eme"]["namespace"] = self.eme_ns - msg["eme"]["name"] = self["xep_0380"].mechanisms[self.eme_ns] - expect_problems: Optional[dict[JID, list[int]]] = {} - while True: - try: - recipients = [mto] - encrypt = await self["xep_0384"].encrypt_message( - body, recipients, expect_problems - ) - msg.append(encrypt) - return msg.send() - except UndecidedException as exn: - await self["xep_0384"].trust(exn.bare_jid, exn.device, exn.ik) - except EncryptionPrepareException as exn: - for error in exn.errors: - if isinstance(error, MissingBundleException): - await self.plain_reply( - mto, - mtype, - f'Could not find keys for device "{error.device}"' - f' of recipient "{error.bare_jid}". Skipping.', - ) - jid: JID = JID(error.bare_jid) - device_list = expect_problems.setdefault(jid, []) - device_list.append(error.device) - except (IqError, IqTimeout) as exn: - await self.plain_reply( - mto, - mtype, - "An error occured while fetching information on a recipient.\n%r" - % exn, - ) - return None - except Exception as exn: - await self.plain_reply( - mto, - mtype, - "An error occured while attempting to encrypt.\n%r" % exn, - ) - raise - def message_to_ollama_server(self, msg: Optional[str]) -> Optional[str]: if msg is not None: client = Client(