# A wimpy FTP Proxy
# JMT April-16-01 (Tax day)
# Don't try to read this code, its a mess

import string, re, thread, select, UserList
from transfer_pair import TransferPair, TransferFilter
from SocketServer import *
from sys import *
from socket import *

PassiveCheck = re.compile('227 .*\(([0-9]+,[0-9]+,[0-9]+,[0-9]+),([0-9]+),([0-9]+)\)')

BLOCKSIZE = 10240

def dbg_msg(msg):
    stderr.write('FTPProxy: %s\n' % msg)
    
class PassiveFilter(TransferFilter):
    def __init__(self, handler):
        self.buf     = ''
        self.handler = handler

    def filter(self, data, tosock, output, sockname):
        res = ''
        while 1:
            idx = string.find(data, '\r\n')
            if idx == -1:
                self.buf = self.buf + data
                break
            self.buf = self.buf + data[:idx+2]
            data = data[idx+2:]

            match = PassiveCheck.match(self.buf)
            if match:
                ip, highport, lowport = map(
                    lambda (beg, end), s=self.buf: s[beg:end],
                    match.regs[1:])
                ip   = string.replace(ip, ',', '.')
                port = (int(highport) << 8) + int(lowport)
                server_sock = socket(AF_INET, SOCK_STREAM)
                server_sock.connect((ip, port))
                
                sock = socket(AF_INET, SOCK_STREAM)
                sock.bind(('', 0))
                sock.listen(10)
                proxy_ip = string.replace(self.handler.conn_ip, ".", ",")
                proxy_port = sock.getsockname()[1]
                res = res + '227 Entering Passive Mode (%s,%d,%d)\r\n' % (
                    proxy_ip, proxy_port >> 8, proxy_port & 0x00ff)
                output('Data sent to %s: "%s"' % (sockname, res))
                tosock.send(res)
                res = ''
                sock, z = sock.accept()

                tpair = TransferPair(server_sock, 'server->proxy(Passive)',
                                     sock, 'proxy->client(Passive)',
                                     output)
                self.handler.add_tpair(server_sock, sock, tpair)
            else:
                res = res + self.buf

            self.buf = ''
                
        if len(res):
            output('Data sent to %s: "%s"' % (sockname, res))
            tosock.send(res)


class FTPRequestHandler(StreamRequestHandler):
    def add_tpair(self, sock1, sock2, tpair):
        self.active_sockets.extend([sock1, sock2])
        self.transfer_pairs.append(tpair)
        
    def handle(self):
        request = self.request
        self.conn_ip = request.getsockname()[0]
        dbg_msg('Client connected (%s)' % self.client_address[0])
        svr = socket(AF_INET, SOCK_STREAM)
        svr.connect((self.server.ftp_host, self.server.ftp_port))
        dbg_msg('Connected to real server')

        tpair = TransferPair(request, 'client->proxy', svr, 'proxy->server',
                             dbg_msg)
        tpair.add_filter(PassiveFilter(self).filter)
        self.transfer_pairs = [tpair]
        self.active_sockets = [request, svr]
        while len(self.active_sockets):
            rfds, z, z = select.select(self.active_sockets, [], [])
            for tpair in self.transfer_pairs:
                if tpair.do_transfer(rfds, self.active_sockets):
                    del self.transfer_pairs[self.transfer_pairs.index(tpair)]
                    
    def finish(self):
        dbg_msg('Client finished (%s)' % self.client_address[0])
        

class FTPProxyServer(ThreadingMixIn, TCPServer):
    def __init__(self, *args, **kwargs):
        self.ftp_host = kwargs['ftp_host']
        self.ftp_port = kwargs['ftp_port']
        del kwargs['ftp_host']
        del kwargs['ftp_port']
        apply(TCPServer.__init__, (self,)+args, kwargs)

    def server_bind(self):
        TCPServer.server_bind(self)
        # Make sure we can reuse this address 
        self.socket.setsockopt(SOL_SOCKET, SO_REUSEADDR, 1)
        
if __name__ == '__main__':
    if len(argv) != 4:
        sys.stderr.write('Syntax: %s <proxy_port> <real_host> <real_port>\n' %(
            argv[0],))
        sys.exit(1)
        
    x = FTPProxyServer(('', int(argv[1])), FTPRequestHandler, ftp_host=argv[2],
                       ftp_port=int(argv[3]))

    dbg_msg('Starting proxy on port %s' % (argv[1]))
    dbg_msg('Real host is %s:%s' % (argv[2], argv[3]))
    x.serve_forever()