src/qweechat/weechat/testproto.py

   1 #!/usr/bin/python
   2 # -*- coding: utf-8 -*-
   3 #
   4 # testproto.py - command-line program for testing protocol WeeChat/relay
   5 #
   6 # Copyright (C) 2013 Sebastien Helleu <flashcode@flashtux.org>
   7 #
   8 # This file is part of QWeeChat, a Qt remote GUI for WeeChat.
   9 #
  10 # QWeeChat is free software; you can redistribute it and/or modify
  11 # it under the terms of the GNU General Public License as published by
  12 # the Free Software Foundation; either version 3 of the License, or
  13 # (at your option) any later version.
  14 #
  15 # QWeeChat is distributed in the hope that it will be useful,
  16 # but WITHOUT ANY WARRANTY; without even the implied warranty of
  17 # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  18 # GNU General Public License for more details.
  19 #
  20 # You should have received a copy of the GNU General Public License
  21 # along with QWeeChat.  If not, see <http://www.gnu.org/licenses/>.
  22 #
  23
  24 #
  25 # Usage:  python testproto.py [-h] [-v] [-6] <hostname> <port>
  26 #
  27 # With initial commands: echo "init password=xxxx" | python testproto.py localhost 5000
  28 #                        python testproto.py localhost 5000 < commands.txt
  29 #
  30 # Return code:
  31 #   0: OK
  32 #   1: missing/invalid arguments (hostname or port)
  33 #   2: connection to WeeChat/relay failed
  34 #   3: I/O error with WeeChat/relay
  35 #
  36
  37 import os, sys, socket, select, struct, time
  38 import protocol  # WeeChat/relay protocol
  39
  40 options = { 'h': 0, 'v': 0, '6': 0 }
  41 hostname = None
  42 port = None
  43
  44 def usage():
  45     """Display usage."""
  46     print('\nSyntax: python %s [-h] [-v] [-6] <hostname> <port>\n' % sys.argv[0])
  47     print('  -h              display this help')
  48     print('  -v              verbose mode: long objects view (two -v: display raw messages)')
  49     print('  -6              connect using IPv6')
  50     print('  hostname, port  hostname (or IP address) and port of machine running WeeChat relay')
  51     print('')
  52     print('Some commands can be piped to the script, for example:')
  53     print('  echo "init password=xxxx" | python %s localhost 5000' % sys.argv[0])
  54     print('  python %s localhost 5000 < commands.txt' % sys.argv[0])
  55     print('')
  56
  57 def connect(address, ipv6):
  58     """Connect to WeeChat/relay."""
  59     inet = socket.AF_INET6 if ipv6 else socket.AF_INET
  60     sock = None
  61     try:
  62         sock = socket.socket(inet, socket.SOCK_STREAM)
  63         sock.connect(address)
  64     except:
  65         if sock:
  66             sock.close()
  67         print('Failed to connect to %s/%d using %s' % (address[0], address[1],
  68                                                        'IPv4' if inet == socket.AF_INET else 'IPv6'))
  69         return (None, None)
  70     print('Connected to %s/%d (%s)' % (hostname, port,
  71                                        'IPv4' if inet == socket.AF_INET else 'IPv6'))
  72     return (sock, inet)
  73
  74 def send(sock, messages):
  75     """Send a text message to WeeChat/relay."""
  76     has_quit = False
  77     try:
  78         for msg in messages.split('\n'):
  79             if msg == 'quit':
  80                 has_quit = True
  81             sock.sendall(msg + '\n')
  82             print('\x1b[33m<-- %s\x1b[0m' % msg)
  83     except:
  84         print('Failed to send message')
  85         return (False, has_quit)
  86     return (True, has_quit)
  87
  88 def decode(message):
  89     """Decode a binary message received from WeeChat/relay."""
  90     global options
  91     try:
  92         proto = protocol.Protocol()
  93         message = proto.decode(message, separator='\n' if options['v'] else ', ')
  94         print('')
  95         if options['v'] >= 2 and message.uncompressed:
  96             # display raw message
  97             print('\x1b[32m--> message uncompressed (%d bytes):\n%s\x1b[0m'
  98                   % (message.size_uncompressed,
  99                      protocol.hex_and_ascii(message.uncompressed, 20)))
 100         # display decoded message
 101         print('\x1b[32m--> %s\x1b[0m' % message)
 102     except:
 103         print('Error while decoding message from WeeChat')
 104         return False
 105     return True
 106
 107 def mainloop(sock):
 108     """Main loop: read keyboard, send commands, read socket and decode and display received binary messages."""
 109     message = ''
 110     recvbuf = ''
 111     prompt = '\x1b[36mrelay> \x1b[0m'
 112     sys.stdout.write(prompt)
 113     sys.stdout.flush()
 114     try:
 115         while True:
 116             inr, outr, exceptr = select.select([sys.stdin, sock], [], [], 1)
 117             for fd in inr:
 118                 if fd == sys.stdin:
 119                     buf = os.read(fd.fileno(), 4096)
 120                     if buf:
 121                         message += buf
 122                         if '\n' in message:
 123                             messages = message.split('\n')
 124                             msgsent = '\n'.join(messages[:-1])
 125                             if msgsent:
 126                                 (send_ok, has_quit) = send(sock, msgsent)
 127                                 if not send_ok:
 128                                     return 3
 129                                 if has_quit:
 130                                     return 0
 131                             message = messages[-1]
 132                             sys.stdout.write(prompt + message)
 133                             sys.stdout.flush()
 134                 else:
 135                     buf = fd.recv(4096)
 136                     if buf:
 137                         recvbuf += buf
 138                         while len(recvbuf) >= 4:
 139                             remainder = None
 140                             length = struct.unpack('>i', recvbuf[0:4])[0]
 141                             if len(recvbuf) < length:
 142                                 # partial message, just wait for end of message
 143                                 break
 144                             # more than one message?
 145                             if length < len(recvbuf):
 146                                 # save beginning of another message
 147                                 remainder = recvbuf[length:]
 148                                 recvbuf = recvbuf[0:length]
 149                             if not decode(recvbuf):
 150                                 return 3
 151                             if remainder:
 152                                 recvbuf = remainder
 153                             else:
 154                                 recvbuf = ''
 155                         sys.stdout.write(prompt + message)
 156                         sys.stdout.flush()
 157     except:
 158         send(sock, 'quit')
 159
 160 # display help if arguments are missing
 161 if len(sys.argv) < 3:
 162     usage()
 163     sys.exit(1)
 164
 165 # read command line arguments
 166 try:
 167     for arg in sys.argv[1:]:
 168         if arg[0] == '-':
 169             for opt in arg[1:]:
 170                 options[opt] = options.get(opt, 0) + 1
 171         elif hostname:
 172             port = int(arg)
 173         else:
 174             hostname = arg
 175 except:
 176     print('Invalid arguments')
 177     sys.exit(1)
 178
 179 if options['h']:
 180     usage()
 181     sys.exit(0)
 182
 183 # connect to WeeChat/relay
 184 (sock, inet) = connect((hostname, port), options['6'])
 185 if not sock:
 186     sys.exit(2)
 187
 188 # send commands from standard input if some data is available
 189 has_quit = False
 190 inr, outr, exceptr = select.select([sys.stdin], [], [], 0)
 191 if inr:
 192     data = os.read(sys.stdin.fileno(), 4096)
 193     if data:
 194         (send_ok, has_quit) = send(sock, data.strip())
 195         if not send_ok:
 196             sock.close()
 197             sys.exit(3)
 198     # open stdin to read user commands
 199     sys.stdin = open('/dev/tty')
 200
 201 # main loop (wait commands, display messages received)
 202 if not has_quit:
 203     mainloop(sock)
 204 time.sleep(0.5)
 205 sock.close()