#! /usr/bin/env python
"""
For crappy high-packets-dropping NAT'ed internet connexion.
Provides a fast-recovery TCP-like pipe.

Usage:

    on a remote server with non-firewalled incoming TCP connexions:
        crappyline.py  --server  LISTENPORT  REALSERVER:PORT

    on your local machine:
        crappyline.py  --client  LISTENPORT  CRAPPYLINESERVER:PORT

    where LISTENPORT can be 0 for stdin/stdout XXX IN-PROGRESS

Example:
        
codespeak$ crappyline.py --server 9022 localhost:22
mylaptop$ ssh -o "ProxyCommand crappyline.py --client 0 %h:9022" codespeak.net
"""

from socket import *
import greensock
import pipelayer
from struct import pack, unpack
from time import time as now


NUM_PARALLEL_TCP_LINES = 7
RETRY_TIMEOUT          = 3             # seconds
PACKET_SIZE            = 1500 - 32
MAGIC                  = 1049700000

class Runner:

    def __init__(self, forward):
        self.forward = forward
        self.outputchannel = greensock.Channel()
        self.notify_queue = greensock.Channel()
        self.pipe = pipelayer.PipeLayer()
        self.pipe.timeout = RETRY_TIMEOUT
        self.nextencoding = None

    def mainloop(self):
        greensock.autogreenlet(self.forwardreader)
        greensock.mainloop()

    def add_transport(self, conn):
        greensock.autogreenlet(self.locallinereader, conn)
        greensock.autogreenlet(self.locallinewriter, conn)

    def locallinereader(self, conn):
        buffer = ''
        while 1:
            data = greensock.recv(conn, 32768)
            if not data:
                break    # EOF
            progress = 0
            buffer += data
            while len(buffer) > 4:
                length, = unpack("!i", buffer[:4])
                length -= MAGIC
                assert 0 < length <= PACKET_SIZE
                if len(buffer) < 4 + length:
                    break
                if not progress:
                    self.pipe.settime(now())
                    progress = 1
                data = self.pipe.decode(buffer[4:4+length])
                #print "QUEUED:", self.pipe.queue_size(), self.pipe.settime(now())
                buffer = buffer[4+length:]
                if data:
                    greensock.sendall(self.forward, data)
            if progress:
                greensock.autogreenlet(self.encoder)

    def locallinewriter(self, conn):
        while 1:
            greensock.wait_output(conn)
            buffer = self.outputchannel.receive()
            buffer = pack("!i", MAGIC + len(buffer)) + buffer
            greensock.sendall(conn, buffer)
            del buffer

    def forwardreader(self):
        while 1:
            while self.pipe.queue_size() > 8192:
                self.notify_queue.receive()
            buffer = greensock.recv(self.forward, 8192)
            if not buffer:
                break     # EOF
            self.pipe.queue(buffer)
            #print "QUEUED:", self.pipe.queue_size(), self.pipe.settime(now())
            del buffer
            greensock.autogreenlet(self.encoder)

    def encoder(self):
        while 1:
            self.outputchannel.wait_sendable()
            t = now()
            delay = self.pipe.settime(t)
            if delay is None:
                return
            if delay == 0.0:
                data = self.pipe.encode(PACKET_SIZE)
                #print "QUEUED:", self.pipe.queue_size(), self.pipe.settime(now())
                self.notify_queue.send(None, block=False)
                if not data:
                    return
                self.outputchannel.send(data)
            else:
                if self.nextencoding is not None:
                    otherdelay = self.nextencoding - t
                    if otherdelay < 1.5 * delay:
                        return
                self.nextencoding = t + delay
                greensock.sleep(delay * 1.1)
                self.nextencoding = None


def run_server(port, forward_addr):
    forward = socket(AF_INET, SOCK_STREAM)
    forward.connect(forward_addr)
    print 'connected to %r.' % (forward.getpeername(),)

    srv = socket(AF_INET, SOCK_STREAM)
    srv.bind(('', port))
    srv.listen(5)
    print 'listening on %r.' % (srv.getsockname(),)

    runner = Runner(forward)

    def accepter():
        while 1:
            greensock.wait_input(srv)
            conn, addr = srv.accept()
            print 'accepted raw TCP connexion from', addr
            runner.add_transport(conn)
            del conn

    greensock.autogreenlet(accepter)

    runner.mainloop()


def run_client(local_port, server_addr):
    local1 = socket(AF_INET, SOCK_STREAM)
    local1.bind(('', local_port))
    local1.listen(1)
    print 'waiting for local connexion on %r...' % (local1.getsockname(),)
    local, addr = local1.accept()
    print 'accepted connexion from %r.' % (addr,)

    runner = Runner(local)

    print 'connecting to server:',
    sys.stdout.flush()
    for i in range(NUM_PARALLEL_TCP_LINES):
        print i,
        sys.stdout.flush()
        conn = socket(AF_INET, SOCK_STREAM)
        conn.connect(server_addr)
        runner.add_transport(conn)
    print 'done.'

    runner.mainloop()


if __name__ == '__main__':
    import sys
    if len(sys.argv) != 4:
        print >> sys.stderr, __doc__
        sys.exit(2)
    mode, listenport, forwardaddr = sys.argv[1:]
    listenport = int(listenport)
    forwardhost, forwardport = forwardaddr.split(':')
    forwardport = int(forwardport)
    forwardaddr = (forwardhost, forwardport)
    if mode == '--server':
        run_server(listenport, forwardaddr)
    elif mode == '--client':
        run_client(listenport, forwardaddr)
    else:
        print >> sys.stderr, __doc__
        sys.exit(2)