dotfiles

My personal shell configs and stuff
git clone git://git.alex.balgavy.eu/dotfiles.git
Log | Files | Refs | Submodules | README | LICENSE

tcprelay.py (4192B)

      1 #!/usr/bin/python
      2 # -*- coding: utf-8 -*-
      3 #
      4 #	tcprelay.py - TCP connection relay for usbmuxd
      5 #
      6 # Copyright (C) 2009	Hector Martin "marcan" <hector@marcansoft.com>
      7 #
      8 # This program is free software; you can redistribute it and/or modify
      9 # it under the terms of the GNU General Public License as published by
     10 # the Free Software Foundation, either version 2 or version 3.
     11 #
     12 # This program is distributed in the hope that it will be useful,
     13 # but WITHOUT ANY WARRANTY; without even the implied warranty of
     14 # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
     15 # GNU General Public License for more details.
     16 #
     17 # You should have received a copy of the GNU General Public License
     18 # along with this program; if not, write to the Free Software
     19 # Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA
     20 
     21 import usbmux
     22 import SocketServer
     23 import select
     24 from optparse import OptionParser
     25 import sys
     26 import threading
     27 
     28 class SocketRelay(object):
     29 	def __init__(self, a, b, maxbuf=65535):
     30 		self.a = a
     31 		self.b = b
     32 		self.atob = ""
     33 		self.btoa = ""
     34 		self.maxbuf = maxbuf
     35 	def handle(self):
     36 		while True:
     37 			rlist = []
     38 			wlist = []
     39 			xlist = [self.a, self.b]
     40 			if self.atob:
     41 				wlist.append(self.b)
     42 			if self.btoa:
     43 				wlist.append(self.a)
     44 			if len(self.atob) < self.maxbuf:
     45 				rlist.append(self.a)
     46 			if len(self.btoa) < self.maxbuf:
     47 				rlist.append(self.b)
     48 			rlo, wlo, xlo = select.select(rlist, wlist, xlist)
     49 			if xlo:
     50 				return
     51 			if self.a in wlo:
     52 				n = self.a.send(self.btoa)
     53 				self.btoa = self.btoa[n:]
     54 			if self.b in wlo:
     55 				n = self.b.send(self.atob)
     56 				self.atob = self.atob[n:]
     57 			if self.a in rlo:
     58 				s = self.a.recv(self.maxbuf - len(self.atob))
     59 				if not s:
     60 					return
     61 				self.atob += s
     62 			if self.b in rlo:
     63 				s = self.b.recv(self.maxbuf - len(self.btoa))
     64 				if not s:
     65 					return
     66 				self.btoa += s
     67 			#print "Relay iter: %8d atob, %8d btoa, lists: %r %r %r"%(len(self.atob), len(self.btoa), rlo, wlo, xlo)
     68 
     69 class TCPRelay(SocketServer.BaseRequestHandler):
     70 	def handle(self):
     71 		print "Incoming connection to %d"%self.server.server_address[1]
     72 		mux = usbmux.USBMux(options.sockpath)
     73 		print "Waiting for devices..."
     74 		if not mux.devices:
     75 			mux.process(1.0)
     76 		if not mux.devices:
     77 			print "No device found"
     78 			self.request.close()
     79 			return
     80 		dev = mux.devices[0]
     81 		print "Connecting to device %s"%str(dev)
     82 		dsock = mux.connect(dev, self.server.rport)
     83 		lsock = self.request
     84 		print "Connection established, relaying data"
     85 		try:
     86 			fwd = SocketRelay(dsock, lsock, self.server.bufsize * 1024)
     87 			fwd.handle()
     88 		finally:
     89 			dsock.close()
     90 			lsock.close()
     91 		print "Connection closed"
     92 
     93 class TCPServer(SocketServer.TCPServer):
     94 	allow_reuse_address = True
     95 
     96 class ThreadedTCPServer(SocketServer.ThreadingMixIn, TCPServer):
     97 	pass
     98 
     99 HOST = "localhost"
    100 
    101 parser = OptionParser(usage="usage: %prog [OPTIONS] RemotePort[:LocalPort] [RemotePort[:LocalPort]]...")
    102 parser.add_option("-t", "--threaded", dest='threaded', action='store_true', default=False, help="use threading to handle multiple connections at once")
    103 parser.add_option("-b", "--bufsize", dest='bufsize', action='store', metavar='KILOBYTES', type='int', default=128, help="specify buffer size for socket forwarding")
    104 parser.add_option("-s", "--socket", dest='sockpath', action='store', metavar='PATH', type='str', default=None, help="specify the path of the usbmuxd socket")
    105 
    106 options, args = parser.parse_args()
    107 
    108 serverclass = TCPServer
    109 if options.threaded:
    110 	serverclass = ThreadedTCPServer
    111 
    112 if len(args) == 0:
    113 	parser.print_help()
    114 	sys.exit(1)
    115 
    116 ports = []
    117 
    118 for arg in args:
    119 	try:
    120 		if ':' in arg:
    121 			rport, lport = arg.split(":")
    122 			rport = int(rport)
    123 			lport = int(lport)
    124 			ports.append((rport, lport))
    125 		else:
    126 			ports.append((int(arg), int(arg)))
    127 	except:
    128 		parser.print_help()
    129 		sys.exit(1)
    130 
    131 servers=[]
    132 
    133 for rport, lport in ports:
    134 	print "Forwarding local port %d to remote port %d"%(lport, rport)
    135 	server = serverclass((HOST, lport), TCPRelay)
    136 	server.rport = rport
    137 	server.bufsize = options.bufsize
    138 	servers.append(server)
    139 
    140 alive = True
    141 
    142 while alive:
    143 	try:
    144 		rl, wl, xl = select.select(servers, [], [])
    145 		for server in rl:
    146 			server.handle_request()
    147 	except:
    148 		alive = False