Initial bot.
OMEMO group chat works, starting to work on actual commands and DB.
This commit is contained in:
commit
c351f7f695
|
|
@ -0,0 +1,346 @@
|
|||
#!/usr/bin/env -S PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION=python python3
|
||||
# -*- coding: utf-8 -*-
|
||||
|
||||
from argparse import ArgumentParser
|
||||
import asyncio
|
||||
from datetime import datetime, timedelta
|
||||
from getpass import getpass
|
||||
import logging
|
||||
import os
|
||||
import re
|
||||
import sys
|
||||
|
||||
import aiosqlite
|
||||
|
||||
from slixmpp import ClientXMPP, JID
|
||||
from slixmpp.exceptions import IqTimeout, IqError
|
||||
from slixmpp.stanza import Message
|
||||
import slixmpp_omemo
|
||||
from slixmpp_omemo import PluginCouldNotLoad, MissingOwnKey, EncryptionPrepareException
|
||||
from slixmpp_omemo import UndecidedException, UntrustedException, NoAvailableSession
|
||||
from omemo.exceptions import MissingBundleException
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
|
||||
LEVEL_DEBUG = 0
|
||||
LEVEL_ERROR = 1
|
||||
|
||||
EME_NS = 'eu.siacs.conversations.axolotl'
|
||||
|
||||
SQL_INIT = """
|
||||
CREATE TABLE rooms_to_join(
|
||||
room TEXT UNIQUE);
|
||||
|
||||
CREATE TABLE budgets(
|
||||
id INTEGER PRIMARY KEY,
|
||||
room TEXT UNIQUE,
|
||||
descr TEXT
|
||||
);
|
||||
|
||||
CREATE TABLE members(
|
||||
id INTEGER PRIMARY KEY,
|
||||
budget INTEGER,
|
||||
jid TEXT,
|
||||
FOREIGN KEY(budget)
|
||||
REFERENCES budgets(id)
|
||||
ON UPDATE CASCADE
|
||||
ON DELETE CASCADE
|
||||
);
|
||||
|
||||
CREATE TABLE bills(
|
||||
id INTEGER PRIMARY KEY,
|
||||
payer INTEGER,
|
||||
amount INTEGER,
|
||||
descr TEXT,
|
||||
date INTEGER,
|
||||
c_date INTEGER,
|
||||
FOREIGN KEY(payer)
|
||||
REFERENCES members(id)
|
||||
ON UPDATE CASCADE
|
||||
ON DELETE CASCADE
|
||||
);
|
||||
|
||||
CREATE TABLE bill_owers(
|
||||
bill INTEGER,
|
||||
ower INTEGER,
|
||||
weight INTEGER,
|
||||
FOREIGN KEY(bill)
|
||||
REFERENCES bills(id)
|
||||
ON UPDATE CASCADE
|
||||
ON DELETE CASCADE,
|
||||
FOREIGN KEY(ower)
|
||||
REFERENCES members(id)
|
||||
ON UPDATE CASCADE
|
||||
ON DELETE CASCADE
|
||||
);"""
|
||||
|
||||
class Albus(ClientXMPP):
|
||||
"""Albus, le bot-chat."""
|
||||
|
||||
nick = 'albus'
|
||||
debug_level = LEVEL_DEBUG
|
||||
|
||||
def __init__(self, jid, password, db_path):
|
||||
ClientXMPP.__init__(self, jid, password)
|
||||
|
||||
self.db_path = db_path
|
||||
self.db = None
|
||||
|
||||
self.auto_authorize = True
|
||||
self.auto_subscribe = True
|
||||
|
||||
self.room_info = {}
|
||||
|
||||
self.add_event_handler('session_start', self.on_start)
|
||||
self.add_event_handler('message', self.on_message)
|
||||
self.add_event_handler('groupchat_direct_invite', self.on_direct_invite)
|
||||
self.add_event_handler('groupchat_presence', self.on_muc_presence)
|
||||
self.add_event_handler('groupchat_message', self.on_muc_message)
|
||||
|
||||
self.cmd_match = re.compile(r'^!')
|
||||
self.cmd_extr = re.compile(r'^!(?P<cmd>(\w|[._-])+)\s*(?P<args>.*)$')
|
||||
|
||||
self.cmd_handlers = {}
|
||||
self.cmd_handlers['budget-init'] = self.cmd_budget_init
|
||||
|
||||
|
||||
async def on_exit(self):
|
||||
if self.db is not None:
|
||||
await self.db.close()
|
||||
log.info('Exiting')
|
||||
|
||||
async def on_start(self, _):
|
||||
self.send_presence()
|
||||
self.get_roster()
|
||||
|
||||
self.xep_0045 = self['xep_0045']
|
||||
self.xep_0384 = self['xep_0384']
|
||||
self.xep_0380 = self['xep_0380']
|
||||
|
||||
log.info('Loading db: %s' % self.db_path)
|
||||
self.db = await aiosqlite.connect(self.db_path)
|
||||
async with self.db.execute("SELECT name FROM sqlite_master WHERE type='table' and name='rooms_to_join'") as c:
|
||||
r = await c.fetchone()
|
||||
if r is None:
|
||||
await self.db_init()
|
||||
else:
|
||||
async with self.db.execute('SELECT room FROM rooms_to_join') as c:
|
||||
async for (room,) in c:
|
||||
await self.process_join(JID(room))
|
||||
|
||||
async def db_init(self):
|
||||
log.info('Initializing db')
|
||||
await self.db.executescript(SQL_INIT)
|
||||
await self.db.commit()
|
||||
|
||||
async def on_direct_invite(self, msg):
|
||||
mroom = msg['groupchat_invite']['jid']
|
||||
await self.db.execute('INSERT INTO rooms_to_join (room) VALUES (?)', (mroom,))
|
||||
await self.db.commit()
|
||||
await self.process_join(mroom)
|
||||
|
||||
async def process_join(self, mroom):
|
||||
log.info('Joining room: %s' % mroom)
|
||||
await self.xep_0045.join_muc_wait(mroom, self.nick, seconds=0)
|
||||
members = self.room_info.setdefault(mroom, set())
|
||||
for a in ('owner', 'admin', 'member', 'none'):
|
||||
for j in await self.xep_0045.get_affiliation_list(mroom, a):
|
||||
members.add(JID(j))
|
||||
|
||||
async def on_muc_presence(self, msg):
|
||||
mtype = msg['type']
|
||||
mroom = msg['muc']['room']
|
||||
maffil = msg['muc']['affiliation']
|
||||
mjid = JID(msg['muc']['jid'].bare)
|
||||
|
||||
members = self.room_info.setdefault(mroom, set())
|
||||
if mtype == 'available':
|
||||
members.add(mjid)
|
||||
elif mtype == 'unavailable' and maffil in ('none', 'outcast'):
|
||||
members.discard(mjid)
|
||||
|
||||
async def on_muc_message(self, msg):
|
||||
mnick = msg['muc']['nick']
|
||||
mroom = msg['muc']['room']
|
||||
mfrom = JID(self.xep_0045.get_jid_property(mroom, mnick, 'jid'))
|
||||
|
||||
if mfrom == self.boundjid:
|
||||
return
|
||||
|
||||
body = await self.decrypt_message(msg, mfrom, eto=mroom, etype='groupchat')
|
||||
if body is not None:
|
||||
await self.process_muc_message(mroom, mfrom, body)
|
||||
|
||||
async def on_message(self, msg):
|
||||
mtype = msg['type']
|
||||
mfrom = msg['from']
|
||||
|
||||
if mtype != 'chat' or mfrom == self.boundjid:
|
||||
return
|
||||
|
||||
body = await self.decrypt_message(msg, mfrom, mfrom, type)
|
||||
if body is not None:
|
||||
await self.process_message(mfrom, body)
|
||||
|
||||
async def process_muc_message(self, mroom, mfrom, body):
|
||||
if not self.cmd_match.match(body):
|
||||
return
|
||||
|
||||
m = self.cmd_extr.match(body)
|
||||
if m is None:
|
||||
await self.room_send(mroom, 'erreur: la commande est mal formée')
|
||||
return
|
||||
|
||||
cmd = m.group('cmd')
|
||||
args = m.group('args')
|
||||
|
||||
h = self.cmd_handlers.get(cmd)
|
||||
if h is None:
|
||||
await self.room_send(mroom, 'erreur: je ne connais pas la commande "%s"' % cmd)
|
||||
return
|
||||
|
||||
await h(mroom, mfrom, args)
|
||||
|
||||
async def cmd_budget_init(self, mroom, mfrom, descr):
|
||||
async with self.db.cursor() as c:
|
||||
await c.execute('INSERT INTO budgets(room, descr) VALUES (?, ?)', (mroom, descr))
|
||||
b_id = c.lastrowid
|
||||
|
||||
members = set()
|
||||
for m in self.room_info[mroom]:
|
||||
if m.bare == self.boundjid.bare or m.bare in members:
|
||||
continue
|
||||
members.add(m.bare)
|
||||
await c.execute('INSERT INTO members(budget, jid) VALUES (?, ?)', (b_id, m.bare))
|
||||
|
||||
await self.db.commit()
|
||||
|
||||
line1 = 'Nouveau budget (%03i): %s\n' % (b_id, descr)
|
||||
line2 = 'Membres: ' + ', '.join(m.split('@')[0] for m in members)
|
||||
await self.room_send(mroom, line1 + line2)
|
||||
|
||||
async def process_message(self, mfrom, body):
|
||||
pass
|
||||
|
||||
async def room_send(self, mroom, mbody):
|
||||
await self.encrypted_send(JID(mroom), 'groupchat', mbody, self.room_info[mroom])
|
||||
|
||||
async def decrypt_message(self, msg, mfrom, eto, etype):
|
||||
if not self.xep_0384.is_encrypted(msg):
|
||||
await self.plain_send(eto, etype, 'attention: ce message n\'est pas chiffré')
|
||||
return msg['body']
|
||||
|
||||
try:
|
||||
encrypted = msg['omemo_encrypted']
|
||||
body = await self.xep_0384.decrypt_message(encrypted, mfrom)
|
||||
if body is not None:
|
||||
body = body.decode('utf8')
|
||||
return body
|
||||
except (MissingOwnKey,):
|
||||
await self.plain_send(eto, etype, 'erreur: le déchiffrement a échoué (missing-own-key)')
|
||||
except (NoAvailableSession,) as exn:
|
||||
await self.plain_send(eto, etype, 'erreur: le déchiffrement a échoué (no-available-session)')
|
||||
except (UndecidedException,) as exn:
|
||||
await self.xep_0384.trust(exn.bare_jid, exn.device, exn.ik)
|
||||
return await self.decrypt_message(msg)
|
||||
except (UntrustedException,) as exn:
|
||||
await self.plain_send(eto, etype, 'erreur: le déchiffrement a échoué (untrusted-key)')
|
||||
except (EncryptionPrepareException,):
|
||||
await self.plain_send(eto, etype, 'erreur: le déchiffrement a échoué (encryption-prepare)')
|
||||
except (Exception,) as exn:
|
||||
await self.plain_send(eto, etype, 'erreur: petit soucis...\n%r' % exn)
|
||||
raise
|
||||
|
||||
async def plain_send(self, mto, mtype, body):
|
||||
msg = self.make_message(mto=mto, mtype=mtype)
|
||||
msg['body'] = body
|
||||
return msg.send()
|
||||
|
||||
async def encrypted_send(self, mto, mtype, body, recipients=None):
|
||||
if recipients is None:
|
||||
recipients = [mto]
|
||||
|
||||
msg = self.make_message(mto=mto, mtype=mtype)
|
||||
msg['eme']['namespace'] = EME_NS
|
||||
msg['eme']['name'] = self.xep_0380.mechanisms[EME_NS]
|
||||
|
||||
expect_problems = {}
|
||||
while True:
|
||||
try:
|
||||
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)
|
||||
# TODO: catch NoEligibleDevicesException
|
||||
except EncryptionPrepareException as exn:
|
||||
for error in exn.errors:
|
||||
if isinstance(error, MissingBundleException):
|
||||
err = f'attention: le chiffrement a échoué pour %s (missing-bundle)' % error.bare_jid
|
||||
self.plain_send(mto, mtype, err)
|
||||
jid = JID(error.bare_jid)
|
||||
device_list = expect_problems.setdefault(jid, [])
|
||||
device_list.append(error.device)
|
||||
except (IqError, IqTimeout) as exn:
|
||||
self.plain_send(mto, mtype, 'erreur: je n\'arrive pas à récupérer les informations de chiffrement pour un participant')
|
||||
return
|
||||
except Exception as exn:
|
||||
await self.plain_send(mto, mtype, 'erreur: petit soucis...\n%r' % exn)
|
||||
raise
|
||||
|
||||
if __name__ == '__main__':
|
||||
parser = ArgumentParser(description=Albus.__doc__)
|
||||
|
||||
parser.add_argument("-q", "--quiet", help="set logging to ERROR",
|
||||
action="store_const", dest="loglevel",
|
||||
const=logging.ERROR, default=logging.INFO)
|
||||
parser.add_argument("-d", "--debug", help="set logging to DEBUG",
|
||||
action="store_const", dest="loglevel",
|
||||
const=logging.DEBUG, default=logging.INFO)
|
||||
|
||||
parser.add_argument("-j", "--jid", dest="jid",
|
||||
help="JID to use")
|
||||
parser.add_argument("-p", "--password", dest="password",
|
||||
help="password to use", default="occino3-adtriana0")
|
||||
|
||||
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("Username: ")
|
||||
if args.password is None:
|
||||
args.password = getpass("Password: ")
|
||||
|
||||
os.makedirs(args.data_dir, exist_ok=True)
|
||||
|
||||
xmpp = Albus(args.jid, args.password, os.path.join(args.data_dir, 'albus.db'))
|
||||
xmpp.register_plugin('xep_0030') # Service Discovery
|
||||
xmpp.register_plugin('xep_0045') # Multi-User Chat
|
||||
xmpp.register_plugin('xep_0199') # XMPP Ping
|
||||
xmpp.register_plugin('xep_0249') # Direct MUC Invitations
|
||||
xmpp.register_plugin('xep_0380') # Explicit Message Encryption
|
||||
xmpp.register_plugin('xep_0437') # Room Activity Indicators
|
||||
xmpp.register_plugin('xep_0313') # Message Archive Management
|
||||
|
||||
try:
|
||||
xmpp.register_plugin('xep_0384', {'data_dir': args.data_dir}, module=slixmpp_omemo)
|
||||
except (PluginCouldNotLoad,):
|
||||
log.exception('And error occured when loading the omemo plugin.')
|
||||
sys.exit(1)
|
||||
|
||||
try:
|
||||
xmpp.connect()
|
||||
xmpp.loop.run_until_complete(xmpp.disconnected)
|
||||
except KeyboardInterrupt:
|
||||
xmpp.loop.run_until_complete(xmpp.disconnect())
|
||||
finally:
|
||||
xmpp.loop.run_until_complete(xmpp.on_exit())
|
||||
|
|
@ -0,0 +1,4 @@
|
|||
aiosqlite >= 0.19.0
|
||||
sqlalchemy >= 2.0.0
|
||||
-e git+https://codeberg.org/poezio/slixmpp.git@slix-1.8.5#egg=slixmpp
|
||||
-e git+https://codeberg.org/poezio/slixmpp-omemo.git@v0.9.1#egg=slixmpp-omemo
|
||||
Loading…
Reference in New Issue