sshroute.py

#!/usr/bin/env python2
"""
Generate a script to create a multi-hop SSH connection.

Usage: sshroute.py [options...] HOST

Options:
  -l SPEC, --local=SPEC   Tunnel a port - local:via:remote
  -r SPEC, --remote=SPEC  Reverse tunnel a port - remote:via:local
  -h HOST, --host=HOST    Go via another host
"""
import sys
from functools import partial

import docopt
import schema


def read_ports(spec, keys):
    parts = map(int, spec.split(':'))
    if len(parts) == 1:
        return dict(zip(keys, parts * len(keys)))
    elif len(parts) == 2:
        return dict(zip(keys, [parts[0]] * (len(keys) - 1) + [parts[1]]))
    elif len(parts) == len(keys):
        return dict(zip(keys, parts))
    else:
        raise ValueError('SPEC must be "port", "{}:{}" or "{}"'.format(
            keys[0], keys[-1], ':'.join(keys)))

def read_local(spec):
    return read_ports(spec, ('local', 'via', 'remote'))

def read_remote(spec):
    return read_ports(spec, ('remote', 'via', 'local'))

def read_host(host):
    parts = host.rsplit(':', 1)
    if len(parts) == 2:
        return dict(host=parts[0], port=int(parts[1]))
    else:
        return dict(host=parts[0], port=None)

def create_dest(host):
    if host['port'] is None:
        return host['host']
    else:
        return '-p {port} {host}'.format(**host)

def build_command(local, remote, hosts, destination):
    if len(hosts) == 0:
        # If there are no intermediate hosts, do it the simple way
        forwards = ['-L{local}:localhost:{remote}'.format(**f) for f in local] + \
                   ['-R{remote}:localhost:{local}'.format(**f) for f in remote]
        return 'ssh ' + ' '.join(forwards) + create_dest(destination)
    else:
        # Forwarding arguments for the first hop of the tunnel
        first = ' '.join(['-L{local}:localhost:{via}'.format(**f) for f in local] +
                         ['-R{via}:localhost:{local}'.format(**f) for f in remote])
        # Forwarding arguments for the last hop of the tunnel
        last = ' '.join(['-L{via}:localhost:{remote}'.format(**f) for f in local] +
                        ['-R{remote}:localhost:{via}'.format(**f) for f in remote])
        # Forwarding arguments for intermediate hops
        via = ' '.join(['-L{via}:localhost:{via}'.format(**f) for f in local] +
                       ['-R{via}:localhost:{via}'.format(**f) for f in remote])

        # Pairs of (host, forwards) for each step in the chain
        parts = zip(hosts + [destination], [first] + [via] * (len(hosts) - 1) + [last])
        # Iteratively build the command from innermost/last to outermost/first
        cmd = ''
        for i, (host, fwds) in reversed(list(enumerate(parts))):
            slashes = '\\' * i
            if cmd:
                cmd = ' {0}"{1}{0}"'.format(slashes, cmd)
            cmd = 'ssh {fwds} {dest}{cmd}'.format(fwds=fwds, dest=create_dest(host), cmd=cmd)

        return cmd


if __name__ == '__main__':
    args = docopt.docopt(__doc__, sys.argv[1:])
    schema = schema.Schema({
        '--local': [schema.Use(read_local)],
        '--remote': [schema.Use(read_remote)],
        '--host': [schema.Use(read_host)],
        'HOST': schema.Use(read_host),
    })
    args = schema.validate(args)

    print '#/bin/bash -i'
    print build_command(args['--local'], args['--remote'], args['--host'], args['HOST'])