Add SASL auth and SSL transport mixins

bd808 · bd808 · commit 34f5b132f7ea · 2017-02-20T10:43:46.000-07:00
* ib3.SASL: Authenticate using SASL before joining channels
* ib3.SSL: Use SSL on connections to servers
diff --git a/examples/nickserv.py b/examples/nickserv.py
@@ -0,0 +1,67 @@
+# -*- coding: utf-8 -*-
+#
+# This file is part of IRC Bot Behavior Bundle (IB3)
+# Copyright (C) 2017 Bryan Davis and contributors
+#
+# This program is free software: you can redistribute it and/or modify it
+# under the terms of the GNU General Public License as published by the Free
+# Software Foundation, either version 3 of the License, or (at your option)
+# any later version.
+#
+# This program is distributed in the hope that it will be useful, but WITHOUT
+# ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
+# FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public License for
+# more details.
+#
+# You should have received a copy of the GNU General Public License along with
+# this program.  If not, see <http://www.gnu.org/licenses/>.
+
+import argparse
+import logging
+
+from ib3 import Bot
+from ib3 import DisconnectOnError
+from ib3 import FreenodePasswdAuth
+from ib3 import Ping
+from ib3 import RejoinOnBan
+from ib3 import RejoinOnKick
+
+
+logging.basicConfig(
+    level=logging.DEBUG,
+    format='%(asctime)s %(name)s %(levelname)s: %(message)s',
+    datefmt='%Y-%m-%dT%H:%M:%SZ'
+)
+logging.captureWarnings(True)
+
+
+class TestBot(
+    FreenodePasswdAuth, Ping, DisconnectOnError,
+    RejoinOnBan, RejoinOnKick, Bot
+):
+    pass
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser(
+        description='Example bot with NickServ auth')
+    parser.add_argument('nick')
+    parser.add_argument('password')
+    parser.add_argument('channel')
+    args = parser.parse_args()
+
+    bot = TestBot(
+        server_list=[('chat.freenode.net', 6667)],
+        nickname=args.nick,
+        realname=args.nick,
+        ident_password=args.password,
+        channels=[args.channel]
+    )
+    try:
+        bot.start()
+    except KeyboardInterrupt:
+        bot.disconnect('KeyboardInterrupt!')
+    except Exception:
+        logging.getLogger('root').exception('Killed by unhandled exception')
+        bot.disconnect('Exception!')
+        raise SystemExit()
diff --git a/examples/saslbot.py b/examples/saslbot.py
@@ -0,0 +1,83 @@
+# -*- coding: utf-8 -*-
+#
+# This file is part of IRC Bot Behavior Bundle (IB3)
+# Copyright (C) 2017 Bryan Davis and contributors
+#
+# This program is free software: you can redistribute it and/or modify it
+# under the terms of the GNU General Public License as published by the Free
+# Software Foundation, either version 3 of the License, or (at your option)
+# any later version.
+#
+# This program is distributed in the hope that it will be useful, but WITHOUT
+# ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
+# FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public License for
+# more details.
+#
+# You should have received a copy of the GNU General Public License along with
+# this program.  If not, see <http://www.gnu.org/licenses/>.
+
+import argparse
+import logging
+
+import irc.strings
+
+import ib3
+
+
+logging.basicConfig(
+    level=logging.DEBUG,
+    format='%(asctime)s %(name)s %(levelname)s: %(message)s',
+    datefmt='%Y-%m-%dT%H:%M:%SZ'
+)
+logging.captureWarnings(True)
+
+logger = logging.getLogger('saslbot')
+
+
+class SaslBot(ib3.SASL, ib3.SSL, ib3.Bot):
+    """Example bot showing use of SASL auth and SSL encryption."""
+    def on_privmsg(self, conn, event):
+        self.do_command(conn, event, event.arguments[0])
+
+    def on_pubmsg(self, conn, event):
+        args = event.arguments[0].split(':', 1)
+        if len(args) > 1:
+            to = irc.strings.lower(args[0])
+            if to == irc.strings.lower(conn.get_nickname()):
+                self.do_command(conn, event, args[1].strip())
+
+    def do_command(self, conn, event, cmd):
+        to = event.target
+        if to == conn.get_nickname():
+            to = event.source.nick
+
+        if cmd == 'disconnect':
+            self.disconnect()
+        elif cmd == 'die':
+            self.die()
+        else:
+            conn.privmsg(to, 'What does "{}" mean?'.format(cmd))
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser(description='Example bot with SASL auth')
+    parser.add_argument('nick')
+    parser.add_argument('password')
+    parser.add_argument('channel')
+    args = parser.parse_args()
+
+    bot = SaslBot(
+        server_list=[('chat.freenode.net', 6697)],
+        nickname=args.nick,
+        realname=args.nick,
+        ident_password=args.password,
+        channels=[args.channel],
+    )
+    try:
+        bot.start()
+    except KeyboardInterrupt:
+        bot.disconnect('KeyboardInterrupt!')
+    except Exception:
+        logger.exception('Killed by unhandled exception')
+        bot.disconnect('Exception!')
+        raise SystemExit()
diff --git a/ib3/__init__.py b/ib3/__init__.py
@@ -17,12 +17,15 @@
 # this program.  If not, see <http://www.gnu.org/licenses/>.
 """Base IRC bot and mixins for commonly desired functionality."""
 
+import base64
 import functools
 import logging
+import ssl
 
 import irc.bot
 import irc.buffer
 import irc.client
+import irc.connection
 
 logger = logging.getLogger(__name__)
 
@@ -112,11 +115,27 @@ def _handle_bannedfromchan(self, conn, event):
             60, functools.partial(conn.join, event.arguments[0]))
 
 
-class FreenodePasswdAuth(object):
+class JoinChannels(object):
+    """Join channels one at a time to avoid flooding."""
+    def join_channels(self, channels):
+        """Join a list of channels, one at a time."""
+        try:
+            car, cdr = channels[0], channels[1:]
+        except (IndexError, TypeError):
+            logger.exception('Failed to find channel to join.')
+        else:
+            logger.info('Joining %s', car)
+            self.connection.join(car)
+            if cdr:
+                self.reactor.scheduler.execute_after(
+                    1, functools.partial(self.join_channels, cdr))
+
+
+class FreenodePasswdAuth(JoinChannels):
     """Authenticate with NickServ before joining channels."""
     def __init__(
             self, server_list, nickname, realname,
-            ident_password, channels, **kwargs):
+            ident_password, channels=[], **kwargs):
         """
         :param server_list: List of servers the bot will use.
         :param nickname: The bot's nickname
@@ -178,8 +197,7 @@ def _handle_privnotice(self, conn, event):
                 self._identify_to_nickserv()
             elif 'You are now identified' in msg:
                 logger.debug('Authentication succeeded')
-                self.reactor.scheduler.execute_after(
-                    1, self._join_next_channel)
+                self.join_channels(self._channels)
             elif 'Invalid password' in msg:
                 logger.error('Password invalid. Check your config!')
                 self.die()
@@ -190,21 +208,6 @@ def _identify_to_nickserv(self):
         self.connection.privmsg('NickServ', 'identify %s %s' % (
             self._primary_nick, self._ident_password))
 
-    def _join_next_channel(self, channels=None):
-        """Join the next channel in our join list."""
-        if channels is None:
-            channels = self._channels
-        try:
-            car, cdr = channels[0], channels[1:]
-        except (IndexError, TypeError):
-            logger.exception('Failed to find channel to join.')
-        else:
-            logger.info('Joining %s', car)
-            self.connection.join(car)
-            if cdr:
-                self.reactor.scheduler.execute_after(
-                    1, functools.partial(self._join_next_channel, cdr))
-
     def _nickserv_regain(self):
         if not self.has_primary_nick():
             # REGAIN disconnects an old user session, or somebody
@@ -218,3 +221,79 @@ def _nickserv_regain(self):
     def has_primary_nick(self):
         """Do we currently have our primary nick?"""
         return self.connection.get_nickname() == self._primary_nick
+
+
+class SASL(JoinChannels):
+    def __init__(
+            self, server_list, nickname, realname,
+            ident_password, channels=[], username=None, **kwargs):
+        """
+        :param server_list: List of servers the bot will use.
+        :param nickname: The bot's nickname
+        :param realname: The bot's realname
+        :param ident_password: The bot's password
+        :param channels: List of channels to join after authenticating
+        """
+        self._primary_nick = nickname
+        self._username = username or nickname
+        self._ident_password = ident_password
+        self._channels = channels
+
+        super(SASL, self).__init__(
+                server_list=server_list,
+                nickname=nickname,
+                realname=realname,
+                username=self._username,
+                **kwargs)
+        self.reactor._on_connect = self._handle_connect
+
+        for event in ['cap', 'authenticate', '903', '908', 'welcome']:
+            logger.debug('Registering for %s', event)
+            self.connection.add_global_handler(
+                event, getattr(self, '_handle_%s' % event))
+
+    def _handle_connect(self, sock):
+        """Send CAP REQ :sasl on connect."""
+        self.connection.cap('REQ', 'sasl')
+
+    def _handle_cap(self, conn, event):
+        """Handle CAP responses."""
+        if event.arguments and event.arguments[0] == 'ACK':
+            conn.send_raw('AUTHENTICATE PLAIN')
+        else:
+            logger.warning('Unexpcted CAP response: %s', event)
+            conn.disconnect()
+
+    def _handle_authenticate(self, conn, event):
+        """Handle AUTHENTICATE responses."""
+        if event.target == '+':
+            creds = '{username}\0{username}\0{password}'.format(
+                username=self._username,
+                password=self._ident_password)
+            conn.send_raw('AUTHENTICATE {}'.format(
+                    base64.b64encode(creds.encode('utf8')).decode('utf8')))
+        else:
+            logger.warning('Unexpcted AUTHENTICATE response: %s', event)
+            conn.disconnect()
+
+    def _handle_903(self, conn, event):
+        """Handle 903 RPL_SASLSUCCESS responses."""
+        self.connection.cap('END')
+
+    def _handle_908(self, conn, event):
+        """Handle 908 RPL_SASLMECHS responses."""
+        logger.warning('SASL PLAIN not supported: %s', event)
+        self.die()
+
+    def _handle_welcome(self, conn, event):
+        """Handle WELCOME message."""
+        logger.info('Connected to server %s', conn.get_server_name())
+        self.join_channels(self._channels)
+
+
+class SSL(object):
+    """Use SSL connections."""
+    def __init__(self, *args, **kwargs):
+        kwargs['connect_factory'] = irc.connection.Factory(
+            wrapper=ssl.wrap_socket)
+        super(SSL, self).__init__(*args, **kwargs)
