echo_client.py 11.4 KB
Newer Older
1 2 3 4 5 6 7 8 9 10 11 12 13 14
#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
    Slixmpp OMEMO plugin
    Copyright (C) 2010  Nathanael C. Fritz
    Copyright (C) 2019 Maxime “pep” Buquet <pep@bouah.net>
    This file is part of slixmpp-omemo.

    See the file LICENSE for copying permission.
"""

import os
import sys
15
import asyncio
16 17 18 19
import logging
from getpass import getpass
from argparse import ArgumentParser

20
from slixmpp import ClientXMPP, JID
21
from slixmpp.exceptions import IqTimeout, IqError
22
from slixmpp.stanza import Message
23
import slixmpp_omemo
24 25
from slixmpp_omemo import PluginCouldNotLoad, MissingOwnKey, EncryptionPrepareException
from slixmpp_omemo import UndecidedException, UntrustedException, NoAvailableSession
26
from omemo.exceptions import MissingBundleException
27 28 29 30 31 32 33 34 35 36

log = logging.getLogger(__name__)


class EchoBot(ClientXMPP):

    """
    A simple Slixmpp bot that will echo encrypted messages it receives, along
    with a short thank you message.

37
    For details on how to build a client with slixmpp, look at examples in the
38 39 40
    slixmpp repository.
    """

41 42
    eme_ns = 'eu.siacs.conversations.axolotl'

43 44 45 46
    def __init__(self, jid, password):
        ClientXMPP.__init__(self, jid, password)

        self.add_event_handler("session_start", self.start)
47
        self.add_event_handler("message", self.message_handler)
48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64

    def start(self, _event) -> None:
        """
        Process the session_start event.

        Typical actions for the session_start event are
        requesting the roster and broadcasting an initial
        presence stanza.

        Arguments:
            event -- An empty dictionary. The session_start
                     event does not provide any additional
                     data.
        """
        self.send_presence()
        self.get_roster()

65 66 67 68
    def message_handler(self, msg: Message) -> None:
        asyncio.ensure_future(self.message(msg))

    async def message(self, msg: Message, allow_untrusted: bool = False) -> None:
69 70 71 72 73 74 75 76 77 78 79
        """
        Process incoming message stanzas. Be aware that this also
        includes MUC messages and error messages. It is usually
        a good idea to check the messages's type before processing
        or sending replies.

        Arguments:
            msg -- The received message stanza. See the documentation
                   for stanza objects and the Message stanza to see
                   how it may be used.
        """
80

81 82 83 84
        if msg['type'] not in ('chat', 'normal'):
            return None

        if not self['xep_0384'].is_encrypted(msg):
85
            await self.plain_reply(msg, 'This message was not encrypted.\n%(body)s' % msg)
86 87
            return None

88
        try:
89 90 91
            mfrom = msg['from']
            encrypted = msg['omemo_encrypted']
            body = self['xep_0384'].decrypt_message(encrypted, mfrom, allow_untrusted)
92
            await self.encrypted_reply(msg, 'Thanks for sending\n%s' % body.decode("utf8"))
93
            return None
94 95 96
        except (MissingOwnKey,):
            # The message is missing our own key, it was not encrypted for
            # us, and we can't decrypt it.
97
            await self.plain_reply(
98 99 100
                msg,
                'I can\'t decrypt this message as it is not encrypted for me.',
            )
101
            return None
102 103 104 105 106 107 108
        except (NoAvailableSession,) as exn:
            # We received a message from that contained a session that we
            # don't know about (deleted session storage, etc.). We can't
            # decrypt the message, and it's going to be lost.
            # Here, as we need to initiate a new encrypted session, it is
            # best if we send an encrypted message directly. XXX: Is it
            # where we talk about self-healing messages?
109
            await self.encrypted_reply(
110 111 112 113
                msg,
                'I can\'t decrypt this message as it uses an encrypted '
                'session I don\'t know about.',
            )
114
            return None
115 116 117 118 119 120 121 122 123
        except (UndecidedException, UntrustedException) as exn:
            # We received a message from an untrusted device. We can
            # choose to decrypt the message nonetheless, with the
            # `allow_untrusted` flag on the `decrypt_message` call, which
            # we will do here. This is only possible for decryption,
            # encryption will require us to decide if we trust the device
            # or not. Clients _should_ indicate that the message was not
            # trusted, or in undecided state, if they decide to decrypt it
            # anyway.
124
            await self.plain_reply(
125 126 127 128
                msg,
                "Your device '%s' is not in my trusted devices." % exn.device,
            )
            # We resend, setting the `allow_untrusted` parameter to True.
129
            await self.message(msg, allow_untrusted=True)
130
            return None
131 132 133 134
        except (EncryptionPrepareException,):
            # Slixmpp tried its best, but there were errors it couldn't
            # resolve. At this point you should have seen other exceptions
            # and given a chance to resolve them already.
135
            await self.plain_reply(msg, 'I was not able to decrypt the message.')
136
            return None
137 138 139
        except (Exception,) as exn:
            await self.plain_reply(msg, 'An error occured while attempting decryption.\n%r' % exn)
            raise
140 141 142

        return None

143
    async def plain_reply(self, original_msg, body):
144 145 146 147 148 149 150 151
        """
        Helper to reply to messages
        """

        mto = original_msg['from']
        mtype = original_msg['type']
        msg = self.make_message(mto=mto, mtype=mtype)
        msg['body'] = body
152
        return msg.send()
153

154
    async def encrypted_reply(self, original_msg, body):
155 156
        """Helper to reply with encrypted messages"""

157 158 159 160 161 162
        mto = original_msg['from']
        mtype = original_msg['type']
        msg = self.make_message(mto=mto, mtype=mtype)
        msg['eme']['namespace'] = self.eme_ns
        msg['eme']['name'] = self['xep_0380'].mechanisms[self.eme_ns]

163 164
        expect_problems = {}  # type: Optional[Dict[JID, List[int]]]

165 166 167 168 169
        while True:
            try:
                # `encrypt_message` excepts the plaintext to be sent, a list of
                # bare JIDs to encrypt to, and optionally a dict of problems to
                # expect per bare JID.
170 171 172 173 174 175
                #
                # Note that this function returns an `<encrypted/>` object,
                # and not a full Message stanza. This combined with the
                # `recipients` parameter that requires for a list of JIDs,
                # allows you to encrypt for 1:1 as well as groupchats (MUC).
                #
176
                # `expect_problems`: See EncryptionPrepareException handling.
177
                recipients = [mto]
178
                encrypt = await self['xep_0384'].encrypt_message(body, recipients, expect_problems)
179 180 181 182 183
                msg.append(encrypt)
                return msg.send()
            except UndecidedException as exn:
                # The library prevents us from sending a message to an
                # untrusted/undecided barejid, so we need to make a decision here.
184 185
                # This is where you prompt your user to ask what to do. In
                # this bot we will automatically trust undecided recipients.
186
                self['xep_0384'].trust(exn.bare_jid, exn.device, exn.ik)
187 188
            # TODO: catch NoEligibleDevicesException
            except EncryptionPrepareException as exn:
189 190 191 192
                # This exception is being raised when the library has tried
                # all it could and doesn't know what to do anymore. It
                # contains a list of exceptions that the user must resolve, or
                # explicitely ignore via `expect_problems`.
193 194 195
                # TODO: We might need to bail out here if errors are the same?
                for error in exn.errors:
                    if isinstance(error, MissingBundleException):
196 197 198 199 200 201 202
                        # We choose to ignore MissingBundleException. It seems
                        # to be somewhat accepted that it's better not to
                        # encrypt for a device if it has problems and encrypt
                        # for the rest, rather than error out. The "faulty"
                        # device won't be able to decrypt and should display a
                        # generic message. The receiving end-user at this
                        # point can bring up the issue if it happens.
203 204 205 206 207 208 209 210 211 212 213 214 215 216
                        self.plain_reply(
                            original_msg,
                            'Could not find keys for device "%d" of recipient "%s". Skipping.' %
                            (error.device, error.bare_jid),
                        )
                        jid = JID(error.bare_jid)
                        device_list = expect_problems.setdefault(jid, [])
                        device_list.append(error.device)
            except (IqError, IqTimeout) as exn:
                self.plain_reply(
                    original_msg,
                    'An error occured while fetching information on a recipient.\n%r' % exn,
                )
                return None
217
            except Exception as exn:
218 219 220 221 222
                await self.plain_reply(
                    original_msg,
                    'An error occured while attempting to encrypt.\n%r' % exn,
                )
                raise
223

224
        return None
225

226

227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290
if __name__ == '__main__':
    # Setup the command line arguments.
    parser = ArgumentParser(description=EchoBot.__doc__)

    # Output verbosity options.
    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)

    # JID and password options.
    parser.add_argument("-j", "--jid", dest="jid",
                        help="JID to use")
    parser.add_argument("-p", "--password", dest="password",
                        help="password to use")

    # Data dir for omemo plugin
    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()

    # Setup logging.
    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: ")

    # Setup the EchoBot and register plugins. Note that while plugins may
    # have interdependencies, the order in which you register them does
    # not matter.

    # Ensure OMEMO data dir is created
    os.makedirs(args.data_dir, exist_ok=True)

    xmpp = EchoBot(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,
        ) # OMEMO
    except (PluginCouldNotLoad,):
        log.exception('And error occured when loading the omemo plugin.')
        sys.exit(1)

    # Connect to the XMPP server and start processing XMPP stanzas.
    xmpp.connect()
    xmpp.process()