vpn-whitelist

#!/usr/bin/env python

import argparse
import os
import socket
import subprocess
import sys

EXIT_FAILURE = 1
DEFAULT_WHITELIST_ADDRESSES_PATH = '/etc/vpn-whitelist.addresses'

argument_parser = argparse.ArgumentParser()
argument_parser.add_argument('addresses', nargs='*')
argument_parser.add_argument('--remove', action='store_true', help='Remove addresses from whitelist.')
argument_parser.add_argument('--file', dest='paths', metavar='PATH', nargs='*', default=[],
                             help='Paths to files containing addresses.')
argument_parser.add_argument('--gateway', help=('Specify the gateway used to bypass the VPN. If not set, the script '
                                                'will attempt to determine the default gateway.'))


def determine_default_gateway():
    try:
        output = subprocess.check_output(('ip', 'route'), universal_newlines=True)
    except subprocess.CalledProcessError as exception:
        return None, exception.output

    for line in output.splitlines():
        # noinspection PyTypeChecker
        if line.startswith('default via'):
            # noinspection PyTypeChecker
            return line.split('default via')[1].split(None, 1)[0], output

    return None, output


def get_line_set(file):
    lines = set()
    for line in file:
        line = line.strip()
        if line and not line.startswith('#'):
            lines.add(line)

    return lines


def main(arguments):
    gateway, output = arguments.gateway or determine_default_gateway()
    if not gateway:
        argument_parser.exit(EXIT_FAILURE, 'Could not determine default gateway:\n%s' % output)

    addresses = set(arguments.addresses)

    # If no addresses were given and no files specified, attempt to load addresses from the default whitelist
    if not addresses and not arguments.paths:
        arguments.paths.append(DEFAULT_WHITELIST_ADDRESSES_PATH)

    for path in (os.path.expanduser(p) for p in arguments.paths):
        if os.path.exists(path):
            with open(path) as file:
                lines = get_line_set(file)

            addresses = addresses.union(lines)

    address_high = len(addresses) - 1
    for index, address in enumerate(addresses):
        print(address + ':')
        for ip in set(i[4][0] for i in socket.getaddrinfo(address, None)):
            if arguments.remove:
                process = subprocess.run(('ip', 'route', 'del', ip, 'via', gateway), stdout=subprocess.PIPE,
                                         stderr=subprocess.PIPE, universal_newlines=True)

                if process.returncode == 0:
                    print('  Removed', ip, 'from whitelist')
                else:
                    print('  Failed to remove %s from whitelist: %s' % (ip, process.stderr.strip()), file=sys.stderr)
            else:
                process = subprocess.run(('ip', 'route', 'add', ip, 'via', gateway), stdout=subprocess.PIPE,
                                         stderr=subprocess.PIPE, universal_newlines=True)

                if process.returncode == 0:
                    print('  Added', ip, 'to whitelist')
                else:
                    print('  Failed to whitelist %s: %s' % (ip, process.stderr.strip()), file=sys.stderr)

        if index != address_high:
            print()


if __name__ == '__main__':
    arguments = argument_parser.parse_args()
    main(arguments)