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

# -*- coding: utf-8 -*-
#
# testproto.py - command-line program for testing WeeChat/relay protocol
#
# Copyright (C) 2013-2019 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
from .. version import qweechat_version

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:  # noqa: E722
            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:  # noqa: E722
            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.debug > 0
                                else ', ')
            print('')
            if self.args.debug >= 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:  # noqa: E722
            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:  # noqa: E722
            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('-d', '--debug', action='count', default=0,
                        help='debug mode: long objects view '
                        '(-dd: display raw messages)')
    parser.add_argument('-v', '--version', action='version',
                        version=qweechat_version())
    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	#
7f19623e	5	# Copyright (C) 2013-2019 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	39	import protocol # WeeChat/relay protocol
438f63fa	40	from .. version import qweechat_version
9b4804b9	41
77df9d06 SH	42	NAME = 'qweechat-testproto'
77df9d06 SH	43
1cf51dd2	44
77df9d06 SH	45	class TestProto(object):
77df9d06 SH	46	"""Test of WeeChat/relay protocol."""
1cf51dd2 SH	47
	48	def __init__(self, args):
	49	self.args = args
	50	self.sock = None
	51	self.has_quit = False
	52	self.address = '{self.args.hostname}/{self.args.port} ' \
	53	'(IPv{0})'.format(6 if self.args.ipv6 else 4, self=self)
	54
	55	def connect(self):
	56	"""
	57	Connect to WeeChat/relay.
	58	Return True if OK, False if error.
	59	"""
	60	inet = socket.AF_INET6 if self.args.ipv6 else socket.AF_INET
	61	try:
	62	self.sock = socket.socket(inet, socket.SOCK_STREAM)
	63	self.sock.connect((self.args.hostname, self.args.port))
524efe3c	64	except: # noqa: E722
1cf51dd2 SH	65	if self.sock:
	66	self.sock.close()
	67	print('Failed to connect to', self.address)
	68	return False
	69	print('Connected to', self.address)
	70	return True
	71
	72	def send(self, messages):
	73	"""
	74	Send a text message to WeeChat/relay.
	75	Return True if OK, False if error.
	76	"""
	77	try:
	78	for msg in messages.split('\n'):
	79	if msg == 'quit':
	80	self.has_quit = True
	81	self.sock.sendall(msg + '\n')
	82	print('\x1b[33m<-- ' + msg + '\x1b[0m')
524efe3c	83	except: # noqa: E722
1cf51dd2 SH	84	traceback.print_exc()
	85	print('Failed to send message')
	86	return False
	87	return True
	88
	89	def decode(self, message):
	90	"""
	91	Decode a binary message received from WeeChat/relay.
	92	Return True if OK, False if error.
	93	"""
	94	try:
	95	proto = protocol.Protocol()
	96	msgd = proto.decode(message,
438f63fa	97	separator='\n' if self.args.debug > 0
1cf51dd2 SH	98	else ', ')
1cf51dd2 SH	99	print('')
438f63fa	100	if self.args.debug >= 2 and msgd.uncompressed:
1cf51dd2 SH	101	# display raw message
	102	print('\x1b[32m--> message uncompressed ({0} bytes):\n'
	103	'{1}\x1b[0m'
	104	''.format(msgd.size_uncompressed,
	105	protocol.hex_and_ascii(msgd.uncompressed, 20)))
	106	# display decoded message
	107	print('\x1b[32m--> {0}\x1b[0m'.format(msgd))
524efe3c	108	except: # noqa: E722
1cf51dd2 SH	109	traceback.print_exc()
	110	print('Error while decoding message from WeeChat')
	111	return False
	112	return True
	113
	114	def send_stdin(self):
	115	"""
	116	Send commands from standard input if some data is available.
	117	Return True if OK (it's OK if stdin has no commands),
	118	False if error.
	119	"""
77df9d06	120	inr = select.select([sys.stdin], [], [], 0)[0]
1cf51dd2 SH	121	if inr:
	122	data = os.read(sys.stdin.fileno(), 4096)
	123	if data:
77df9d06	124	if not self.send(data.strip()):
403f36d9	125	# self.sock.close()
1cf51dd2 SH	126	return False
	127	# open stdin to read user commands
	128	sys.stdin = open('/dev/tty')
	129	return True
	130
	131	def mainloop(self):
	132	"""
	133	Main loop: read keyboard, send commands, read socket,
	134	decode/display binary messages received from WeeChat/relay.
	135	Return 0 if OK, 4 if send error, 5 if decode error.
	136	"""
	137	if self.has_quit:
	138	return 0
	139	message = ''
	140	recvbuf = ''
	141	prompt = '\x1b[36mrelay> \x1b[0m'
	142	sys.stdout.write(prompt)
	143	sys.stdout.flush()
	144	try:
	145	while not self.has_quit:
77df9d06 SH	146	inr = select.select([sys.stdin, self.sock], [], [], 1)[0]
	147	for _file in inr:
	148	if _file == sys.stdin:
	149	buf = os.read(_file.fileno(), 4096)
1cf51dd2 SH	150	if buf:
	151	message += buf
	152	if '\n' in message:
	153	messages = message.split('\n')
	154	msgsent = '\n'.join(messages[:-1])
	155	if msgsent and not self.send(msgsent):
	156	return 4
	157	message = messages[-1]
	158	sys.stdout.write(prompt + message)
	159	sys.stdout.flush()
	160	else:
77df9d06	161	buf = _file.recv(4096)
1cf51dd2 SH	162	if buf:
	163	recvbuf += buf
	164	while len(recvbuf) >= 4:
	165	remainder = None
	166	length = struct.unpack('>i', recvbuf[0:4])[0]
	167	if len(recvbuf) < length:
	168	# partial message, just wait for the
	169	# end of message
	170	break
	171	# more than one message?
	172	if length < len(recvbuf):
	173	# save beginning of another message
	174	remainder = recvbuf[length:]
	175	recvbuf = recvbuf[0:length]
	176	if not self.decode(recvbuf):
	177	return 5
	178	if remainder:
	179	recvbuf = remainder
	180	else:
	181	recvbuf = ''
9b4804b9 SH	182	sys.stdout.write(prompt + message)
9b4804b9 SH	183	sys.stdout.flush()
524efe3c	184	except: # noqa: E722
1cf51dd2 SH	185	traceback.print_exc()
	186	self.send('quit')
	187	return 0
	188
	189	def __del__(self):
	190	print('Closing connection with', self.address)
	191	time.sleep(0.5)
	192	self.sock.close()
	193
	194
77df9d06 SH	195	def main():
77df9d06 SH	196	"""Main function."""
1cf51dd2 SH	197	# parse command line arguments
	198	parser = argparse.ArgumentParser(
	199	formatter_class=argparse.RawDescriptionHelpFormatter,
	200	fromfile_prefix_chars='@',
77df9d06	201	description='Command-line program for testing WeeChat/relay protocol.',
1cf51dd2	202	epilog='''
77df9d06	203	Environment variable "QWEECHAT_PROTO_OPTIONS" can be set with default options.
1cf51dd2 SH	204	Argument "@file.txt" can be used to read default options in a file.
	205
	206	Some commands can be piped to the script, for example:
77df9d06 SH	207	echo "init password=xxxx" \| {name} localhost 5000
77df9d06 SH	208	{name} localhost 5000 < commands.txt
1cf51dd2 SH	209
	210	The script returns:
	211	0: OK
	212	2: wrong arguments (command line)
	213	3: connection error
	214	4: send error (message sent to WeeChat)
	215	5: decode error (message received from WeeChat)
77df9d06	216	'''.format(name=NAME))
1cf51dd2 SH	217	parser.add_argument('-6', '--ipv6', action='store_true',
1cf51dd2 SH	218	help='connect using IPv6')
438f63fa SH	219	parser.add_argument('-d', '--debug', action='count', default=0,
	220	help='debug mode: long objects view '
	221	'(-dd: display raw messages)')
	222	parser.add_argument('-v', '--version', action='version',
	223	version=qweechat_version())
1cf51dd2 SH	224	parser.add_argument('hostname',
	225	help='hostname (or IP address) of machine running '
	226	'WeeChat/relay')
	227	parser.add_argument('port', type=int,
	228	help='port of machine running WeeChat/relay')
	229	if len(sys.argv) == 1:
	230	parser.print_help()
	231	sys.exit(0)
77df9d06 SH	232	_args = parser.parse_args(
77df9d06 SH	233	shlex.split(os.getenv('QWEECHAT_PROTO_OPTIONS') or '') + sys.argv[1:])
1cf51dd2	234
77df9d06	235	test = TestProto(_args)
1cf51dd2 SH	236
	237	# connect to WeeChat/relay
	238	if not test.connect():
	239	sys.exit(3)
	240
	241	# send commands from standard input if some data is available
	242	if not test.send_stdin():
	243	sys.exit(4)
	244
	245	# main loop (wait commands, display messages received)
77df9d06	246	returncode = test.mainloop()
1cf51dd2	247	del test
77df9d06 SH	248	sys.exit(returncode)
	249
	250
	251	if __name__ == "__main__":
	252	main()