## SnorkyCore -- This is the main library of Snorky,
## aimed to provide the forward over SSH of arbitrary
## ports.
 
import select, SocketServer, threading, paramiko
 
# Set verbosity (only useful for debugging)
g_verbose = False
 
def verbose(string):
    if(g_verbose):
        print "=> " + string
 
class SnorkyHandler(SocketServer.BaseRequestHandler):
    """The handler for the data exchanged by Snorky and the
    server on the SSH server contacted by Snorky"""
 
    def handle(self):
        """This the overloaded method the actually handles the data. It is
        called when the connection is open (and actually it opens the connection
        to the ssh server!)."""
 
        # We first need to try to obtain the forwarded channel
        try:
            self.channel = self.transport.open_channel('direct-tcpip',
                                                       (self.remote_host, self.remote_port),
                                                       (self.local_host , self.local_port))
            verbose("Channel opened")
        except Exception, e:
            verbose("Unable to open the requested channel (%s)" % e)
            return
 
        # Keep the server going :)
        while(True):
 
            # Now that we have the channel, we need to build up our handler
            # Asynchronous reading from the server and the forwarded
            # channel is necessary
            r , w, x = select.select([self.request, self.channel], [], [])
 
            # Check if we need to read a request from the server, and if it is so,
            # write into the channel
            if self.request in r:
                data = self.request.recv(1024)
                if( len(data) == 0 ):
                    break # Nothing to do, exit...
 
                # Ok, transmit on the channel
                self.channel.send(data)
 
            # Check if we need to read from forwarded channel
            if self.channel in r:
                data = self.channel.recv(1024)
                if( len(data) == 0 ):
                        break # Nothing to do
 
                # Ok, so send to the server!
                self.request.send(data)
 
        # Connection complete... so
        self.channel.close()
        self.request.close()
 
 
class SnorkyForwardServer(SocketServer.ThreadingTCPServer):
    """This is the server aimed to handle all traffic"""
    daemon_threads = True
    allow_reuse_address = True
 
 
class SnorkyServerThread(threading.Thread):
    """The thread running the server"""
 
    def __init__(self, local_host, local_port, remote_host, remote_port, transport):
        threading.Thread.__init__(self)
        self.transport = transport
        self.remote_host = remote_host
        self.remote_port = remote_port
        self.local_host = local_host
        self.local_port = local_port
 
    def run(self):
        """Overloaded run method so that it runs our customized server"""
 
        # We need to define a new SubHandler with tranport in it!
        class SnorkySubHandler(SnorkyHandler):
            transport = self.transport
            remote_host = self.remote_host
            remote_port = self.remote_port
            local_host = self.local_host
            local_port = self.local_port
 
        # Start the server
        self.server = SnorkyForwardServer((self.local_host, self.local_port), SnorkySubHandler)
        self.server.serve_forever()
 
 
class SnorkySSH():
    """Provide an abstract interface to the SSH work (forwarding)
    using the Paramiko library"""
 
    def __init__(self):
        # Create the Paramiko Client and tell it we don't bother about
        # keys...
        self.SSHClient = paramiko.SSHClient()
        self.SSHClient.set_missing_host_key_policy(paramiko.AutoAddPolicy())
 
        # Our status
        self.connected = False
 
    def authenticate(self, server, username, password):
        # Check that we are not connected
        if(self.connected):
            verbose("We are connected, so doing nothing")
            return
 
        # Try connecting with Paramiko
        try:
            self.SSHClient.connect(server, 22, username, password)
            self.connected = True
        except Exception, e:
            verbose("Connection to %s@%s failed" % username, server)
            self.connected = False
 
    def openTunnel(self, remote_host, local_port, remote_port):
        # Check that we are connected
        if(self.connected == False):
            verbose("We are not connected, please call SnorkySSH.authenticate() first")
            return
 
        # Get the transport
        transport = self.SSHClient.get_transport()
 
        # Create the server thread
        self.server_thread = SnorkyServerThread('', local_port, remote_host, remote_port, transport)
 
        # Start the server
        self.server_thread.start()
 
    def closeTunnel(self):
        """Disconnect from the remote server"""
        # Just destroy anything!
        self.server_thread.server.shutdown()
        del self.server_thread
        self.connected = False