tunnel

#!/usr/bin/env python

# Copyright Anne Archibald 2011
# based on code from paramiko
# This code is licensed under the LGPL

import sys, getpass
import socket, select
import SocketServer
import threading
import time
import struct
from optparse import OptionParser

from paramiko import SSHClient, AutoAddPolicy, Transport
from paramiko.resource import ResourceManager

SSH_PORT=22

def verbose(s):
    if g_verbose:
        print s

def user_machine_port(m):
    if '@' in m:
        u, r = m.split('@')
    else:
        u, r = None, m
    if ':' in m:
        m, p = r.split(':')
    else:
        m, p = r, SSH_PORT
    return u, m, int(p)

class SSHSockClient(SSHClient):
    """subclass paramiko.client.SSHClient that uses an existing socket-like object

    Our base class (SSHClient) always creates a new socket
    object. This subclass lets us re-use an existing socket-like
    object."""

    def connect(self, hostname,
            port=SSH_PORT, username=None,
            password=None, pkey=None,
            key_filename=None,
            timeout=None,
            allow_agent=True,
            look_for_keys=True,
            compress=False,
            sock=None):
        """
        Connect to a socket connected SSH server and authenticate to the
        server.  The server's host key
        is checked against the system host keys (see L{load_system_host_keys})
        and any local host keys (L{load_host_keys}).  If the server's hostname
        is not found in either set of host keys, the missing host key policy
        is used (see L{set_missing_host_key_policy}).  The default policy is
        to reject the key and raise an L{SSHException}.

        Authentication is attempted in the following order of priority:

            - The C{pkey} or C{key_filename} passed in (if any)
            - Any key we can find through an SSH agent
            - Any "id_rsa" or "id_dsa" key discoverable in C{~/.ssh/}
            - Plain username/password auth, if a password was given

        If a private key requires a password to unlock it, and a password is
        passed in, that password will be used to attempt to unlock the key.

        @param hostname: the server to connect to
        @type hostname: str
        @param port: the server port to connect to
        @type port: int
        @param username: the username to authenticate as (defaults to the
            current local username)
        @type username: str
        @param password: a password to use for authentication or for unlocking
            a private key
        @type password: str
        @param pkey: an optional private key to use for authentication
        @type pkey: L{PKey}
        @param key_filename: the filename, or list of filenames, of optional
            private key(s) to try for authentication
        @type key_filename: str or list(str)
        @param timeout: an optional timeout (in seconds) for the TCP connect
        @type timeout: float
        @param allow_agent: set to False to disable connecting to the SSH agent
        @type allow_agent: bool
        @param look_for_keys: set to False to disable searching for discoverable
            private key files in C{~/.ssh/}
        @type look_for_keys: bool
        @param compress: set to True to turn on compression
        @type compress: bool
        @param sock: the socket to connect to
        @type sock: socket-like object

        @raise BadHostKeyException: if the server's host key could not be
            verified
        @raise AuthenticationException: if authentication failed
        @raise SSHException: if there was any other error connecting or
            establishing an SSH session
        """

        if sock is None:
            for (family, socktype, proto, canonname, sockaddr) in socket.getaddrinfo(hostname, port, socket.AF_UNSPEC, socket.SOCK_STREAM):
                if socktype == socket.SOCK_STREAM:
                    af = family
                    addr = sockaddr
                    break
            else:
                # some OS like AIX don't indicate SOCK_STREAM support, so just guess. :(
                af, _, _, _, addr = socket.getaddrinfo(hostname, port, socket.AF_UNSPEC, socket.SOCK_STREAM)
            sock = socket.socket(af, socket.SOCK_STREAM)
            if timeout is not None:
                try:
                    sock.settimeout(timeout)
                except:
                    pass
            sock.connect(addr)

        t = self._transport = Transport(sock)
        t.use_compression(compress=compress)
        if self._log_channel is not None:
            t.set_log_channel(self._log_channel)
        t.start_client()
        ResourceManager.register(self, t)

        server_key = t.get_remote_server_key()
        keytype = server_key.get_name()

        if port == SSH_PORT:
            server_hostkey_name = hostname
        else:
            server_hostkey_name = "[%s]:%d" % (hostname, port)
        our_server_key = self._system_host_keys.get(server_hostkey_name, {}).get(keytype, None)
        if our_server_key is None:
            our_server_key = self._host_keys.get(server_hostkey_name, {}).get(keytype, None)
        if our_server_key is None:
            # will raise exception if the key is rejected; let that fall out
            self._policy.missing_host_key(self, server_hostkey_name, server_key)
            # if the callback returns, assume the key is ok
            our_server_key = server_key

        if server_key != our_server_key:
            raise BadHostKeyException(hostname, server_key, our_server_key)

        if username is None:
            username = getpass.getuser()

        if key_filename is None:
            key_filenames = []
        elif isinstance(key_filename, (str, unicode)):
            key_filenames = [ key_filename ]
        else:
            key_filenames = key_filename
        self._auth(username, password, pkey, key_filenames, allow_agent, look_for_keys)


def setup_tunnel(remote_machine_list):
    clients = []
    for m in remote_machine_list:
        user, machine, port = user_machine_port(m)
        if clients:
            verbose("opening channel to %s" % m)
            t = clients[-1].get_transport()
            # I don't think these lport numbers necessarily mean anything
            ch = t.open_channel('direct-tcpip',
                    dest_addr=(machine,port),
                    src_addr=('localhost',0))
        else:
            ch = None

        verbose("starting tunneled connection")
        c = SSHSockClient()
        clients.append(c)
        c.load_system_host_keys()
        c.set_missing_host_key_policy(AutoAddPolicy())
        kwargs=dict(sock=ch, hostname=machine, username=user, port=port,
                    compress=len(clients)==len(remote_machine_list))
        verbose("  arguments to connect(): "+str(kwargs))
        c.connect(sock=ch, hostname=machine, username=user, port=port,
                compress=len(clients)==len(remote_machine_list))
        verbose("Connected to %s" % m)

    clients[-1].get_transport().set_keepalive(120)
    return clients


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.chain_host, self.chain_port),
                                                   self.request.getpeername())
        except Exception, e:
            verbose('Incoming request to %s:%d failed: %s' % (self.chain_host,
                                                              self.chain_port,
                                                              repr(e)))
            return
        if chan is None:
            verbose('Incoming request to %s:%d was rejected by the SSH server.' %
                    (self.chain_host, self.chain_port))
            return

        verbose('Connected!  Tunnel open %r -> %r -> %r' % (self.request.getpeername(),
                                                            chan.getpeername(), (self.chain_host, self.chain_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()
        verbose('Tunnel closed from %r' % (self.request.getpeername(),))
        self.request.close()

class SOCKS4aHandler (SocketServer.BaseRequestHandler):
    def handle(self):
        data = ""
        while True:
            r, w, x = select.select([self.request], [], [])
            if self.request in r:
                data = self.request.recv(1024)
                if len(data) == 0:
                    continue
                if data[0]!=chr(4):
                    # not the right version of SOCKS
                    verbose('Failed: Not a SOCKS4 connection')
                    return
                if len(data)<8:
                    # packet too short!?
                    verbose('Failed: Packet too short')
                    return
                v, request, port, ip = struct.unpack("!BBHI",data[:8])
                if request!=1:
                    verbose('Failed: did not request outgoing connection')
                    return
                user, rest = data[8:].split(chr(0),1)
                verbose("User claims to be %s" % repr(user))
                if ip<256:
                    host, rest = rest.split(chr(0),1)
                else:
                    host = socket.inet_ntoa(data[4:8])
                data=rest
                break
        try:
            chan = self.ssh_transport.open_channel('direct-tcpip',
                                                   (host, port),
                                                   self.request.getpeername())
        except Exception, e:
            verbose('Incoming request to %s:%d failed: %s' % (host,
                                                              port,
                                                              repr(e)))
            self.request.send(struct.pack('!BBHI',0,0x5b,0,0))
            return
        self.request.send(struct.pack('!BBHI',0,0x5a,0,0))

        verbose('Connected!  Tunnel open %r -> %r -> %r' %
                (self.request.getpeername(),
                 chan.getpeername(), (host, port)))

        if data:
            verbose("Dropping trailing characters on SOCKS request: %s" % repr(data))
        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()
        verbose('Tunnel closed from %r' % (self.request.getpeername(),))
        self.request.close()


def forward_tunnel(local_port, remote_host, remote_port, transport):
    # this is a little convoluted, but lets me configure things for the Handler
    # object.  (SocketServer doesn't give Handlers any way to access the outer
    # server normally.)
    class SubHandler (Handler):
        chain_host = remote_host
        chain_port = remote_port
        ssh_transport = transport
    s = threading.Thread(target=ForwardServer(('localhost', local_port), SubHandler).serve_forever)
    s.setDaemon(True)
    s.start()
    verbose("started forwarding from %d on localhost to %d on remote host" %
            (local_port, remote_port))

def socks(port, transport):
    # this is a little convoluted, but lets me configure things for the Handler
    # object.  (SocketServer doesn't give Handlers any way to access the outer
    # server normally.)
    class SOCKSSubHandler (SOCKS4aHandler):
        ssh_transport = transport
    s = threading.Thread(target=ForwardServer(('localhost', port), SOCKSSubHandler).serve_forever)
    s.setDaemon(True)
    s.start()
    verbose("started SOCKS4a forwarding from %d on localhost" % port)


if __name__=='__main__':
    parser = OptionParser(
            usage="%prog [options] machine_1 machine_2 ... machine_n",
            epilog=("Open an SSH connection to each machine_i, "
                    "tunneled through the connection to "
                    "machine_(i-1). Forwardings then transfer "
                    "connections to localhost to connections from "
                    "machine_n to itself; this way no unencrypted "
                    "traffic is sent across the network. SOCKS "
                    "forwarding allows a SOCKS-aware client to connect "
                    "to arbitrary machines while apparently coming from "
                    "machine_n. DNS requests may also be handled by the "
                    "SOCKS server. Note that all ports opened on the "
                    "local machine are accessible to any user on the "
                    "local machine, but to no one else."))
    parser.add_option("-v", "--verbose", action="store_true",
            help="Verbose output.")
    parser.add_option("-T", "--taken", action="store_true", 
            help="Report success if the local port is already taken.")
    parser.add_option("-D", "--socks", 
            metavar="PORT", type=int,
            help="Act as a SOCKS4a proxy listening on PORT (to localhost only)")
    parser.add_option("-f", "--forward",
            action="append", metavar="LOCAL:REMOTE",
            help="forward connections from localhost port LOCAL to final host port REMOTE")
    options, args = parser.parse_args(sys.argv[1:])
    # FIXME: add an option to forward stdin/stdout

    global g_verbose
    g_verbose = options.verbose

    clients = setup_tunnel(args)
    for s in options.forward:
        l, r = s.split(":")
        l, r = int(l), int(r)
        try:
            forward_tunnel(l, "localhost", r, clients[-1].get_transport())
        except socket.error, e:
            if not (e.errno==98 and options.taken):
                raise

    if options.socks:
        try:
            socks(options.socks, clients[-1].get_transport())
        except socket.error, e:
            if not (e.errno==98 and options.taken):
                raise
    while True:
        time.sleep(1)