Quickly generate a server and several clients for Wireguard.

wg-genconf.py

#!/usr/bin/env python3

from typing import List, Tuple
from ipaddress import _BaseAddress, _BaseNetwork

import io
import math
import random
import argparse
import ipaddress
import subprocess

# region exception types definition


class MyError(Exception):
    pass

# region argparse


def _arg_parser() -> None:
    parser = argparse.ArgumentParser(
        description='Generate Wireguard config files',
        epilog=f'''Examples:

  wg-genconf -p 12345 -w   1.2.3.4  10.0.0.123/24  4
      Generate 10.0.0.1/24 for gateway, and 10.0.0.2/24 to 10.0.0.5/24 as clients.
      All configs are written to /etc/wireguard/temp-{{1-5}}.conf .

  wg-genconf   1.2.3.4  10.0.0.0/24  Alice Bob
      Generate 10.0.0.1/24 for gateway, 10.0.0.2/24 for Alice and 10.0.0.3/24 for Bob.
      All configs are written to stdout, with random listen port on gateway.

  wg-genconf -p 8888 -n 4 -u the_gw_public_key -w   1.2.3.4  10.0.0.0/24  Eve
      Generate 10.0.0.4/24 for Eve. Mostly used for adding users to an existing configuration.
      Gateway's extra config is written to /etc/wireguard/temp-0.conf, and Eve's config to /etc/wireguard/temp-3.conf .
''',
        formatter_class=argparse.RawTextHelpFormatter,
    )
    parser.add_argument('-p', '--port', type=int, nargs=None,
                        help='Listen port. If not set, use random port in 10000-65534.')
    parser.add_argument('-w', '--write-to-file', action='store_true',
                        help='Write generated config files to /etc/wireguard/temp-*.conf. If not set, print to screen.')
    parser.add_argument('-n', '--ip-start-num', type=int, nargs=None,
                        help='Client ip starts at X, usually used when adding clients to existing wg conf.', metavar='X')
    parser.add_argument('-u', '--pubkey', type=str, nargs=None,
                        help='Specify pubkey manually. Note when set, no gw node header config is generated.', metavar='PUBKEY')
    parser.add_argument('gw_addr', metavar='gw_public_addr', type=str,
                        help='The public address of gateway, that clients connect to.')
    parser.add_argument('addr', metavar='vpn_addr', type=str,
                        help='The network of the private subnet.')
    parser.add_argument('names', metavar='name', type=str, nargs='*',  default=None,
                        help='Name of clients. However, for single number, it\'ll be treated as number of clients. See examples.')
    return parser.parse_args()


args = _arg_parser()

# region functions


def gen_addresses(ADDR: _BaseNetwork, CLIENTS: List[str], ip_startnum=2) -> Tuple[_BaseAddress, List[_BaseAddress]]:
    gw_addr = ADDR.network_address+1
    c_addrs = []
    c_addrs_gen = range(ip_startnum, ADDR.num_addresses-2)
    if len(CLIENTS) > ADDR.num_addresses-3:
        raise MyError("Subnet capacity too small.")
    for (i, _) in zip(c_addrs_gen, CLIENTS):
        c_addrs.append(ADDR.network_address+i)
    return gw_addr, c_addrs


def gen_keypair() -> Tuple[str, str]:
    try:
        pri_raw = subprocess.check_output(['wg', 'genkey'])
        pri = pri_raw.decode('ascii').rstrip()
        pub = subprocess.check_output(['wg', 'pubkey'], input=pri_raw).decode('ascii').rstrip()
    except subprocess.CalledProcessError as e:
        raise MyError(f"Error while generating keys. {e}")
    return pri, pub


def gen_preshared() -> str:
    try:
        return subprocess.check_output(['wg', 'genpsk']).decode('ascii').rstrip()
    except subprocess.CalledProcessError as e:
        raise MyError(f"Error while generating preshared keys. {e}")


def gen_conf(port: int, gw_publicip: str, network: _BaseNetwork, c_names: List[str], write_to_file=False, append_mode_pubkey=None, ip_startnum=2) -> None:
    gw_addr, c_addrs = gen_addresses(network, c_names, ip_startnum)
    netmaskbits = network.max_prefixlen - round(math.log(network.num_addresses, 2))
    client_conf_files = []

    f = io.StringIO() if not write_to_file else open('/etc/wireguard/temp-0.conf', 'w')
    if not append_mode_pubkey:
        gw_private, gw_public = gen_keypair()
        print(f"""[Interface]
    #Table = off
    Address = {gw_addr}/{netmaskbits}
    ListenPort = {port}
    PrivateKey = {gw_private}""",
              file=f)
    else:
        gw_public = append_mode_pubkey
    for i, (c_name, c_addr) in enumerate(zip(c_names, c_addrs)):
        c_private, c_public = gen_keypair()
        psk = gen_preshared()
        print(f"""
[Peer]
# Client {c_name}
PublicKey = {c_public}
PresharedKey = {psk}
AllowedIPs = {c_addr}
""", file=f)
        g = io.StringIO() if not write_to_file else open(f'/etc/wireguard/temp-{c_name}.conf', 'w')
        client_conf_files.append(g)
        print(f"""[Interface]
#Table = off
Address = {c_addr}/{netmaskbits}
#ListenPort = {port}
PrivateKey = {c_private}

[Peer]
# Gateway
PublicKey = {gw_public}
PresharedKey = {psk}
AllowedIPs = {network}
Endpoint = {gw_publicip}:{port}
""", file=g)

    if not write_to_file:
        print("#################### Gateway File ####################")
        print(f.getvalue())
        print()
        for g in client_conf_files:
            print("#################### Client File ####################")
            print(g.getvalue())
            print()
    f.close()
    for g in client_conf_files:
        g.close()

# region main


def __main__():
    GW_PUBLICIP = args.gw_addr
    ADDR = ipaddress.ip_network(args.addr, strict=False)
    PORT = args.port if args.port else random.randint(10000, 65534)
    IP_START_NUM = args.ip_start_num if args.ip_start_num else 2
    CLIENTS = ([str(i-1+IP_START_NUM) for i in range(int(args.names[0]))]
               if len(args.names) == 1 and args.names[0].isdigit()
               else args.names
               )
    WRITE_TO_FILE = args.write_to_file
    GW_PUBKEY = args.pubkey

    if IP_START_NUM < 2:
        print(f"ip-start-num too small. It should be >=2, rather than {IP_START_NUM}.")
        exit(1)

    print(f"args: ADDR={ADDR}  PORT={PORT}  writetofile={WRITE_TO_FILE}  CLIENTS={','.join(CLIENTS)}")
    if IP_START_NUM or GW_PUBKEY:
        print(f"      GW_PUBKEY={GW_PUBKEY} IP_START_NUM={IP_START_NUM}")
    try:
        gen_conf(PORT, GW_PUBLICIP, ADDR, CLIENTS, WRITE_TO_FILE, GW_PUBKEY, IP_START_NUM)
    except MyError as e:
        print(e.args[0])


if __name__ == '__main__':
    __main__()