otr-bot.py 7.44 KB
Newer Older
1 2 3 4
#!/usr/bin/python
import sys
import jabberbot
import xmpp
Tails developers's avatar
Tails developers committed
5
import potr
kytv's avatar
kytv committed
6
import logging
Tails developers's avatar
Tails developers committed
7
from argparse import ArgumentParser
8

Tails developers's avatar
Tails developers committed
9
class OtrContext(potr.context.Context):
10

Tails developers's avatar
Tails developers committed
11 12
    def __init__(self, account, peer):
        super(OtrContext, self).__init__(account, peer)
13

Tails developers's avatar
Tails developers committed
14 15
    def getPolicy(self, key):
        return True
16

Tails developers's avatar
Tails developers committed
17 18 19 20
    def inject(self, msg, appdata = None):
        mess = appdata["base_reply"]
        mess.setBody(msg)
        appdata["send_raw_message_fn"](mess)
21 22


Tails developers's avatar
Tails developers committed
23
class BotAccount(potr.context.Account):
24

Tails developers's avatar
Tails developers committed
25 26 27 28 29
    def __init__(self, jid, keyFilePath):
        protocol = 'xmpp'
        max_message_size = 10*1024
        super(BotAccount, self).__init__(jid, protocol, max_message_size)
        self.keyFilePath = keyFilePath
30

Tails developers's avatar
Tails developers committed
31 32 33
    def loadPrivkey(self):
        with open(self.keyFilePath, 'rb') as keyFile:
            return potr.crypt.PK.parsePrivateKey(keyFile.read())[0]
34 35


Tails developers's avatar
Tails developers committed
36 37 38 39 40 41 42 43 44 45 46 47 48 49
class OtrContextManager:

    def __init__(self, jid, keyFilePath):
        self.account = BotAccount(jid, keyFilePath)
        self.contexts = {}

    def start_context(self, other):
        if not other in self.contexts:
            self.contexts[other] = OtrContext(self.account, other)
        return self.contexts[other]

    def get_context_for_user(self, other):
        return self.start_context(other)

50 51 52 53 54

class OtrBot(jabberbot.JabberBot):

    PING_FREQUENCY = 60

anonym's avatar
anonym committed
55 56
    def __init__(self, account, password, otr_key_path,
                 connect_server = None, log_file = None):
57 58
        self.__connect_server = connect_server
        self.__password = password
anonym's avatar
anonym committed
59
        self.__log_file = log_file
Tails developers's avatar
Tails developers committed
60
        super(OtrBot, self).__init__(account, password)
Tails developers's avatar
Tails developers committed
61 62 63 64
        self.__otr_manager = OtrContextManager(account, otr_key_path)
        self.send_raw_message_fn = super(OtrBot, self).send_message
        self.__default_otr_appdata = {
            "send_raw_message_fn": self.send_raw_message_fn
65
            }
Tails developers's avatar
Tails developers committed
66 67 68 69 70

    def __otr_appdata_for_mess(self, mess):
        appdata = self.__default_otr_appdata.copy()
        appdata["base_reply"] = mess
        return appdata
71

72 73 74 75 76
    # Unfortunately Jabberbot's connect() is not very friendly to
    # overriding in subclasses so we have to re-implement it
    # completely (copy-paste mostly) in order to add support for using
    # an XMPP "Connect Server".
    def connect(self):
anonym's avatar
anonym committed
77 78
        logging.basicConfig(filename = self.__log_file,
                            level = logging.DEBUG)
79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101
        if not self.conn:
            conn = xmpp.Client(self.jid.getDomain(), debug=[])
            if self.__connect_server:
                try:
                    conn_server, conn_port = self.__connect_server.split(":", 1)
                except ValueError:
                    conn_server = self.__connect_server
                    conn_port = 5222
                conres = conn.connect((conn_server, int(conn_port)))
            else:
                conres = conn.connect()
            if not conres:
                return None
            authres = conn.auth(self.jid.getNode(), self.__password, self.res)
            if not authres:
                return None
            self.conn = conn
            self.conn.sendInitPresence()
            self.roster = self.conn.Roster.getRoster()
            for (handler, callback) in self.handlers:
                self.conn.RegisterHandler(handler, callback)
        return self.conn

102 103 104 105 106
    # Wrap OTR encryption around Jabberbot's most low-level method for
    # sending messages.
    def send_message(self, mess):
        body = str(mess.getBody())
        user = str(mess.getTo().getStripped())
Tails developers's avatar
Tails developers committed
107 108 109 110 111 112
        otrctx = self.__otr_manager.get_context_for_user(user)
        if otrctx.state == potr.context.STATE_ENCRYPTED:
            otrctx.sendMessage(potr.context.FRAGMENT_SEND_ALL, body,
                               appdata = self.__otr_appdata_for_mess(mess))
        else:
            self.send_raw_message_fn(mess)
113 114 115 116 117

    # Wrap OTR decryption around Jabberbot's callback mechanism.
    def callback_message(self, conn, mess):
        body = str(mess.getBody())
        user = str(mess.getFrom().getStripped())
Tails developers's avatar
Tails developers committed
118 119 120 121 122 123 124 125 126 127 128 129 130 131 132
        otrctx = self.__otr_manager.get_context_for_user(user)
        if mess.getType() == "chat":
            try:
                appdata = self.__otr_appdata_for_mess(mess.buildReply())
                decrypted_body, tlvs = otrctx.receiveMessage(body,
                                                             appdata = appdata)
                otrctx.processTLVs(tlvs)
            except potr.context.NotEncryptedError:
                otrctx.authStartV2(appdata = appdata)
                return
            except (potr.context.UnencryptedMessage, potr.context.NotOTRMessage):
                decrypted_body = body
        else:
            decrypted_body = body
        if decrypted_body == None:
133
            return
134
        if mess.getType() == "groupchat":
Tails developers's avatar
Tails developers committed
135
            bot_prefix = self.jid.getNode() + ": "
136 137 138 139
            if decrypted_body.startswith(bot_prefix):
                decrypted_body = decrypted_body[len(bot_prefix):]
            else:
                return
140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158
        mess.setBody(decrypted_body)
        super(OtrBot, self).callback_message(conn, mess)

    # Override Jabberbot quitting on keep alive failure.
    def on_ping_timeout(self):
        self.__lastping = None

    @jabberbot.botcmd
    def ping(self, mess, args):
        """Why not just test it?"""
        return "pong"

    @jabberbot.botcmd
    def say(self, mess, args):
        """Unleash my inner parrot"""
        return args

    @jabberbot.botcmd
    def clear_say(self, mess, args):
Tails developers's avatar
Tails developers committed
159
        """Make me speak in the clear even if we're in an OTR chat"""
Tails developers's avatar
Tails developers committed
160
        self.send_raw_message_fn(mess.buildReply(args))
161 162 163 164 165
        return ""

    @jabberbot.botcmd
    def start_otr(self, mess, args):
        """Make me *initiate* (but not refresh) an OTR session"""
166 167
        if mess.getType() == "groupchat":
            return
168 169 170 171 172
        return "?OTRv2?"

    @jabberbot.botcmd
    def end_otr(self, mess, args):
        """Make me gracefully end the OTR session if there is one"""
173 174
        if mess.getType() == "groupchat":
            return
175
        user = str(mess.getFrom().getStripped())
Tails developers's avatar
Tails developers committed
176 177
        self.__otr_manager.get_context_for_user(user).disconnect(appdata =
            self.__otr_appdata_for_mess(mess.buildReply()))
178 179 180
        return ""

if __name__ == '__main__':
Tails developers's avatar
Tails developers committed
181 182 183 184 185 186 187
    parser = ArgumentParser()
    parser.add_argument("account",
                        help = "the user account, given as user@domain")
    parser.add_argument("password",
                        help = "the user account's password")
    parser.add_argument("otr_key_path",
                        help = "the path to the account's OTR key file")
188 189 190
    parser.add_argument("-c", "--connect-server", metavar = 'ADDRESS',
                        help = "use a Connect Server, given as host[:port] " +
                        "(port defaults to 5222)")
Tails developers's avatar
Tails developers committed
191 192
    parser.add_argument("-j", "--auto-join", nargs = '+', metavar = 'ROOMS',
                        help = "auto-join multi-user chatrooms on start")
anonym's avatar
anonym committed
193 194
    parser.add_argument("-l", "--log-file", metavar = 'LOGFILE',
                        help = "Log to file instead of stderr")
Tails developers's avatar
Tails developers committed
195
    args = parser.parse_args()
196 197 198
    otr_bot_opt_args = dict()
    if args.connect_server:
        otr_bot_opt_args["connect_server"] = args.connect_server
anonym's avatar
anonym committed
199 200
    if args.log_file:
        otr_bot_opt_args["log_file"] = args.log_file
201 202
    otr_bot = OtrBot(args.account, args.password, args.otr_key_path,
                     **otr_bot_opt_args)
Tails developers's avatar
Tails developers committed
203 204 205
    if args.auto_join:
        for room in args.auto_join:
            otr_bot.join_room(room)
206
    otr_bot.serve_forever()