tcprelay.py

#!/usr/bin/python
# -*- coding: utf-8 -*-
#
# tcprelay.py - TCP connection relay for usbmuxd
#
# Copyright (C) 2009	Hector Martin "marcan" <hector@marcansoft.com>
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 2 or version 3.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program; if not, write to the Free Software
# Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA
#
# @2020 Modified by ptcong to support multiple devices and multiple ports per device

import usbmux
import SocketServer
import select
from optparse import OptionParser
import sys
import traceback
import time


class SocketRelay(object):

    def __init__(self, a, b, maxBuffer=65535):
        self.a = a
        self.b = b
        self.atob = ""
        self.btoa = ""
        self.maxBuffer = maxBuffer

    def handle(self):
        while True:
            rlist = []
            wlist = []
            xlist = [self.a, self.b]
            if self.atob:
                wlist.append(self.b)
            if self.btoa:
                wlist.append(self.a)
            if len(self.atob) < self.maxBuffer:
                rlist.append(self.a)
            if len(self.btoa) < self.maxBuffer:
                rlist.append(self.b)
            rlo, wlo, xlo = select.select(rlist, wlist, xlist)
            if xlo:
                return
            if self.a in wlo:
                n = self.a.send(self.btoa)
                self.btoa = self.btoa[n:]
            if self.b in wlo:
                n = self.b.send(self.atob)
                self.atob = self.atob[n:]
            if self.a in rlo:
                s = self.a.recv(self.maxBuffer - len(self.atob))
                if not s:
                    return
                self.atob += s
            if self.b in rlo:
                s = self.b.recv(self.maxBuffer - len(self.btoa))
                if not s:
                    return
                self.btoa += s
            # print "Relay iter: %8d atob, %8d btoa, lists: %r %r %r"%(len(self.atob), len(self.btoa), rlo, wlo, xlo)


class TCPRelay(SocketServer.BaseRequestHandler):

    def handle(self):
        def show_msg(msg):
            print "{0} ### {1}".format(self.server.server_address[1], msg)

        show_msg("Incoming connection to local port {0}:{1}".format(self.server.server_address[0], self.server.server_address[1]))
        show_msg("Waiting for device {0}:{1}".format(self.server.deviceUDID, self.server.remotePort))

        device = None
        if self.server.deviceUDID is None:
            device = mux.devices[0]
        else:
            for dev in mux.devices:
                if dev.serial == self.server.deviceUDID:
                    device = dev
                    break

        if not device:
            self.request.close()
            mux.process(0.1)
            return

        try:
            show_msg("Connecting to device {0}:{1}".format(device.serial, self.server.remotePort))

            dsock = mux.connect(device, self.server.remotePort)
            lsock = self.request

            show_msg("Connection established, relaying data")

            try:
                fwd = SocketRelay(dsock, lsock, self.server.bufferSize * 1024)
                fwd.handle()
            finally:
                dsock.close()
                lsock.close()

            show_msg("Connection closed")

        except KeyboardInterrupt:
            quit()

        except:
            pass
        mux.process(0.1)


class TCPServer(SocketServer.TCPServer):
    allow_reuse_address = True


class ThreadedTCPServer(SocketServer.ThreadingMixIn, TCPServer):
    pass

parser = OptionParser(usage="usage: %prog [OPTIONS] [Lhost:]Lport::deviceUDID:Rport [[Lhost:]Lport::deviceUDID:Rport] ...")
parser.add_option("-t", "--threaded", dest='threaded', action='store_true', default=True,
                  help="use threading to handle multiple connections at once")
parser.add_option("-b", "--bufsize", dest='bufsize', action='store', metavar='KILOBYTES', type='int', default=128,
                  help="specify buffer size for socket forwarding")
parser.add_option("-s", "--socket", dest='sockpath', action='store', metavar='PATH', type='str', default=None,
                  help="specify the path of the usbmuxd socket")

options, args = parser.parse_args()

serverClass = ThreadedTCPServer if options.threaded else TCPServer

if len(args) == 0:
    parser.print_help()
    sys.exit(1)


ports = []
for arg in args:
    try:
        localPort, remotePort = arg.split("::")
        localHost, localPort = localPort.split(":") if len(localPort.split(":")) > 1 else ("localhost", localPort)
        deviceUDID, remotePort = remotePort.split(":") if len(remotePort.split(":")) > 1 else (None, remotePort)

        ports.append((localHost, int(localPort), deviceUDID, int(remotePort)))

    except:
        parser.print_help()
        sys.exit(1)

servers = []

for localHost, localPort, deviceUDID, remotePort in ports:
    print "Forwarding {0}:{1} ==> {2}:{3}".format(localHost, localPort, deviceUDID, remotePort)

    server = serverClass((localHost, localPort), TCPRelay)
    server.deviceUDID = deviceUDID
    server.remotePort = remotePort
    server.bufferSize = options.bufsize
    servers.append(server)

mux = usbmux.USBMux(options.sockpath)
for i in range(1, len(ports)):
    mux.process(0.1)

alive = True
while alive:
    try:
        rl, wl, xl = select.select(servers, [], [])
        for server in rl:
            server.handle_request()
    except KeyboardInterrupt:
        quit()
    except:
        traceback.print_exc()
        alive = False