xfer-test

#!/usr/bin/python2
import os
import re
from socket import *
import subprocess
import string
import sys

bufsize = 64*1024
rx_pasv = re.compile(r'(\d+),(\d+),(\d+),(\d+),(\d+),(\d+)')
rx_pid = re.compile(r'\[\d+\]')

def pasv_socket(pasv):
	global sock
	match = rx_pasv.search(pasv)
	sa = map(int, match.groups())
	sock = socket(AF_INET, SOCK_STREAM)
	sock.connect((string.join(map(str, sa[:4]), '.'), (sa[4] << 8 | sa[5])))
	return pasv[:match.start()] + '#,#,#,#,#,#' + pasv[match.end():]

def port_socket():
	global sock, ip
	sock = socket(AF_INET, SOCK_STREAM)
	sock.bind((ip, 0))
	addr = sock.getsockname()
	sock.listen(1)
	return "PORT %s,%d,%d\r\n" % (string.replace(addr[0], '.', ','),
								  addr[1] >> 8, addr[1] & 0xff)

def read_stdin(bufsize):
	return os.read(0, bufsize)

def recv_socket(bufsize):
	global sock
	return sock.recv(bufsize)

def write_stdout(buf):
	return os.write(1, buf)

def send_socket(buf):
	global sock
	sock.send(buf)

try:
	program = sys.argv[1]
	pasv = { 'pasv': 1, 'port': 0 }[sys.argv[2]]
	ifn,ofn = {
		'retr': (recv_socket, write_stdout),
		'store': (read_stdin, send_socket),
		}[sys.argv[3]]
	fmtcmd = { 'asc': 'TYPE A', 'bin': 'TYPE I' }[sys.argv[4]]
	try:
		arg = sys.argv[5].split(':',1)
		maxbytes = int(arg[0])
		if len(arg) > 1:
			intrmode = arg[1]
		else:
			intrmode = 'close'
		sys.argv = sys.argv[:5] + sys.argv[6:]
	except ValueError:
		maxbytes = 2**32
		intrmode = 'close'
	xcommands = sys.argv[5:-1]
	command = sys.argv[-1]
except:
	print "usage: xfer-test program [pasv|port] [retr|store] [asc|bin] [maxbytes[:abort|close]] [setup commands] command"
	sys.exit(1)

ip = os.environ.setdefault('TCPREMOTEIP', '127.0.0.1')
os.environ.setdefault('TCPLOCALIP', ip)
os.environ.setdefault('UID', str(os.getuid()))
os.environ.setdefault('GID', str(os.getgid()))

devnull = os.open('/dev/null', os.O_RDWR)
pipe = subprocess.Popen(program,
						stdin=subprocess.PIPE,
						stdout=subprocess.PIPE,
						stderr=devnull)
def readline():
	global pipe
	line = pipe.stdout.readline()
	line = re.sub(', \d+[^\d]?B/s', '', line)
	num = int(line[:3])
	if num >= 400:
		os.write(1, line)
		sys.exit(1)
	return line

os.write(1, readline())

xcommands[0:0] = [fmtcmd]
for cmd in xcommands:
	pipe.stdin.write(cmd)
	pipe.stdin.write("\r\n")
	pipe.stdin.flush()
	os.write(1, readline())

if pasv:
	pipe.stdin.write("PASV\r\n")
	pipe.stdin.flush()
	line = pasv_socket(readline())
	os.write(1, line)
else:
	pipe.stdin.write(port_socket())
	pipe.stdin.flush()
	os.write(1, readline())

pipe.stdin.write(command)
pipe.stdin.write('\r\n')
pipe.stdin.flush()
os.write(1, readline())

if not pasv:
	(fd,addr) = sock.accept()
	sock.close()
	sock = fd

while 1:
	if bufsize > maxbytes:
		bufsize = maxbytes
	if bufsize <= 0:
		break
	buf = ifn(bufsize)
	if not buf:
		break
	ofn(buf)
	maxbytes -= len(buf)
if intrmode == 'abort':
	pipe.stdin.write("ABOR\r\n")
	pipe.stdin.flush()
elif intrmode == 'drop':
	pipe.stdin.close()
else:
	sock.close()

os.write(1, readline())