This commit is contained in:
hueso 2025-03-06 20:55:17 -03:00
parent 193aef1d0b
commit a43f5a9ae2
2 changed files with 14 additions and 129 deletions

View File

@ -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()

View File

@ -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(