forwarder.py

#!/usr/bin/env python

# This code is licensed according to the main package that it depends on:
# http://pypi.python.org/pypi/ssh
# License :: OSI Approved :: GNU Library or Lesser General Public License (LGPL)
# http://www.gnu.org/licenses/lgpl.html

import sys, os, os.path
import select
import SocketServer
import logging
import threading
from logging import info, debug, error
from types import ListType

import ssh
from ssh import SSHClient, SSHConfig, WarningPolicy, client
from ssh.client import SSH_PORT

class ForwardServer (SocketServer.ThreadingTCPServer):
    daemon_threads = True
    allow_reuse_address = True

class Handler (SocketServer.BaseRequestHandler):
    def handle(self):
        try:
            chan = self.ssh_transport.open_channel('direct-tcpip', 
                                                   (self.target_host, self.target_port), 
                                                   self.request.getpeername())
        except Exception, e:
            error('Incoming request to %s:%d failed: %s' % 
                  (self.target_host, self.target_port, repr(e)))
            return
        if chan is None:
            error('Incoming request to %s:%d was rejected by the SSH server.' % 
                  (self.target_host, self.target_port))
            return

        info('Connected!  Tunnel open %r -> %r -> %r' % 
              (self.request.getpeername(), chan.getpeername(), 
               (self.target_host, self.target_port)))
        while True:
            r, w, x = select.select([self.request, chan], [], [])
            if self.request in r:
                data = self.request.recv(1024)
                if len(data) == 0:
                    break
                chan.send(data)
            if chan in r:
                data = chan.recv(1024)
                if len(data) == 0:
                    break
                self.request.send(data)
        chan.close()
        peer = self.request.getpeername()
        self.request.close()
        info('Tunnel closed from %r' % (peer,))


class SSHForwarder(SSHClient):
    def __init__(self, proxy_host, port, target_host='localhost', user=None, config=None):
        super(SSHForwarder, self).__init__()
        if config:
            self.config = config
        else:
            self.config = SSHConfig()
            try:
                f = open(os.path.join(os.environ['HOME'], ".ssh/config"))
                self.config.parse(f)
            except Exception:
                sys.exc_clear()
                
        hconfig = self.config.lookup(proxy_host)
        user = hconfig['user'] if (hconfig and hconfig['user']) else os.environ['USER'] 
        self.host = proxy_host
        self.user = user
        self.load_system_host_keys()
        self.set_missing_host_key_policy(WarningPolicy)

        info("Connecting to proxy host '{host}' as user '{user}'".format(**self.__dict__))
        self.connect(self.host, SSH_PORT, username=self.user)
        
        self.forwarders = {}
        self.daemons = {}
        for p in set(port if (type(port) is ListType) else [port]):
            target = target_host
            class SubHandler (Handler):
                target_host = target
                target_port = p
                ssh_transport = self.get_transport()

            self.forwarders[p] = ForwardServer(('', p), SubHandler)
            self.daemons[p] = threading.Thread(target=self.forwarders[p].serve_forever)
            self.daemons[p].daemon = True
            forward_info = {'p':str(p), 'target':target, 'host':self.host}
            info("Forwarding port '{p}' to target host '{target}' via proxy host '{host}'".format(**forward_info))
            self.daemons[p].start()

    def close(self):
        try:
            for f in self.forwarders.values():
                f.shutdown()
                f.server_close()
        finally:
            super(SSHForwarder, self).close()

def main():
    try:
        proxy_host = 'proxyhost'
        target_port = 4000
        target_host = 'targethost'
        f = SSHForwarder(proxy_host, target_port, target_host)
        # .... DO SOME STUFF ....
    finally:
        if f:
            f.close()
        
if __name__ == "__main__":
    main()