Initial commit

weercd.py

+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+#
+# Copyright (C) 2011-2014 Sébastien Helleu <flashcode@flashtux.org>
+#
+# This file is part of weercd, the WeeChat IRC testing server.
+#
+# weercd 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.
+#
+# weercd 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 weercd.  If not, see <http://www.gnu.org/licenses/>.
+#
+
+"""
+WeeChat IRC testing server.
+"""
+
+from __future__ import division, print_function
+
+import argparse
+import os
+import random
+import re
+import select
+import shlex
+import socket
+import string
+import sys
+import time
+import traceback
+
+NAME = 'weercd'
+VERSION = '0.8'
+
+
+def fuzzy_string(minlength=1, maxlength=50, spaces=False):
+    """Return a fuzzy string (random length and content)."""
+    length = random.randint(minlength, maxlength)
+    strspace = ''
+    if spaces:
+        strspace = ' '
+    return ''.join(random.choice(string.ascii_uppercase +
+                                 string.ascii_lowercase +
+                                 string.digits + strspace)
+                   for x in range(length))
+
+
+def fuzzy_host():
+    """Return a fuzzy host name."""
+    return '{0}@{1}'.format(fuzzy_string(1, 10), fuzzy_string(1, 10))
+
+
+def fuzzy_channel():
+    """Return a fuzzy channel name."""
+    return '#{0}'.format(fuzzy_string(1, 25))
+
+
+class Client(object):  # pylint: disable=too-many-instance-attributes
+    """A client of weercd server."""
+
+    def __init__(self, sock, addr, args):
+        self.sock, self.addr = sock, addr
+        self.args = args
+        self.name = NAME
+        self.version = VERSION
+        self.nick = ''
+        self.nicknumber = 0
+        self.channels = {}
+        self.lastbuf = ''
+        self.incount, self.outcount, self.inbytes, self.outbytes = 0, 0, 0, 0
+        self.quit, self.endmsg, self.endexcept = False, '', None
+        self.starttime = time.time()
+        self.connect()
+
+    def run(self):
+        """Execute the action asked for the client."""
+        if self.quit:
+            return
+
+        # send commands from file (which can be stdin)
+        if self.args.file:
+            self.send_file()
+            return
+
+        # flood the client
+        if self.args.wait > 0:
+            print('Waiting', self.args.wait, 'seconds')
+            time.sleep(self.args.wait)
+        sys.stdout.write('Flooding client..')
+        sys.stdout.flush()
+        try:
+            while not self.quit:
+                self.flood()
+        except Exception as exc:
+            if self.quit:
+                self.endmsg = 'quit received'
+            else:
+                self.endmsg = 'connection lost'
+            self.endexcept = exc
+        except KeyboardInterrupt:
+            self.endmsg = 'interrupted'
+        else:
+            self.endmsg = 'quit received'
+
+    def fuzzy_nick(self, with_number=False):
+        """Return a fuzzy nick name."""
+        if with_number:
+            self.nicknumber += 1
+            return '{0}{1}'.format(fuzzy_string(1, 5), self.nicknumber)
+        else:
+            return fuzzy_string(1, 10)
+
+    def send(self, data):
+        """Send one message to client."""
+        if self.args.debug:
+            print('<--', data)
+        msg = data + '\r\n'
+        self.outbytes += len(msg)
+        self.sock.send(msg.encode('UTF-8'))
+        self.outcount += 1
+
+    # pylint: disable=too-many-arguments
+    def send_cmd(self, cmd, data, nick='{self.name}', host='',
+                 target='{self.nick}'):
+        """Send an IRC command to the client."""
+        self.send(':{0}{1}{2} {3}{4}{5}{6}{7}'
+                  ''.format(nick,
+                            '!' if host else '',
+                            host,
+                            cmd,
+                            ' ' if target else '',
+                            target,
+                            ' :' if data else '',
+                            data).format(self=self))
+
+    def recv(self, data):
+        """Read one IRC message from client."""
+        if self.args.debug:
+            print('-->', data)
+        if data.startswith('PING '):
+            args = data[5:]
+            if args[0] == ':':
+                args = args[1:]
+            self.send('PONG :{0}'.format(args))
+        elif data.startswith('NICK '):
+            self.nick = data[5:]
+        elif data.startswith('PART '):
+            match = re.search('^PART :?(#[^ ]+)', data)
+            if match:
+                channel = match.group(1)
+                if channel in self.channels:
+                    del self.channels[channel]
+        elif data.startswith('QUIT '):
+            self.quit = True
+        self.incount += 1
+
+    def read(self, timeout):
+        """Read raw data received from client."""
+        inr = select.select([self.sock], [], [], timeout)[0]
+        if inr:
+            data = self.sock.recv(4096)
+            if data:
+                data = data.decode('UTF-8')
+                self.inbytes += len(data)
+                data = self.lastbuf + data
+                while True:
+                    pos = data.find('\r\n')
+                    if pos < 0:
+                        break
+                    self.recv(data[0:pos])
+                    data = data[pos + 2:]
+                self.lastbuf = data
+
+    def connect(self):
+        """Inform the client that the connection is OK."""
+        try:
+            count = self.args.nickused
+            while self.nick == '':
+                self.read(0.1)
+                if self.nick and count > 0:
+                    self.send_cmd('433', 'Nickname is already in use.',
+                                  target='* {self.nick}')
+                    self.nick = ''
+                    count -= 1
+            self.send_cmd('001', 'Welcome to the WeeChat IRC server')
+            self.send_cmd('002', 'Your host is {self.name}, running version '
+                          '{self.version}')
+            self.send_cmd('003', 'Are you solid like a rock?')
+            self.send_cmd('004', 'Let\'s see!')
+        except KeyboardInterrupt:
+            self.quit = True
+            self.endmsg = 'interrupted'
+            return
+
+    def channel_random_nick(self, channel):
+        """Return a random nick of a channel."""
+        if len(self.channels[channel]) < 2:
+            return None
+        rnick = self.nick
+        while rnick == self.nick:
+            rnick = self.channels[channel][
+                random.randint(0, len(self.channels[channel]) - 1)]
+        return rnick
+
+    def flood_self_join(self):
+        """Self join on a new channel."""
+        channel = fuzzy_channel()
+        if channel in self.channels:
+            return
+        self.send_cmd('JOIN', channel,
+                      nick=self.nick, host=self.addr[0], target='')
+        self.send_cmd('353', '@{self.nick}',
+                      target='{0} = {1}'.format(self.nick, channel))
+        self.send_cmd('366', 'End of /NAMES list.',
+                      target='{0} {1}'.format(self.nick, channel))
+        self.channels[channel] = [self.nick]
+
+    def flood_user_notice(self):
+        """Notice for the user."""
+        self.send_cmd('NOTICE', fuzzy_string(1, 400, spaces=True),
+                      nick=self.fuzzy_nick(), host=fuzzy_host())
+
+    def flood_channel_join(self, channel):
+        """Join of a user in a channel."""
+        if len(self.channels[channel]) >= self.args.maxnicks:
+            return
+        newnick = self.fuzzy_nick(with_number=True)
+        self.send_cmd('JOIN', channel,
+                      nick=newnick, host=fuzzy_host(), target='')
+        self.channels[channel].append(newnick)
+
+    def flood_channel_part(self, channel):
+        """Part or quit of a user in a channel."""
+        if len(self.channels[channel]) == 0:
+            return
+        rnick = self.channel_random_nick(channel)
+        if not rnick:
+            return
+        if random.randint(1, 2) == 1:
+            self.send_cmd('PART', channel,
+                          nick=rnick, host=fuzzy_host(), target='')
+        else:
+            self.send_cmd('QUIT', fuzzy_string(1, 30),
+                          nick=rnick, host=fuzzy_host(), target='')
+        self.channels[channel].remove(rnick)
+
+    def flood_channel_kick(self, channel):
+        """Kick of a user in a channel."""
+        if len(self.channels[channel]) == 0:
+            return
+        rnick1 = self.channel_random_nick(channel)
+        rnick2 = self.channel_random_nick(channel)
+        if rnick1 and rnick2 and rnick1 != rnick2:
+            self.send_cmd('KICK', fuzzy_string(1, 50),
+                          nick=rnick1, host=fuzzy_host(),
+                          target='{0} {1}'.format(channel, rnick2))
+            self.channels[channel].remove(rnick2)
+
+    def flood_channel_message(self, channel):
+        """Message from a user in a channel."""
+        if len(self.channels[channel]) == 0:
+            return
+        rnick = self.channel_random_nick(channel)
+        if not rnick:
+            return
+        msg = fuzzy_string(1, 400, spaces=True)
+        if 'channel' in self.args.notice and random.randint(1, 100) == 100:
+            # notice for channel
+            self.send_cmd('NOTICE', msg,
+                          nick=rnick, host=fuzzy_host(), target=channel)
+        else:
+            # add random highlight
+            if random.randint(1, 100) == 100:
+                msg = '{0}: {1}'.format(self.nick, msg)
+            action2 = random.randint(1, 50)
+            if action2 == 1:
+                # CTCP action (/me)
+                msg = '\x01ACTION {0}\x01'.format(msg)
+            elif action2 == 2:
+                # CTCP version
+                msg = '\x01VERSION\x01'
+            self.send_cmd('PRIVMSG', msg,
+                          nick=rnick, host=fuzzy_host(), target=channel)
+
+    def flood(self):
+        """Yay, funny stuff here! Flood the client!"""
+        self.read(self.args.sleep)
+        # global actions
+        action = random.randint(1, 2)
+        if action == 1 and len(self.channels) < self.args.maxchans:
+            self.flood_self_join()
+        elif action == 2 and 'user' in self.args.notice:
+            self.flood_user_notice()
+        # actions for each channel
+        for channel in self.channels:
+            action = random.randint(1, 50)
+            if 1 <= action <= 10:
+                self.flood_channel_join(channel)
+            elif action == 11:
+                self.flood_channel_part(channel)
+            elif action == 12:
+                self.flood_channel_kick(channel)
+            else:
+                self.flood_channel_message(channel)
+        # display progress
+        if self.outcount % 1000 == 0:
+            sys.stdout.write('.')
+            sys.stdout.flush()
+
+    def send_file(self):
+        """Send messages from a file to client."""
+        stdin = self.args.file == sys.stdin
+        count = 0
+        self.read(0.2)
+        try:
+            while True:
+                # display the prompt if we are reading in stdin
+                if stdin:
+                    sys.stdout.write('Message to send to client: ')
+                    sys.stdout.flush()
+                message = self.args.file.readline()
+                if not message:
+                    break
+                if sys.version_info < (3,):
+                    message = message.decode('UTF-8')
+                message = message.rstrip('\n')
+                if message and not message.startswith('//'):
+                    self.send(message.format(self=self))
+                    count += 1
+                self.read(0.1 if stdin else self.args.sleep)
+        except IOError as exc:
+            self.endmsg = 'unable to read file {0}'.format(self.args.file)
+            self.endexcept = exc
+            return
+        except Exception as exc:
+            traceback.print_exc()
+            self.endmsg = 'connection lost'
+            return
+        except KeyboardInterrupt:
+            self.endmsg = 'interrupted'
+            return
+        finally:
+            sys.stdout.write('\n')
+            sys.stdout.write('{0} messages sent from {1}, press Enter to exit'
+                             .format(count, 'stdin' if stdin else 'file'))
+            sys.stdout.flush()
+            try:
+                sys.stdin.readline()
+            except Exception:  # pylint: disable=broad-except
+                pass
+
+    def stats(self):
+        """Display some statistics about data exchanged with the client."""
+        msgexcept = ''
+        if self.endexcept:
+            msgexcept = '({0})'.format(self.endexcept)
+        print(self.endmsg, msgexcept)
+        elapsed = time.time() - self.starttime
+        countrate = self.outcount / elapsed
+        bytesrate = self.outbytes / elapsed
+        print('Elapsed: {elapsed:.1f}s - '
+              'packets: in:{self.incount}, out:{self.outcount} '
+              '({countrate:.0f}/s) - '
+              'bytes: in:{self.inbytes}, out: {self.outbytes} '
+              '({bytesrate:.0f}/s)'
+              ''.format(self=self,
+                        elapsed=elapsed,
+                        countrate=countrate,
+                        bytesrate=bytesrate))
+        if self.endmsg == 'connection lost':
+            print('Uh-oh! No quit received, client has crashed? Ahah \\o/')
+
+    def __del__(self):
+        self.stats()
+        print('Closing connection with', self.addr)
+        self.sock.close()
+
+
+def main():
+    """Main function."""
+    # parse command line arguments
+    parser = argparse.ArgumentParser(
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+        fromfile_prefix_chars='@',
+        description='The WeeChat IRC testing server.',
+        epilog='Note: the environment variable "WEERCD_OPTIONS" can be '
+        'set with default options. Argument "@file.txt" can be used to read '
+        'default options in a file.')
+    parser.add_argument('-H', '--host', help='host for socket bind')
+    parser.add_argument('-p', '--port', type=int, default=7777,
+                        help='port for socket bind')
+    parser.add_argument('-f', '--file', type=argparse.FileType('r'),
+                        help='send messages from file, instead of flooding '
+                        'the client (use "-" for stdin)')
+    parser.add_argument('-c', '--maxchans', type=int, default=5,
+                        help='max number of channels to join')
+    parser.add_argument('-n', '--maxnicks', type=int, default=100,
+                        help='max number of nicks per channel')
+    parser.add_argument('-u', '--nickused', type=int, default=0,
+                        help='send 433 (nickname already in use) this number '
+                        'of times before accepting nick')
+    parser.add_argument('-N', '--notice', metavar='NOTICE_TYPE',
+                        choices=['user', 'channel'],
+                        default=['user', 'channel'], nargs='*',
+                        help='notices to send: "user" (to user), "channel" '
+                        '(to channel)')
+    parser.add_argument('-s', '--sleep', type=float, default=0,
+                        help='sleep for select: delay between 2 messages sent '
+                        'to client (float, in seconds)')
+    parser.add_argument('-w', '--wait', type=float, default=0,
+                        help='time to wait before flooding client (float, '
+                        'in seconds)')
+    parser.add_argument('-d', '--debug', action='store_true',
+                        help='debug output')
+    parser.add_argument('-v', '--version', action='version', version=VERSION)
+    args = parser.parse_args(shlex.split(os.getenv('WEERCD_OPTIONS') or '') +
+                             sys.argv[1:])
+
+    # welcome message, with options
+    print(NAME, VERSION, '- WeeChat IRC testing server')
+    print('Options:', vars(args))
+
+    # main loop
+    while True:
+        servsock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+        try:
+            servsock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
+            servsock.bind((args.host or '', args.port))
+            servsock.listen(1)
+        except Exception as exc:  # pylint: disable=broad-except
+            print('Socket error: {0}'.format(exc))
+            sys.exit(1)
+        print('Listening on port', args.port, '(ctrl-C to exit)')
+        clientsock = None
+        addr = None
+        try:
+            clientsock, addr = servsock.accept()
+        except KeyboardInterrupt:
+            servsock.close()
+            sys.exit(0)
+        print('Connection from', addr)
+        client = Client(clientsock, addr, args)
+        client.run()
+        del client
+        # no loop if message were sent from a file
+        if args.file:
+            break
+
+if __name__ == "__main__":
+    main()