albus/albus.py

#!/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())