Python script to make layer 3 VPN using openssh

ssh-vpn.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
SSH VPN using pure python and openssh-client/server

Server configurations:
    * Enable ip forwarding. (linux: /etc/sysctl.conf)
    * Enable tunnel in openssh server config file: $ echo "PermitTunnel yes" >> /etc/ssh/sshd_config

Example:

    ssh-vpn.py -lroot example.com -mgvi8 -L10.8.8.19/24 -R10.8.8.20/24 -e8.8.8.8 -e8.8.4.4
    

ChangeLog:
    [2015-04-19] 1.0.1
        * Pretty code and simple document
        * Terminate and wait for process, instead of kill



"""

__version__ = '1.0.1'

import sys
import platform
import time
import argparse
import socket
import struct
import getpass
import subprocess as sb
import logging

parser = argparse.ArgumentParser(description='SSH-VPN managed by python script.')

parser.add_argument('host', metavar='HOST',
                    help='ssh connects and logs into the specified hostname.')

parser.add_argument('-l', '--user', metavar='USER', dest='user', default=getpass.getuser(),
                    help='Username')

parser.add_argument('-p', '--port', metavar='PORT', type=int, default=22, dest='port',
                    help="Remote port number.")

parser.add_argument('-i', '--tunnel-id', metavar='NUMBER', type=int, dest='tunnel_id', default='0',
                    help='The argument must be local_tun[:remote_tun].  The devices may be specified '
                         'by numerical ID or the keyword “any”, which uses the next available tunnel '
                         'device. If remote_tun is not specified, it defaults to “any”.  '
                         'The default is “any:any”.')

parser.add_argument('-L', '--local-address', metavar='ADDRESS/SUBNET', default='10.8.8.1/24', dest='local_addr',
                    help='Local address to set on tunnel(point-to-point) device.')

parser.add_argument('-R', '--remote-address', metavar='ADDRESS/SUBNET', default='10.8.8.2/24', dest='remote_addr',
                    help='Remote address to set on tunnel(point-to-point) device')

parser.add_argument('-t', '--timeout', metavar='SECONDS', type=int, default=30, dest='timeout',
                    help="SSH Timeout(second).default 30.")

parser.add_argument('-v', '--verbose', default=False, dest='verbose', action='store_true',
                    help='Verbosity.')

parser.add_argument('-c', '--compress', default=False, dest='compress', action='store_true',
                    help='Enable compression.')

parser.add_argument('-s', '--simulate', default=False, dest='simulate', action='store_true',
                    help='Just simulates the process and prints comannds instead of executing them.')

# parser.add_argument('--pid-file', default='/var/lock/ssh-vpn/pid.lock',
# help='pid lock file to be created.default: /var/lock/ssh-vpn/pid.lock .')

parser.add_argument('--alive-interval', metavar='SECONDS', type=int, default=45, dest='alive_interval',
                    help='Sets a timeout interval in seconds after which if no data has been received '
                         'from the server. default is 45.')

parser.add_argument('--attempts', metavar='NUMBER', type=int, default=3, dest='attempts',
                    help='Specifies the number of tries (one per second) to make before exiting.'
                         ' The argument must be an integer.  This may be useful in scripts if the'
                         ' connection sometimes fails.  The default is 3.')

parser.add_argument('-e', '--exception', metavar='EXCEPTION', action='append', dest='exceptions', default=[],
                    help='Do not pass packets to this destination over tunnel.')

parser.add_argument('--server-command', metavar='COMMAND', action='append', dest='server_commands', default=[],
                    help='These commands will be executed on server after connection established.')

parser.add_argument('--local-command', metavar='COMMAND', action='append', dest='local_commands', default=[],
                    help='These commands will be executed on client after connection established.')

parser.add_argument('-m', '--masquerade', default=False, dest='masquerade', action='store_true',
                    help='Masquerade all outgoing packets(NAT).')

parser.add_argument('-g', '--default-gateway', default=False, dest='default_gateway', action='store_true',
                    help="Replace the system's default gateway.")

parser.add_argument('-S', '--control-path', metavar='CONTROL_PATH', default='/var/run/ssh-vpn-tunnel-control',
                    dest='control_path',
                    help='Specifies the location of a control socket for connection sharing, or the '
                         'string “none” to disable connection sharing.  Refer to the description of '
                         'ControlPath and ControlMaster in ssh_config(5) for details.')

logger = logging.getLogger()
logging.basicConfig(level=10)
is_darwin = platform.system() == 'Darwin'


def get_default_gateway_linux():
    """
    Returns the current default gateway from `/proc`
    """
    with open("/proc/net/route") as fh:
        for line in fh:
            fields = line.strip().split()
            if fields[1] != '00000000' or not int(fields[3], 16) & 2:
                continue
            return socket.inet_ntoa(struct.pack("<L", int(fields[2], 16)))


def get_default_gateway_darwin():
    p = sb.Popen("netstat -nr | grep default | awk '{print $2}'", shell=True, stdout=sb.PIPE)
    o = p.communicate()[0].strip()
    return o


def get_default_gateway():
    if is_darwin:
        return get_default_gateway_darwin()
    else:
        return get_default_gateway_linux()


def call_cmd(cmd):
    if args.simulate:
        logger.info('Calling: %s' % cmd)
        return 0
    else:
        return sb.call(cmd.split())


def get_network_id(ip):
    ip, subnet = ip.split('/')
    octets = ip.split('.')
    octets[3] = 0
    return '%s/%s' % ('.'.join([str(i) for i in octets]), subnet)


if __name__ == '__main__':
    args = parser.parse_args()

    host = args.host.strip()
    try:
        host_addr = socket.gethostbyname(host)
    except socket.gaierror:
        log.error("Connection error")
        log.info("Check your internet connection and system dns servers")
        sys.exit(0)
    tun_name = 'tun%d' % args.tunnel_id
    default_gw = get_default_gateway()
    ssh_process = None
    exceptions_rollback = []

    try:
        logger.info('Trying to connect to : %s' % host)
        # Adding route for server

        if is_darwin:
            call_cmd('sudo route delete %s' % host_addr)
            call_cmd('sudo route add %s %s' % (host_addr, default_gw))
            exceptions = []
        else:
            call_cmd("sudo ip route replace %s via %s" % (host_addr, default_gw))
            exceptions = ['ip route replace %s via %s' % (addr, default_gw) for addr in args.exceptions]
            exceptions_rollback += ['sudo ip route del %s via %s' % (addr, default_gw) for addr in args.exceptions]

        if is_darwin:
            ssh_local_commands = ['ifconfig %s up' % tun_name]
        else:
            ssh_local_commands = ['ip link set up dev %s' % tun_name]

        if args.default_gateway:
            if is_darwin:
                ssh_local_commands += [
                    'ifconfig %s %s %s' % (tun_name, args.local_addr.split('/')[0], args.remote_addr.split('/')[0]),
                    'route delete default',
                    'route add -net 0.0.0.0 %s' % (args.local_addr.split('/')[0])]
                # TODO add exceptions for darwin
            else:
                ssh_local_commands += [
                    'ip addr replace %s remote %s dev %s' % (args.local_addr, args.remote_addr, tun_name),
                    'ip route replace default via %s' % args.remote_addr.split('/')[0]]
            ssh_local_commands += exceptions
        else:
            if exceptions:
                raise Exception('-e just allowed if you use -g.')

        ssh_local_commands += args.local_commands

        ssh_options = {'Tunnel': 'point-to-point',
                       'TunnelDevice': '%d:%d' % (args.tunnel_id, args.tunnel_id),
                       'ConnectionAttempts': args.attempts,
                       'PermitLocalCommand': 'yes',
                       'ConnectTimeout': str(args.timeout),
                       'ServerAliveInterval': str(args.alive_interval),
                       'ControlPersist': 'no',
                       'LocalCommand': ';'.join(ssh_local_commands)}

        ssh_remote_commands = ['ip link set up dev %s' % tun_name,
                               'ip addr replace %s remote %s dev %s' % (args.remote_addr, args.local_addr, tun_name),
                               'ip route replace %s via %s dev %s' % (
                                   args.local_addr.split('/')[0],
                                   args.remote_addr.split('/')[0],
                                   tun_name),
                               'ip route del %s dev %s src %s' % (
                                   get_network_id(args.local_addr),
                                   tun_name,
                                   args.remote_addr.split('/')[0])]

        ssh_remote_commands += args.server_commands

        if args.masquerade:
            if is_darwin:
                raise Exception('Masquerade was not implemented on darwin')
            else:
                ssh_remote_commands += [
                    'iptables -t nat -D POSTROUTING -s %s -j MASQUERADE' % args.local_addr.split('/')[0]]
                ssh_remote_commands += [
                    'iptables -t nat -A POSTROUTING -s %s -j MASQUERADE' % args.local_addr.split('/')[0]]

                # ssh_remote_commands += ['iptables -AINPUT -i %s -j ACCEPT' % (tun_name, )]

        ssh_cmd = 'sudo ssh ' \
                  '-MS%(control_path)s ' \
                  '%(verbose)s ' \
                  '%(compress)s ' \
                  '-p%(port)d ' \
                  '%(options)s ' \
                  '-l%(user)s ' \
                  '%(host)s ' \
                  '"%(remote_commands)s"' % dict(
                      verbose='-v' if args.verbose else '',
                      compress='-C' if args.compress else '',
                      port=args.port,
                      options=' '.join(['-o %s="%s"' % (k, v) for k, v in list(ssh_options.items())]),
                      user=args.user,
                      host=args.host,
                      remote_commands=';'.join(ssh_remote_commands),
                      control_path=args.control_path)

        if args.simulate:
            for c in ssh_local_commands:
                logging.info('Local Calling: %s' % c)
            for c in ssh_remote_commands:
                logging.info('Remote Calling: %s' % c)
            logging.info('Calling: %s' % ssh_cmd)
        else:
            ssh_process = sb.Popen(ssh_cmd, shell=True, stdout=sys.stdout, stderr=sys.stderr)
            while not ssh_process.poll():
                time.sleep(1)
    except KeyboardInterrupt:
        logger.info('You pressed CTRL+C')
    finally:
        ssh_process.terminate()
        ssh_process.wait()

        if args.default_gateway:
            if is_darwin:
                call_cmd('sudo route delete default')
                call_cmd('sudo route add default %s' % default_gw)
                call_cmd('sudo route delete %s' % host_addr)
                # TODO: remove exceptions in dirty darwin
            else:
                call_cmd('sudo ip route replace default via %s' % default_gw)
                call_cmd('sudo ip route del %s via %s' % (host_addr, default_gw))
                for er in exceptions_rollback:
                    call_cmd(er)

    print('\n')
    sys.exit(0)