# This code is public domain
# Anant Narayanan <anant@kix.in>

import os
import sys
import time
import socket

# Edit to your liking
EXTERNAL_HOST = "192.168.1.1"
EXTERNAL_PORT = 5990
INTERNAL_HOST = "127.0.0.1"
INTERNAL_SSH_PORT = 22
INTERNAL_HTTP_PORT = 80

TIMEOUT = 5
BLOCK_SIZE = 1
MAX_CONNECTIONS = 64
HTTP_METHODS = ['OPT', 'GET', 'HEA', 'PUT', 'DEL', 'TRA', 'CON']

class Pipe:
    def __init__(self, sock, lport, dat = None):
        pSock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        pSock.connect((INTERNAL_HOST, lport))
        pSock.setblocking(0)
        if (dat):
            pSock.send(dat)

        # begin the "piping". FIXME: UGLY!
        while (1):
            c2sdat = None
            s2cdat = None
            # client -> server
            try:
                c2sdat = sock.recv(BLOCK_SIZE)
                if (c2sdat):
                    pSock.send(c2sdat)
                else:
                    break
            except:
                pass
            # server -> client
            try:
                s2cdat = pSock.recv(BLOCK_SIZE)
                if (s2cdat):
                    sock.send(s2cdat)
                else:
                    break
            except:
                pass
        pSock.close()
        sock.close()

class Muxer:
    def __init__(self):
        mSock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        mSock.bind((EXTERNAL_HOST, EXTERNAL_PORT))
        mSock.listen(MAX_CONNECTIONS)
        
        while (1):
            conn, addr = mSock.accept()
            pid = os.fork()
            if pid:
                # Parent. Do nothing
                pass
            else:
                # Child. Set socket to nonblocking mode
                conn.setblocking(0)
                sclock = time.time()
                ishttp = False
                while (time.time() - sclock < TIMEOUT):
                    try:
                        dat = conn.recv(3)
                        if (dat in HTTP_METHODS):
                            # it's an HTTP connection
                            ishttp = True
                            break
                    except:
                        # wait some more until we know it's SSH
                        pass
                if (ishttp):
                    print 'Piping to HTTP'
                    Pipe(conn, INTERNAL_HTTP_PORT, dat)
                else:
                    print 'Piping to SSH'
                    Pipe(conn, INTERNAL_SSH_PORT)
                os._exit(1)
        
        mSock.close()

Muxer()