qweechat/weechat/testproto.py

   1 # -*- coding: utf-8 -*-
   2 #
   3 # testproto.py - command-line program for testing WeeChat/relay protocol
   4 #
   5 # Copyright (C) 2013-2014 Sébastien Helleu <flashcode@flashtux.org>
   6 #
   7 # This file is part of QWeeChat, a Qt remote GUI for WeeChat.
   8 #
   9 # QWeeChat is free software; you can redistribute it and/or modify
  10 # it under the terms of the GNU General Public License as published by
  11 # the Free Software Foundation; either version 3 of the License, or
  12 # (at your option) any later version.
  13 #
  14 # QWeeChat is distributed in the hope that it will be useful,
  15 # but WITHOUT ANY WARRANTY; without even the implied warranty of
  16 # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  17 # GNU General Public License for more details.
  18 #
  19 # You should have received a copy of the GNU General Public License
  20 # along with QWeeChat.  If not, see <http://www.gnu.org/licenses/>.
  21 #
  22
  23 """
  24 Command-line program for testing WeeChat/relay protocol.
  25 """
  26
  27 from __future__ import print_function
  28
  29 import argparse
  30 import os
  31 import select
  32 import shlex
  33 import socket
  34 import struct
  35 import sys
  36 import time
  37 import traceback
  38
  39 import protocol  # WeeChat/relay protocol
  40
  41 NAME = 'qweechat-testproto'
  42
  43
  44 class TestProto(object):
  45     """Test of WeeChat/relay protocol."""
  46
  47     def __init__(self, args):
  48         self.args = args
  49         self.sock = None
  50         self.has_quit = False
  51         self.address = '{self.args.hostname}/{self.args.port} ' \
  52             '(IPv{0})'.format(6 if self.args.ipv6 else 4, self=self)
  53
  54     def connect(self):
  55         """
  56         Connect to WeeChat/relay.
  57         Return True if OK, False if error.
  58         """
  59         inet = socket.AF_INET6 if self.args.ipv6 else socket.AF_INET
  60         try:
  61             self.sock = socket.socket(inet, socket.SOCK_STREAM)
  62             self.sock.connect((self.args.hostname, self.args.port))
  63         except:
  64             if self.sock:
  65                 self.sock.close()
  66             print('Failed to connect to', self.address)
  67             return False
  68         print('Connected to', self.address)
  69         return True
  70
  71     def send(self, messages):
  72         """
  73         Send a text message to WeeChat/relay.
  74         Return True if OK, False if error.
  75         """
  76         try:
  77             for msg in messages.split('\n'):
  78                 if msg == 'quit':
  79                     self.has_quit = True
  80                 self.sock.sendall(msg + '\n')
  81                 print('\x1b[33m<-- ' + msg + '\x1b[0m')
  82         except:
  83             traceback.print_exc()
  84             print('Failed to send message')
  85             return False
  86         return True
  87
  88     def decode(self, message):
  89         """
  90         Decode a binary message received from WeeChat/relay.
  91         Return True if OK, False if error.
  92         """
  93         try:
  94             proto = protocol.Protocol()
  95             msgd = proto.decode(message,
  96                                 separator='\n' if self.args.verbose > 0
  97                                 else ', ')
  98             print('')
  99             if self.args.verbose >= 2 and msgd.uncompressed:
 100                 # display raw message
 101                 print('\x1b[32m--> message uncompressed ({0} bytes):\n'
 102                       '{1}\x1b[0m'
 103                       ''.format(msgd.size_uncompressed,
 104                                 protocol.hex_and_ascii(msgd.uncompressed, 20)))
 105             # display decoded message
 106             print('\x1b[32m--> {0}\x1b[0m'.format(msgd))
 107         except:
 108             traceback.print_exc()
 109             print('Error while decoding message from WeeChat')
 110             return False
 111         return True
 112
 113     def send_stdin(self):
 114         """
 115         Send commands from standard input if some data is available.
 116         Return True if OK (it's OK if stdin has no commands),
 117         False if error.
 118         """
 119         inr = select.select([sys.stdin], [], [], 0)[0]
 120         if inr:
 121             data = os.read(sys.stdin.fileno(), 4096)
 122             if data:
 123                 if not self.send(data.strip()):
 124                     #self.sock.close()
 125                     return False
 126             # open stdin to read user commands
 127             sys.stdin = open('/dev/tty')
 128         return True
 129
 130     def mainloop(self):
 131         """
 132         Main loop: read keyboard, send commands, read socket,
 133         decode/display binary messages received from WeeChat/relay.
 134         Return 0 if OK, 4 if send error, 5 if decode error.
 135         """
 136         if self.has_quit:
 137             return 0
 138         message = ''
 139         recvbuf = ''
 140         prompt = '\x1b[36mrelay> \x1b[0m'
 141         sys.stdout.write(prompt)
 142         sys.stdout.flush()
 143         try:
 144             while not self.has_quit:
 145                 inr = select.select([sys.stdin, self.sock], [], [], 1)[0]
 146                 for _file in inr:
 147                     if _file == sys.stdin:
 148                         buf = os.read(_file.fileno(), 4096)
 149                         if buf:
 150                             message += buf
 151                             if '\n' in message:
 152                                 messages = message.split('\n')
 153                                 msgsent = '\n'.join(messages[:-1])
 154                                 if msgsent and not self.send(msgsent):
 155                                     return 4
 156                                 message = messages[-1]
 157                                 sys.stdout.write(prompt + message)
 158                                 sys.stdout.flush()
 159                     else:
 160                         buf = _file.recv(4096)
 161                         if buf:
 162                             recvbuf += buf
 163                             while len(recvbuf) >= 4:
 164                                 remainder = None
 165                                 length = struct.unpack('>i', recvbuf[0:4])[0]
 166                                 if len(recvbuf) < length:
 167                                     # partial message, just wait for the
 168                                     # end of message
 169                                     break
 170                                 # more than one message?
 171                                 if length < len(recvbuf):
 172                                     # save beginning of another message
 173                                     remainder = recvbuf[length:]
 174                                     recvbuf = recvbuf[0:length]
 175                                 if not self.decode(recvbuf):
 176                                     return 5
 177                                 if remainder:
 178                                     recvbuf = remainder
 179                                 else:
 180                                     recvbuf = ''
 181                             sys.stdout.write(prompt + message)
 182                             sys.stdout.flush()
 183         except:
 184             traceback.print_exc()
 185             self.send('quit')
 186         return 0
 187
 188     def __del__(self):
 189         print('Closing connection with', self.address)
 190         time.sleep(0.5)
 191         self.sock.close()
 192
 193
 194 def main():
 195     """Main function."""
 196     # parse command line arguments
 197     parser = argparse.ArgumentParser(
 198         formatter_class=argparse.RawDescriptionHelpFormatter,
 199         fromfile_prefix_chars='@',
 200         description='Command-line program for testing WeeChat/relay protocol.',
 201         epilog='''
 202 Environment variable "QWEECHAT_PROTO_OPTIONS" can be set with default options.
 203 Argument "@file.txt" can be used to read default options in a file.
 204
 205 Some commands can be piped to the script, for example:
 206   echo "init password=xxxx" | {name} localhost 5000
 207   {name} localhost 5000 < commands.txt
 208
 209 The script returns:
 210   0: OK
 211   2: wrong arguments (command line)
 212   3: connection error
 213   4: send error (message sent to WeeChat)
 214   5: decode error (message received from WeeChat)
 215 '''.format(name=NAME))
 216     parser.add_argument('-6', '--ipv6', action='store_true',
 217                         help='connect using IPv6')
 218     parser.add_argument('-v', '--verbose', action='count', default=0,
 219                         help='verbose mode: long objects view '
 220                         '(-vv: display raw messages)')
 221     parser.add_argument('hostname',
 222                         help='hostname (or IP address) of machine running '
 223                         'WeeChat/relay')
 224     parser.add_argument('port', type=int,
 225                         help='port of machine running WeeChat/relay')
 226     if len(sys.argv) == 1:
 227         parser.print_help()
 228         sys.exit(0)
 229     _args = parser.parse_args(
 230         shlex.split(os.getenv('QWEECHAT_PROTO_OPTIONS') or '') + sys.argv[1:])
 231
 232     test = TestProto(_args)
 233
 234     # connect to WeeChat/relay
 235     if not test.connect():
 236         sys.exit(3)
 237
 238     # send commands from standard input if some data is available
 239     if not test.send_stdin():
 240         sys.exit(4)
 241
 242     # main loop (wait commands, display messages received)
 243     returncode = test.mainloop()
 244     del test
 245     sys.exit(returncode)
 246
 247
 248 if __name__ == "__main__":
 249     main()