Chatter

chatter.py

import os,sys,argparse
import socket
import struct
import traceback
import logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger()

class Message(object):
	TYPE_FILE = 1
	TYPE_STRING = 2

	def __init__(self, length, type):
		self.length = length
		self.type = type

	def build_header(self):
		header = struct.pack('i', self.type)
		header += struct.pack('Q', self.length)
		return header

	@staticmethod
	def decode_header(sock):
		type = struct.unpack('i', sock.recv(4))[0]
		if type == Message.TYPE_FILE:
			return File.decode_header(sock)
		elif type == Message.TYPE_STRING:
			return String.decode_header(sock)
		else:
			raise Exception("Unknown message type: %d" % (type))

	@staticmethod
	def build_payload(msg):
		length = len(msg)
		return struct.pack('%ds' % (length), msg)

	@staticmethod
	def decode_payload(msg):
		return struct.unpack('%ds' % (len(msg)), msg)[0]

class String(Message):
	def __init__(self, length):
		super(String, self).__init__(length, Message.TYPE_STRING)

	@staticmethod
	def decode_header(sock):
		length = struct.unpack('Q', sock.recv(8))[0]
		return String(length)

class File(Message):
	def __init__(self, path, length=0):
		if not length:
			if not os.path.exists(path):
				raise Exception("Invalid path: %s" % (path))
			length = os.stat(path).st_size
		super(File, self).__init__(length, Message.TYPE_FILE)
		self.path = path

	def build_header(self):
		'''
		Format:
		type: 4
		length: 8
		filename_len: 4
		filename: filename_len
		'''
		header = super(File, self).build_header()
		header += struct.pack('i', len(self.path))
		header += struct.pack('%ds' % (len(self.path)), self.path)
		return header

	@staticmethod
	def decode_header(sock):
		# We've already finished parsing type
		file_length = struct.unpack('Q', sock.recv(8))[0]
		filename_len = struct.unpack('i', sock.recv(4))[0]
		filename = struct.unpack('%ds' % (filename_len), sock.recv(filename_len))[0]
		msg = File(filename, file_length)
		return msg

def clear_screen():
	cmd = ''
	if os.name == 'nt':
		cmd = 'cls'
	else:
		cmd = 'clear'
	os.system(cmd)

def server(args):
	server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
	server_socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)

	try:
		server_socket.bind(('', args.port))
	except socket.error as e:
		logger.error('Bind failed! :' + e[1])
		sys.exit(-1)

	server_socket.listen(10)


	out = sys.stdout
	if args.output:
		out = open(args.output, 'wb')

	bufsize = args.bufsize
	while 1:
		sock, addr = server_socket.accept()
		print str(addr)
		message	= Message.decode_header(sock)
		if args.clear:
			clear_screen()

		read = 0
		length = message.length
		remaining = length
		while remaining > 0:
			bufsize = bufsize if bufsize < remaining else remaining
			payload = Message.decode_payload(sock.recv(bufsize))
			out.write(payload)
			if args.tee:
				print payload
			remaining -= len(payload)
			read += len(payload)
			logger.debug("Received: %d/%d (%d%%)" % (read, length, (read * 100) / length))
		out.flush()

def send_file(sock, input_file, bufsize):
	msg = File(input_file)
	header = msg.build_header()
	try:
		sock.send(header)
	except Exception:
		logger.critical("Could not send header to server!")
		return

	with open(input_file, 'r') as f:
		read = 0
		length = msg.length
		remaining = length
		while remaining > 0:
			bufsize = bufsize if bufsize < remaining else remaining
			fbytes = f.read(bufsize)
			payload = Message.build_payload(fbytes)
			assert len(payload) == len(fbytes)
			try:
				sock.send(payload)
			except Exception:
				logger.critical("Could not send payload to server!")
				return
			remaining -= len(fbytes)
			read += len(fbytes)
			logger.debug("Sent: %d/%d (%d%%)" % (read, length, (read * 100) / length))
	logger.debug("Sent file")

def send_message(sock, data, bufsize):
	if isinstance(data, list):
		data = ''.join(data)
	if data[-1] != '\n':
		data += '\n'
	message = String(len(data))
	header = message.build_header()
	payload = Message.build_payload(data)
	try:
		sock.send(header)
		sock.send(payload)
	except Exception:
		logger.critical('Could not send header to server!')
		return
	logger.debug("Sent")

def client(args):
	try:
		client_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
		client_socket.connect((args.host, args.port))
	except Exception:
		logger.critical('Could not connect to server!')
		logger.critical(traceback.format_exc())
		return

	if args.file:
		send_file(client_socket, args.file, args.bufsize)
	elif args.message:
		send_message(client_socket, args.message, args.bufsize)
	client_socket.close()

def setup_parser():
	parser = argparse.ArgumentParser()
	subparsers = parser.add_subparsers(dest='command', title='subcommands', description='server | client', help='server | client')

	server_parser = subparsers.add_parser('server')
	# Common stuff
	server_parser.add_argument('-p', '--port', action='store', type=int, default=30000, help='Port to use')
	server_parser.add_argument('-c', '--clear', action='store_true', default=False, help='Clear screen on every new message')
	server_parser.add_argument('-o', '--output', action='store', type=str,
			default=None, help='Store output in file')
	server_parser.add_argument('-tee', action='store_true', default=False,
			help='Print and pipe output to file (if specified)')
	server_parser.add_argument('-b', '--bufsize', type=int, default=1024 * 1024,
			help='Buffer size')
	server_parser.set_defaults(func=server)

	client_parser = subparsers.add_parser('client')
	# Common stuff
	client_parser.add_argument('-p', '--port', action='store', type=int, default=30000, help='Port to use')
	client_parser.add_argument('--host', action='store', type=str, default='127.0.0.1', help='Host to connect to')
	client_parser.add_argument('-b', '--bufsize', type=int, default=1024 * 1024,
			help='Buffer size')

	mutex_input = client_parser.add_mutually_exclusive_group(required=True)
	mutex_input.add_argument('-f', '--file', action='store', type=str, help='File data to send')
	mutex_input.add_argument('-m', '--message', action='store', type=str, help='Data to send')
	client_parser.set_defaults(func=client)

	return parser

def main(argv):
	parser = setup_parser()
	args   = parser.parse_args(argv[1:])

	args.func(args)

if __name__ == '__main__':
	main(sys.argv)