[irc/weechat/qweechat.git] / qweechat / weechat / testproto.py

# -*- coding: utf-8 -*-
#
# testproto.py - command-line program for testing WeeChat/relay protocol
#
# Copyright (C) 2013-2014 Sébastien Helleu <flashcode@flashtux.org>
#
# This file is part of QWeeChat, a Qt remote GUI for WeeChat.
#
# QWeeChat 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.
#
# QWeeChat 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 QWeeChat.  If not, see <http://www.gnu.org/licenses/>.
#

"""
Command-line program for testing WeeChat/relay protocol.
"""

from __future__ import print_function

import argparse
import os
import select
import shlex
import socket
import struct
import sys
import time
import traceback

import protocol  # WeeChat/relay protocol

NAME = 'qweechat-testproto'


class TestProto(object):
    """Test of WeeChat/relay protocol."""

    def __init__(self, args):
        self.args = args
        self.sock = None
        self.has_quit = False
        self.address = '{self.args.hostname}/{self.args.port} ' \
            '(IPv{0})'.format(6 if self.args.ipv6 else 4, self=self)

    def connect(self):
        """
        Connect to WeeChat/relay.
        Return True if OK, False if error.
        """
        inet = socket.AF_INET6 if self.args.ipv6 else socket.AF_INET
        try:
            self.sock = socket.socket(inet, socket.SOCK_STREAM)
            self.sock.connect((self.args.hostname, self.args.port))
        except:
            if self.sock:
                self.sock.close()
            print('Failed to connect to', self.address)
            return False
        print('Connected to', self.address)
        return True

    def send(self, messages):
        """
        Send a text message to WeeChat/relay.
        Return True if OK, False if error.
        """
        try:
            for msg in messages.split('\n'):
                if msg == 'quit':
                    self.has_quit = True
                self.sock.sendall(msg + '\n')
                print('\x1b[33m<-- ' + msg + '\x1b[0m')
        except:
            traceback.print_exc()
            print('Failed to send message')
            return False
        return True

    def decode(self, message):
        """
        Decode a binary message received from WeeChat/relay.
        Return True if OK, False if error.
        """
        try:
            proto = protocol.Protocol()
            msgd = proto.decode(message,
                                separator='\n' if self.args.verbose > 0
                                else ', ')
            print('')
            if self.args.verbose >= 2 and msgd.uncompressed:
                # display raw message
                print('\x1b[32m--> message uncompressed ({0} bytes):\n'
                      '{1}\x1b[0m'
                      ''.format(msgd.size_uncompressed,
                                protocol.hex_and_ascii(msgd.uncompressed, 20)))
            # display decoded message
            print('\x1b[32m--> {0}\x1b[0m'.format(msgd))
        except:
            traceback.print_exc()
            print('Error while decoding message from WeeChat')
            return False
        return True

    def send_stdin(self):
        """
        Send commands from standard input if some data is available.
        Return True if OK (it's OK if stdin has no commands),
        False if error.
        """
        inr = select.select([sys.stdin], [], [], 0)[0]
        if inr:
            data = os.read(sys.stdin.fileno(), 4096)
            if data:
                if not self.send(data.strip()):
                    #self.sock.close()
                    return False
            # open stdin to read user commands
            sys.stdin = open('/dev/tty')
        return True

    def mainloop(self):
        """
        Main loop: read keyboard, send commands, read socket,
        decode/display binary messages received from WeeChat/relay.
        Return 0 if OK, 4 if send error, 5 if decode error.
        """
        if self.has_quit:
            return 0
        message = ''
        recvbuf = ''
        prompt = '\x1b[36mrelay> \x1b[0m'
        sys.stdout.write(prompt)
        sys.stdout.flush()
        try:
            while not self.has_quit:
                inr = select.select([sys.stdin, self.sock], [], [], 1)[0]
                for _file in inr:
                    if _file == sys.stdin:
                        buf = os.read(_file.fileno(), 4096)
                        if buf:
                            message += buf
                            if '\n' in message:
                                messages = message.split('\n')
                                msgsent = '\n'.join(messages[:-1])
                                if msgsent and not self.send(msgsent):
                                    return 4
                                message = messages[-1]
                                sys.stdout.write(prompt + message)
                                sys.stdout.flush()
                    else:
                        buf = _file.recv(4096)
                        if buf:
                            recvbuf += buf
                            while len(recvbuf) >= 4:
                                remainder = None
                                length = struct.unpack('>i', recvbuf[0:4])[0]
                                if len(recvbuf) < length:
                                    # partial message, just wait for the
                                    # end of message
                                    break
                                # more than one message?
                                if length < len(recvbuf):
                                    # save beginning of another message
                                    remainder = recvbuf[length:]
                                    recvbuf = recvbuf[0:length]
                                if not self.decode(recvbuf):
                                    return 5
                                if remainder:
                                    recvbuf = remainder
                                else:
                                    recvbuf = ''
                            sys.stdout.write(prompt + message)
                            sys.stdout.flush()
        except:
            traceback.print_exc()
            self.send('quit')
        return 0

    def __del__(self):
        print('Closing connection with', self.address)
        time.sleep(0.5)
        self.sock.close()


def main():
    """Main function."""
    # parse command line arguments
    parser = argparse.ArgumentParser(
        formatter_class=argparse.RawDescriptionHelpFormatter,
        fromfile_prefix_chars='@',
        description='Command-line program for testing WeeChat/relay protocol.',
        epilog='''
Environment variable "QWEECHAT_PROTO_OPTIONS" can be set with default options.
Argument "@file.txt" can be used to read default options in a file.

Some commands can be piped to the script, for example:
  echo "init password=xxxx" | {name} localhost 5000
  {name} localhost 5000 < commands.txt

The script returns:
  0: OK
  2: wrong arguments (command line)
  3: connection error
  4: send error (message sent to WeeChat)
  5: decode error (message received from WeeChat)
'''.format(name=NAME))
    parser.add_argument('-6', '--ipv6', action='store_true',
                        help='connect using IPv6')
    parser.add_argument('-v', '--verbose', action='count', default=0,
                        help='verbose mode: long objects view '
                        '(-vv: display raw messages)')
    parser.add_argument('hostname',
                        help='hostname (or IP address) of machine running '
                        'WeeChat/relay')
    parser.add_argument('port', type=int,
                        help='port of machine running WeeChat/relay')
    if len(sys.argv) == 1:
        parser.print_help()
        sys.exit(0)
    _args = parser.parse_args(
        shlex.split(os.getenv('QWEECHAT_PROTO_OPTIONS') or '') + sys.argv[1:])

    test = TestProto(_args)

    # connect to WeeChat/relay
    if not test.connect():
        sys.exit(3)

    # send commands from standard input if some data is available
    if not test.send_stdin():
        sys.exit(4)

    # main loop (wait commands, display messages received)
    returncode = test.mainloop()
    del test
    sys.exit(returncode)


if __name__ == "__main__":
    main()
Commit	Line	Data
9b4804b9 SH	1	# -- coding: utf-8 --
9b4804b9 SH	2	#
77df9d06	3	# testproto.py - command-line program for testing WeeChat/relay protocol
9b4804b9	4	#
da74afdb	5	# Copyright (C) 2013-2014 Sébastien Helleu <flashcode@flashtux.org>
9b4804b9 SH	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
77df9d06 SH	23	"""
	24	Command-line program for testing WeeChat/relay protocol.
	25	"""
	26
1cf51dd2 SH	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
9b4804b9	38
9b4804b9 SH	39	import protocol # WeeChat/relay protocol
9b4804b9 SH	40
77df9d06 SH	41	NAME = 'qweechat-testproto'
77df9d06 SH	42
1cf51dd2	43
77df9d06 SH	44	class TestProto(object):
77df9d06 SH	45	"""Test of WeeChat/relay protocol."""
1cf51dd2 SH	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	"""
77df9d06	119	inr = select.select([sys.stdin], [], [], 0)[0]
1cf51dd2 SH	120	if inr:
	121	data = os.read(sys.stdin.fileno(), 4096)
	122	if data:
77df9d06	123	if not self.send(data.strip()):
1cf51dd2 SH	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:
77df9d06 SH	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)
1cf51dd2 SH	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:
77df9d06	160	buf = _file.recv(4096)
1cf51dd2 SH	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 = ''
9b4804b9 SH	181	sys.stdout.write(prompt + message)
9b4804b9 SH	182	sys.stdout.flush()
1cf51dd2 SH	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
77df9d06 SH	194	def main():
77df9d06 SH	195	"""Main function."""
1cf51dd2 SH	196	# parse command line arguments
	197	parser = argparse.ArgumentParser(
	198	formatter_class=argparse.RawDescriptionHelpFormatter,
	199	fromfile_prefix_chars='@',
77df9d06	200	description='Command-line program for testing WeeChat/relay protocol.',
1cf51dd2	201	epilog='''
77df9d06	202	Environment variable "QWEECHAT_PROTO_OPTIONS" can be set with default options.
1cf51dd2 SH	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:
77df9d06 SH	206	echo "init password=xxxx" \| {name} localhost 5000
77df9d06 SH	207	{name} localhost 5000 < commands.txt
1cf51dd2 SH	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)
77df9d06	215	'''.format(name=NAME))
1cf51dd2 SH	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)
77df9d06 SH	229	_args = parser.parse_args(
77df9d06 SH	230	shlex.split(os.getenv('QWEECHAT_PROTO_OPTIONS') or '') + sys.argv[1:])
1cf51dd2	231
77df9d06	232	test = TestProto(_args)
1cf51dd2 SH	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)
77df9d06	243	returncode = test.mainloop()
1cf51dd2	244	del test
77df9d06 SH	245	sys.exit(returncode)
	246
	247
	248	if __name__ == "__main__":
	249	main()