vpngen/vpngen-cli.py

#! /usr/bin/env python3

from __future__ import print_function
import argparse
import json
import os
import sys
import glob
import json
import os
import os.path
import re
import shutil
from enum import Enum
from subprocess import call


class VpnGenError(Enum):
    Success = 0,
    VpnAlreadyExists = 1,
    VpnDoesNotExists = 2,
    ClientAlreadyExists = 3,
    ClientDoesNotExists = 4


class VpnGen:
    default_config_base_dir = ""
    default_config_file = ""
    default_client_config_file = ""
    ovpn_config_path = ""

    def __init__(self, default_config_path, ovpn_config_path):
        self.default_config_base_dir = os.path.abspath(default_config_path)
        self.default_config_file = "%s.conf" % self.default_config_base_dir
        self.default_client_config_file = "%s%sclients%sclient.conf" % (self.default_config_base_dir, os.sep, os.sep)
        self.ovpn_config_path = os.path.abspath(ovpn_config_path)

    def f7(self, seq):
        seen = set()
        seen_add = seen.add
        return [x for x in seq if not (x in seen or seen_add(x))]

    def get_vpn_vars(self):
        with open(self.default_config_file, "r") as f:
            default_config = f.read()

        variables = re.findall('\$\{([^}]+)}', default_config)
        variables += ["KEY_COUNTRY", "KEY_PROVINCE", "KEY_CITY", "KEY_ORG", "KEY_EMAIL"]
        variables = self.f7(variables)

        return variables

    def get_client_vars(self, vpn_name):
        default_client_config_path = self.get_client_default_config_path(vpn_name)
        if not os.path.exists(default_client_config_path):
            return None

        with open(default_client_config_path, "r") as f:
            default_config = f.read()
        variables = re.findall('\$\{([^}]+)}', default_config)
        variables = self.f7(variables)

        vpn_variables = self.get_vpn_vars()
        real_variables = []

        for var in variables:
            if var not in vpn_variables and var != "client":
                real_variables.append(var)

        return real_variables

    def get_base_dir(self, vpn_name):
        return "%s%s%s%s" % (self.ovpn_config_path, os.sep, vpn_name, os.sep)

    def get_config_path(self, vpn_name):
        return "%s%s%s.conf" % (self.ovpn_config_path, os.sep, vpn_name)

    def get_vpn_variables_path(self, vpn_name):
        base_dir = self.get_base_dir(vpn_name)
        return "%svpngen.json" % base_dir

    def get_easy_rsa_dir(self, vpn_name):
        base_dir = self.get_base_dir(vpn_name)
        return "%seasy-rsa%s" % (base_dir, os.sep)

    def get_easy_rsa_key_dir(self, vpn_name):
        easyrsadir = self.get_easy_rsa_dir(vpn_name)
        return "%skeys%s" % (easyrsadir, os.sep)

    def get_pkitool_path(self, vpn_name):
        easyrsadir = self.get_easy_rsa_dir(vpn_name)
        return "%spkitool" % easyrsadir

    def get_client_default_config_path(self, vpn_name):
        base_dir = self.get_base_dir(vpn_name)
        return "%s%sclients%sclient.conf" % (base_dir, os.sep, os.sep)

    def get_client_dir(self, vpn_name, client_name):
        base_dir = self.get_base_dir(vpn_name)
        return "%sclients%s%s-%s%s" % (base_dir, os.sep, client_name, vpn_name, os.sep)

    def get_client_config_path(self, vpn_name, client_name):
        client_dir = self.get_client_dir(vpn_name, client_name)
        return "%s%s-%s.conf" % (client_dir, client_name, vpn_name)

    def get_client_variables_path(self, vpn_name, client_name):
        client_dir = self.get_client_dir(vpn_name, client_name)
        return "%svpngen.json" % client_dir

    def get_client_generated_files_paths(self, vpn_name, client_name):
        keys_dir = self.get_easy_rsa_key_dir(vpn_name,)
        return [
            "%s%s.crt" % (keys_dir, client_name),
            "%s%s.key" % (keys_dir, client_name)
        ]

    def get_client_tarball_path(self, vpn_name, client_name):
        base_dir = self.get_base_dir(vpn_name)
        return "%sclients%s%s-%s.tar.bz2" % (base_dir, os.sep, client_name, vpn_name)

    def get_server_needed_files_paths(self, vpn_name):
        keys_dir = self.get_easy_rsa_key_dir(vpn_name)
        return [
            "%sca.crt" % keys_dir,
            "%sta.key" % keys_dir
        ]

    def get_client_misc_files_paths(self, vpn_name):
        base_dir = self.get_base_dir(vpn_name)
        return glob.glob("%smisc-files%s*" % (base_dir, os.sep))

    def get_all_needed_files_paths(self, vpn_name, client_name):
        return self.get_client_generated_files_paths(vpn_name, client_name) +\
               self.get_server_needed_files_paths(vpn_name) +\
               self.get_client_misc_files_paths(vpn_name)

    def get_server_variables(self, vpn_name):
        with open(self.get_vpn_variables_path(vpn_name), "r") as f:
            return json.load(f)['variables']

    def get_client_variables(self, vpn_name, client_name):
        with open(self.get_client_variables_path(vpn_name, client_name), "r") as f:
            return json.load(f)['variables']

    def get_client_list(self, vpn_name):
        base_dir = self.get_base_dir(vpn_name)
        files_paths = glob.glob("%sclients%s*" % (base_dir, os.sep))
        files_names = list(map(lambda file_path: os.path.basename(file_path), files_paths))

        clients = []
        for file_name in files_names:
            if file_name != 'client.conf' and not file_name.endswith(".tar.bz2"):
                clients.append(file_name[0:len(file_name) - len(vpn_name) - 1])
        return clients

    def setup_vars(self, vpn_name, variables):
        os.environ["KEY_COUNTRY"] = variables['KEY_COUNTRY']
        os.environ["KEY_PROVINCE"] = variables['KEY_PROVINCE']
        os.environ["KEY_CITY"] = variables['KEY_CITY']
        os.environ["KEY_ORG"] = variables['KEY_ORG']
        os.environ["KEY_OU"] = variables['KEY_ORG']
        os.environ["KEY_CN"] = variables['KEY_ORG']
        os.environ["KEY_NAME"] = variables['KEY_ORG']
        os.environ["KEY_EMAIL"] = variables['KEY_EMAIL']
        os.environ["KEY_SIZE"] = variables['KEY_SIZE']
        os.environ["CA_EXPIRE"] = variables['CA_EXPIRE']
        os.environ["KEY_EXPIRE"] = variables['KEY_EXPIRE']

        self.setup_vars_openssl(vpn_name)

    def setup_vars_openssl(self, vpn_name):
        easyrsadir = self.get_easy_rsa_dir(vpn_name)
        os.environ["EASY_RSA"] = easyrsadir
        os.environ["OPENSSL"] = "openssl"
        os.environ["PKCS11TOOL"] = "pkcs11-tool"
        os.environ["GREP"] = "grep"
        os.environ["KEY_CONFIG"] = "%s%s" % (easyrsadir, "openssl.cnf")
        os.environ["KEY_DIR"] = "%s%s" % (easyrsadir, "keys")
        os.environ["PKCS11_MODULE_PATH"] = "dummy"
        os.environ["PKCS11_PIN"] = "dummy"

    def create_vpn(self, vpn_name, variables):
        base_dir = self.get_base_dir(vpn_name)
        conf_file = self.get_config_path(vpn_name)
        conf_vpngen_file = self.get_vpn_variables_path(vpn_name)
        if os.path.exists(base_dir) or os.path.exists(conf_file):
            return VpnGenError.VpnAlreadyExists

        with open(self.default_config_file, "r") as f:
            default_config = f.read()

        variables['name'] = vpn_name
        for variable in variables:
            default_config = default_config.replace("${%s}" % variable, variables[variable])

        os.makedirs(base_dir)
        with open(conf_file, "w") as f:
            f.write(default_config)

        os.rmdir(base_dir)
        shutil.copytree(self.default_config_base_dir, base_dir)

        curdir = os.curdir
        easyrsadir = self.get_easy_rsa_dir(vpn_name)
        pkitool = self.get_pkitool_path(vpn_name)
        os.chdir(easyrsadir)

        self.setup_vars(vpn_name, variables)

        call([".%sclean-all" % os.sep])
        call([pkitool, "--initca", "-batch"])
        call([pkitool, "--server", "server", "-batch"])
        call([".%sbuild-dh" % os.sep])
        call(["openssl", "ca", "-gencrl",
              "-keyfile", "keys%sca.key" % os.sep,
              "-cert", "keys%sca.crt" % os.sep,
              "-out", "keys%scrl.pem" % os.sep,
              "-config", "openssl.cnf"])

        del os.environ["KEY_OU"]
        del os.environ["KEY_CN"]
        del os.environ["KEY_NAME"]

        call(["openvpn", "--genkey", "--secret", "keys%sta.key" % os.sep])

        with open(conf_vpngen_file, "w") as f:
            json.dump({'variables': variables}, f, indent=4, separators=(',', ': '))

        os.chdir(curdir)

        return VpnGenError.Success

    def remove_vpn(self, vpn_name):
        base_dir = self.get_base_dir(vpn_name)
        conf_file = self.get_config_path(vpn_name)
        if not os.path.exists(base_dir) and not os.path.exists(conf_file):
            return VpnGenError.VpnDoesNotExists
        os.remove(conf_file)
        shutil.rmtree(base_dir)
        return VpnGenError.Success

    def create_client(self, vpn_name, client_name, variables):
        base_dir = self.get_base_dir(vpn_name)
        if not os.path.exists(base_dir):
            return VpnGenError.VpnDoesNotExists
        client_dir = self.get_client_dir(vpn_name, client_name)
        if os.path.exists(client_dir):
            return VpnGenError.ClientAlreadyExists

        curdir = os.curdir
        easyrsadir = self.get_easy_rsa_dir(vpn_name)
        pkitool = self.get_pkitool_path(vpn_name)
        os.chdir(easyrsadir)

        self.setup_vars(vpn_name, variables)
        os.environ["KEY_CN"] = client_name
        os.environ["KEY_NAME"] = client_name

        call([pkitool, client_name])

        os.chdir(curdir)

        os.makedirs(client_dir)

        return self.rebuild_client(vpn_name, client_name, variables)

    def remove_client(self, vpn_name, client_name):
        base_dir = self.get_base_dir(vpn_name)
        if not os.path.exists(base_dir):
            return VpnGenError.VpnDoesNotExists
        client_dir = self.get_client_dir(vpn_name, client_name)
        if not os.path.exists(client_dir):
            return VpnGenError.ClientDoesNotExists

        self.setup_vars_openssl(vpn_name)

        curdir = os.curdir
        easyrsadir = self.get_easy_rsa_dir(vpn_name)
        os.chdir(easyrsadir)

        call(["./revoke-full", client_name])

        os.chdir(curdir)

        return VpnGenError.Success

    def rebuild_client(self, vpn_name, client_name, variables):
        variables["client"] = client_name
        client_dir = self.get_client_dir(vpn_name, client_name)
        client_conf_file = self.get_client_config_path(vpn_name, client_name)
        client_default_config_path = self.get_client_default_config_path(vpn_name)
        with open(client_default_config_path, "r") as f:
            client_default_config = f.read()

        for variable in variables:
            client_default_config = client_default_config.replace("${%s}" % variable, variables[variable])

        files_names = glob.glob("%s%s*" % (client_dir, os.sep))
        for file_name in files_names:
            os.remove(file_name)

        with open(client_conf_file, "w") as f:
            f.write(client_default_config)

        files_paths = self.get_all_needed_files_paths(vpn_name, client_name)
        for file_path in files_paths:
            dest = "%s%s-%s-%s" % (client_dir, client_name, vpn_name, os.path.basename(file_path))
            shutil.copy(file_path, dest)

        split = os.path.splitext(client_conf_file)
        client_ovpn_file = "%s.ovpn" % split[0]
        shutil.copy(client_conf_file, client_ovpn_file)

        files_names = glob.glob("%s%s*" % (client_dir, os.sep))
        files_names = list(map(lambda file_path: os.path.basename(file_path), files_names))

        call(["tar", "cfj", self.get_client_tarball_path(vpn_name, client_name),
              "-C", client_dir] + files_names)

        client_variables = {}
        for variable in self.get_client_vars(vpn_name):
            client_variables[variable] = variables[variable]

        with open(self.get_client_variables_path(vpn_name, client_name), "w") as f:
            json.dump({'variables': client_variables}, f, indent=4, separators=(',', ': '))

        return VpnGenError.Success



def eprint(*args, **kwargs):
    print(*args, file=sys.stderr, **kwargs)


def create_variables(variables, defaults):
    variables_set = defaults.copy()
    for variable in variables:
        if variable == 'name' or variable == 'client':
            continue
        default = variables_set[variable] if variable in variables_set else ''
        print("Enter a value for '%s' [%s]: " % (variable, default), end='', flush=True)
        value = sys.stdin.readline()[:-1]
        if value != '' or variable not in variables_set:
            variables_set[variable] = value
    return variables_set


def main():
    parser = argparse.ArgumentParser(description='Manage OpenVPN VPNs')
    parser.add_argument('--vpn', help='The VPN to use', required=True)
    parser.add_argument('--config', dest='config', default='/etc/vpngen/vpngen.json', help='Configuration file path')

    parser.add_argument('--create', help='Create a VPN', action='store_true')
    parser.add_argument('--remove', help='Remove a VPN', action='store_true')
    parser.add_argument('--create-client', help='Create a client for the VPN', metavar='CLIENT')
    parser.add_argument('--remove-client', help='Remove a client for the VPN', metavar='CLIENT')
    parser.add_argument('--rebuild-client', help='Rebuild a client configuration', metavar='CLIENT')
    parser.add_argument('--rebuild-clients', help='Rebuild clients configurations', action='store_true')

    args = parser.parse_args()

    with open(args.config, "r") as f:
        config = json.load(f)

    vpn_name = config['vpnPrefix'] + args.vpn + config['vpnSuffix']
    if args.create_client is not None:
        client_name = args.create_client
    elif args.remove_client is not None:
        client_name = args.remove_client
    elif args.rebuild_client is not None:
        client_name = args.rebuild_client
    else:
        client_name = None
    if client_name is not None:
        client_name = config['clientPrefix'] + client_name + config['clientSuffix']

    vpng = VpnGen(config['defaultConfigPath'], config['ovpnConfigPath'])

    config_path = vpng.get_vpn_variables_path(vpn_name)
    if os.path.exists(config_path):
        with open(config_path, "r") as f:
            data = json.load(f)
        config['defaults'].update(data['variables'])

    if args.create:
        default_variables = config['defaults'].copy()
        variables = create_variables(vpng.get_vpn_vars(), default_variables)
        res = vpng.create_vpn(vpn_name, variables)
        if res == VpnGenError.Success:
            print("VPN %s created successfully" % vpn_name)
        else:
            eprint("Failed to create VPN %s: %s" % (vpn_name, res))
            exit(1)
    elif args.remove:
        res = vpng.remove_vpn(vpn_name)
        if res == VpnGenError.Success:
            print("VPN %s removed successfully" % vpn_name)
        else:
            eprint("Failed to remove VPN %s: %s" % (vpn_name, res))
            exit(1)
    elif args.create_client:
        default_variables = config['defaults'].copy()
        default_variables.update(vpng.get_server_variables(vpn_name))
        variables = create_variables(vpng.get_client_vars(vpn_name), default_variables)
        if variables is None:
            res = VpnGenError.VpnDoesNotExists
        else:
            res = vpng.create_client(vpn_name, client_name, variables)
        if res == VpnGenError.Success:
            print("Client %s created successfully on VPN %s" % (client_name, vpn_name))
        else:
            eprint("Failed to create client %s on VPN %s: %s" % (client_name, vpn_name, res))
            exit(1)
    elif args.remove_client:
        res = vpng.remove_client(vpn_name, client_name)
        if res == VpnGenError.Success:
            print("Client %s removed successfully on VPN %s" % (client_name, vpn_name))
        else:
            eprint("Failed to remove client %s on VPN %s: %s" % (client_name, vpn_name, res))
            exit(1)
    elif args.rebuild_client:
        default_variables = config['defaults'].copy()
        default_variables.update(vpng.get_server_variables(vpn_name))
        default_variables.update(vpng.get_client_variables(vpn_name, client_name))

        res = vpng.rebuild_client(vpn_name, client_name, default_variables)
        if res == VpnGenError.Success:
            print("Client %s configuration rebuilt successfully on VPN %s" % (client_name, vpn_name))
        else:
            eprint("Failed to rebuild client %s configuration on VPN %s: %s" % (client_name, vpn_name, res))
            exit(1)
    elif args.rebuild_clients:
        default_variables = config['defaults'].copy()
        default_variables.update(vpng.get_server_variables(vpn_name))

        for client_name in vpng.get_client_list(vpn_name):
            variables = default_variables.copy()
            variables.update(vpng.get_client_variables(vpn_name, client_name))
            res = vpng.rebuild_client(vpn_name, client_name, variables)
            if res == VpnGenError.Success:
                print("Client %s configuration rebuilt successfully on VPN %s" % (client_name, vpn_name))
            else:
                eprint("Failed to rebuild client %s configuration on VPN %s: %s" % (client_name, vpn_name, res))
                exit(1)


main()