#!/usr/bin/env python

import sys
import posixpath
import socket
import select
import re
import time
import string
import signal

if hasattr(signal, 'SIGPIPE'):
	signal.signal(signal.SIGPIPE, signal.SIG_IGN)

total_bytes=0L # infinite by default
blocksize=1024*1024L # one megabyte by default
client=0
server=0
verbose=0
use_stdio=0
tcp_nodelay_selected=0
tcp_nodelay=-1
default_tcp_window_size=64*1024
tcp_window_size=-1

output_tcp=1
input_tcp=1

if not use_stdio:
	import os

def usage(retval):
	# old usage:
	# seki-strombrg> pnetcat
	# You must specify exactly one of -c and -s
	# Usage: /Dcs/seki/strombrg/bin/pnetcat [-b blocksize_in_bytes] [-t total_bytes_to_write] [-c hostname port] [-s port] [-v]
	# Wed Sep 21 16:37:08

	sys.stderr.write('Usage: %s [-b blocksize_in_bytes] [-t total_bytes_to_transfer] [-i|-I port] [-o|-O hostname port|-n] [-w window] [-v]\n' % sys.argv[0])
	sys.stderr.write('\n')
	sys.stderr.write('-i\t\tsays to use file I/O to read from stdin\n')
	sys.stderr.write('-I port\t\tsays to use sockets to get input from "port" on the local machine\n')
	sys.stderr.write('\n')
	sys.stderr.write('-o\t\tsays to use file I/O to write to stdout\n')
	sys.stderr.write('-O\t\thostname port\tsays to use sockets to connect to hostname on port\n')
	sys.stderr.write('-n\t\tsays to only read - do not write the data anywhere (should be faster than writing to /dev/null)\n')
	sys.stderr.write('-w size\t\tsays to set the TCP window size\n')
	sys.stderr.write('\n')
	sys.stderr.write('--O-udp\t\tSays to use UDP sockets, but only for the producer - IE modify -O mode\n')
	sys.stderr.write('--I-udp\t\tSays to use UDP sockets, but only for the consumer - IE modify -I mode\n')
	sys.stderr.write('-u\t\tSays to use UDP sockets, not the default, which is TCP (implies both of the above)\n')
	sys.stderr.write('\n')
	sys.stderr.write('Naturally, -i conflicts with -I, and -o conflicts with -O\n')
	sys.stderr.write('\n')
	sys.stderr.write("-w window conflicts with -u, because it's for setting the TCP window size\n")
	sys.stderr.write('\n')
	sys.stderr.write("-N [0|1] says to set TCP_NODELAY to 0 or 1 (0 enables Nagel, 1 disables) (_N_agel Algorithm)\n")
	sys.stderr.write('\n')
	sys.stderr.write('Also, -u of course requires -I or -O, --O-udp requires -O and --I-tcp requires -I\n')
	sys.stderr.write('Also note that -u may require small blocksizes compared to TCP\n')
	sys.stderr.write('For example 65508 was once the highest UDP blocksize on a\n')
	sys.stderr.write('run of the mill Fedora Core 4 system\n')
	sys.stderr.write('\n')
	sys.stderr.write('Among -i/-I and -o/-O, the lower case letter does file I/O, and the capital does socket I/O\n')

	sys.exit(retval)

def protocol_string(is_tcp):
	if is_tcp:
		return 'tcp'
	else:
		return 'udp'

def portno(service, is_tcp):
	numeric=re.compile('^[0-9][0-9]*$')
	if not numeric.match(service):
		try:
			result = socket.getservbyname(service, protocol_string(is_tcp))
		except (socket.error, msg):
			sys.stderr.write("getservbyname(%s,'%s') failed: %s\n" % (service, protocol_string(is_tcp), msg))
	else:
		result = string.atoi(service)
	return result

def getip(host):
	numeric=re.compile('^[0-9][0-9]\.*$')
	if not numeric.match(host):
		dict={}
		# test for round robin DNS or other forms of variability
		reps=10
		for i in range(reps):
			res = socket.gethostbyname(host)
			if dict.has_key(res):
				dict[res] += 1
			else:
				dict[res] = 1
		if len(dict.keys()) == 1:
			return dict.keys()[0]
		else:
			sys.stderr.write('Sorry, there is some variability in hostname lookups for %s\n' % host)
			sys.stderr.write('Please use an IP address instead\n')
			sys.exit(1)
	else:
		return host

file_in=0
file_out=0
socket_in=0
socket_out=0
null_out=0

#sys.stderr.write('You gave: %s\n' % string.join(sys.argv))

while sys.argv[1:]:
	# note the del sys.argv[1] at the bottom of the loop - so we only delete sys.argv[1] if there's one or more arguments to the option
	if sys.argv[1] == '-b' and sys.argv[2:]:
		blocksize=string.atol(sys.argv[2])
		del sys.argv[1]
	elif sys.argv[1] == '-i':
		file_in=1
	elif sys.argv[1] == '-I' and sys.argv[2:]:
		socket_in=1
		input_port_description = sys.argv[2]
		del sys.argv[1]
	elif sys.argv[1] == '-o':
		file_out=1
	elif sys.argv[1] == '-O' and sys.argv[3:]:
		socket_out=1
		hostname=getip(sys.argv[2])
		output_port_description = sys.argv[3]
		del sys.argv[1]
		del sys.argv[1]
	elif sys.argv[1] == '-n':
		null_out=1
	elif sys.argv[1] == '-N' and sys.argv[2:]:
		tcp_nodelay_selected = 1
		tcp_nodelay_flag = string.atoi(sys.argv[2])
		del sys.argv[1]
	elif sys.argv[1] == '-t' and sys.argv[2:]:
		total_bytes=string.atol(sys.argv[2])
		del sys.argv[1]
	elif sys.argv[1] == '-h':
		usage(0)
	elif sys.argv[1] == '-v':
		verbose=1
	elif sys.argv[1] == '-u':
		output_tcp = 0
		input_tcp = 0
	elif sys.argv[1] == '--O-udp':
		output_tcp = 0
	elif sys.argv[1] == '--I-udp':
		input_tcp = 0
	elif sys.argv[1] == '-w' and sys.argv[2:]:
		tcp_window_size = string.atoi(sys.argv[2])
		del sys.argv[1]
	else:
		sys.stderr.write('Illegal option: %s\n' % sys.argv[1])
		usage(1)
	del sys.argv[1]

if file_in + socket_in != 1:
	sys.stderr.write('You must specify exactly one of -i and -I (file_in is %d, socket_in is %d)\n' % (file_in, socket_in))
	usage(1)

if file_out + socket_out + null_out != 1:
	sys.stderr.write('You must specify exactly one of -o, -O and -n\n')
	usage(1)

if not input_tcp and not socket_in:
	sys.stderr.write('-u requires -I and/or -O\n--O-udp requires -I\n')
	usage(1)

if not output_tcp and not socket_out:
	sys.stderr.write('-u requires -I and/or -O\n--I-udp requires -O\n')
	usage(1)

if tcp_window_size != -1 and not input_tcp and not output_tcp:
	sys.stderr.write('-u and -w conflict\n')
	usage(1)

if tcp_window_size == -1:
	tcp_window_size = default_tcp_window_size

if socket_in:
	input_portno=portno(input_port_description, input_tcp)

if socket_out:
	output_portno=portno(output_port_description, output_tcp)

def percent(numerator,denominator):
	return "%3.5s" % (numerator*100.0/denominator)

def xfer(mode,total_to_xfer,block_len,readfn,writefn):
	if verbose:
		sys.stderr.write('%s, %s: total_to_xfer is %d, block_len is %d\n' % (sys.argv[0], mode, total_to_xfer, block_len))
	total_xferred = 0
	if total_to_xfer == 0:
		# read and write forever (or until something interrupts us :)
		# this method is a little faster than the other - less CPU
		if verbose:
			sys.stderr.write('%s, %s: Infinite loop starting\n' % (sys.argv[0], mode))
		while 1:
			block = readfn(block_len)
			if verbose:
				sys.stderr.write('%s, %s: read block (no len for speed)\n' % (sys.argv[0], mode))
			if not block:
				if verbose:
					sys.stderr.write('%s, %s: 0 length block received\n' % (sys.argv[0], mode))
				break
			done=0
			try:
				writefn(block)
			except (socket.error, OSError, IOError), (errno, detail):
				if errno == 32:
					# Broken pipe - just ignore it
					done=1
				else:
					sys.stderr.write('Error: %d %s\n' % (errno, detail))
					sys.exit(1)
			if done:
				break
			#writefn(block[0:actual_block_len])
			if verbose:
				sys.stderr.write('%s, %s: wrote block (no len for speed)\n' % (sys.argv[0], mode))
	else:
		# read and write only up to total_to_xfer bytes
		# this method is a little slower than the other - more CPU
		if verbose:
			sys.stderr.write('%s, %s: Transferring %d bytes (all told)\n' % (sys.argv[0], mode, total_to_xfer))
		while 1:
			if total_xferred + block_len > total_bytes:
				intended_block_len = total_to_xfer - total_xferred
			else:
				intended_block_len = block_len
			if intended_block_len == 0:
				if verbose:
					sys.stderr.write('%s, %s: All bytes transferred\n' % (sys.argv[0], mode))
				break
			block = readfn(intended_block_len)
			if not block:
				if verbose:
					sys.stderr.write('%s, %s: 0 length block received\n' % (sys.argv[0], mode))
				break
			actual_block_len = len(block)
			total_xferred += actual_block_len
			done=0
			try:
				writefn(block)
			except (socket.error, OSError, IOError), (errno, detail):
				if errno == 32:
					# Broken pipe - just ignore it
					done=1
				else:
					sys.stderr.write('Error: %d %s\n' % (errno, detail))
					sys.exit(1)
			#writefn(block[0:actual_block_len])
			if done:
				break
			if verbose:
				sys.stderr.write('%s, %s: Transmitted block of size %d, transferred %d of %d, %s%% complete\n' % (sys.argv[0], mode, actual_block_len, total_xferred, total_to_xfer, percent(total_xferred,total_to_xfer)))

def read_from_0(length):
	return os.read(0,length)

def write_to_1(buffer):
	offset = 0
	buffer_length = len(buffer)
	while 1:
		if offset == 0:
			length_actually_written = os.write(1,buffer)
		else:
			# dang slicing...  This is going to duplicate part of the buffer, which is slow...
			length_actually_written = os.write(1,buffer[offset:])
		offset += length_actually_written
		if offset >= buffer_length:
			return

def udpreadfn(sockin,length):
	data, remotehost = sockin.recvfrom(length)
	if verbose:
		sys.stderr.write('Received UDP packet from %s %s\n' % (remotehost))
	return data

def udpwritefn(sockout,buffer):
	sockout.sendto(buffer, (hostname, output_portno))

def noop(buffer):
	pass

def type_of_socket(is_tcp):
	if is_tcp:
		return socket.SOCK_STREAM
	else:
		return socket.SOCK_DGRAM

if file_in:
	if use_stdio:
		readfn=sys.stdin.read
	else:
		readfn=read_from_0
	mode = 'input from stdin'
elif socket_in:
	input_socket = socket.socket(socket.AF_INET,type_of_socket(input_tcp))
	input_socket.setsockopt(socket.SOL_SOCKET,socket.SO_REUSEADDR,1)
	# this would only bind to the primary hostname of the machine:
	# input_socket.bind((socket.gethostname(),port))
	#
	# '' is supposed to mean "bind to all interfaces"
	input_socket.bind(('',input_portno))
	if input_tcp:
		input_socket.setsockopt(socket.SOL_SOCKET, socket.SO_RCVBUF, tcp_window_size)
	if input_tcp:
		if tcp_nodelay_selected:
			sys.stderr.write('Setting TCP_NODELAY to %d\n' % tcp_nodelay_flag)
			input_socket.setsockopt(socket.SOL_TCP,socket.TCP_NODELAY, tcp_nodelay_flag)
		else:
			sys.stderr.write('Accepting system default for TCP_NODELAY\n')
		input_socket.listen(0)
		if verbose:
			sys.stderr.write('Waiting for connection...\n')
		input_connection, (remotehost, remoteport) = input_socket.accept()
		if verbose:
			sys.stderr.write('Received connection from %s %s\n' % (remotehost, remoteport))
		readfn=input_connection.recv
	else:
		readfn=lambda length: udpreadfn(input_socket,length)
	mode = 'input from socket (irrespective of TCP or UDP)'
else:
	sys.stderr.write('%s: Weird 1\n' % sys.argv[0])
	sys.exit(1)

if file_out:
	if use_stdio:
		writefn=sys.stdout.write
	else:
		writefn=write_to_1
	mode += ', output to stdout (irrespective of TCP or UDP)'
elif socket_out:
	output_socket = socket.socket(socket.AF_INET,type_of_socket(output_tcp))
	output_socket.connect((hostname,output_portno))
	if output_tcp:
		writefn=output_socket.send
	else:
		writefn=lambda length: udpwritefn(output_socket,length)
	if output_tcp:
		# this one's assumed to only be appropriate for TCP, though there may be other reliable types someday that want this too
		output_socket.setsockopt(socket.SOL_SOCKET, socket.SO_SNDBUF, tcp_window_size)
		if tcp_nodelay_selected:
			sys.stderr.write('Setting TCP_NODELAY to %d\n' % tcp_nodelay_flag)
			output_socket.setsockopt(socket.SOL_TCP,socket.TCP_NODELAY, tcp_nodelay_flag)
		else:
			sys.stderr.write('Accepting system default for TCP_NODELAY\n')
	mode += ', output to socket'
elif null_out:
	writefn=noop
	mode += ', no output'
else:
	sys.stderr.write('%s: Weird 2\n' % sys.argv[0])
	sys.exit(1)

if verbose:
	sys.stderr.write('%s starting up: %s\n' % (sys.argv[0], mode))

xfer(mode,total_bytes,blocksize,readfn,writefn)

if verbose:
	sys.stderr.write('Terminating normally: %s\n' % mode)

sys.exit(0)