Add python alternative to 'nc' for SSH connections

#!/bin/sh

die(){ echo "$*" >&2; exit 1; }

has(){ type $1 >/dev/null 2>/dev/null; }

has nc || die "nc is needed to connect to $1:$2"

nc $1 $2 || exec ssh -W $1:$2 $3

if has python; then

	python `dirname $0`/forward.py "$@" && exit

elif has nc; then

	nc -w 10 $1 $2 && exit

fi

exec ssh -W $1:$2 $3

import io

import os

import sys

import socket

import select

class Stream:

	def __init__(self, read_handle, write_handle):

		self.read_handle = read_handle

		self.write_handle = write_handle

		self.dest_stream = None

		self.read_buffer = None

		self.write_buffer = None

		self.closed = False

	@classmethod

	def open_socket(cls, host, port):

		sock = None

		addrinfo = socket.getaddrinfo(host, port, 0, socket.SOCK_STREAM)

		for af, stype, proto, name, addr in addrinfo:

			sock = socket.socket(af, stype, proto)

			sock.settimeout(2)

			try:

				sock.connect(addr)

			except (OSError, socket.timeout):

				sock.close()

				sock = None

			else:

				break

		if sock is None:

			raise OSError("Unable to connect to %s:%s" % (host, port))

		sock.settimeout(0)

		return cls(sock, sock)

	@classmethod

	def open_stdpipe(cls):

		return cls(io.open(0, 'rb', 0), io.open(1, 'wb', 0))

	def read(self, maxsize):

		try:

			return self.read_handle.recv(maxsize)

		except AttributeError:

			return self.read_handle.read(maxsize)

	def write(self, data):

		try:

			return self.write_handle.send(data)

		except AttributeError:

			handle = self.write_handle

			written = os.write(handle.fileno(), data)

			handle.flush()

			return written

	def connect_stream(self, to_stream):

		self.dest_stream = to_stream

	def transfer(self, data, wlist):

		if self.write_buffer:

			return False

		self.write_buffer = data

		if self.write_handle not in wlist:

			wlist.append(self.write_handle)

		return True

	def action(self, rlist, wlist):

		rhandle, whandle = self.read_handle, self.write_handle

		if whandle in wlist:

			assert self.write_buffer

			wlist.remove(whandle)

			written = self.write(self.write_buffer)

			self.write_buffer = self.write_buffer[written:]

		if rhandle in rlist:

			assert not self.read_buffer

			rlist.remove(rhandle)

			self.read_buffer = data = self.read(1024)

			if not data:

				self.closed = True

				self.dest_stream.closed = True

		if self.closed:

			return

		if self.read_buffer and self.dest_stream.transfer(self.read_buffer, wlist):

			self.read_buffer = None

		if not self.read_buffer:

			rlist.append(rhandle)

		if self.write_buffer:

			wlist.append(whandle)

def main():

	sock_stream = Stream.open_socket(*sys.argv[1:3])

	pipe_stream = Stream.open_stdpipe()

	sock_stream.connect_stream(pipe_stream)

	pipe_stream.connect_stream(sock_stream)

	rlist = []

	wlist = []

	xlist = []

	while 1:

		sock_stream.action(rlist, wlist)

		pipe_stream.action(rlist, wlist)

		if not rlist and not wlist:

			break

		rlist, wlist, _ = select.select(rlist, wlist, xlist, None)

if __name__ == '__main__':

	try:

		main()

	except Exception as exc:

		sys.stderr.write("Connection failed: %s\n" % str(exc))

		sys.exit(1)