wgtables: a configuration tool for wireguard

wgtables

#!/usr/bin/python3
import os, sys, re

#---------------------------------------------------------------------------
#Thank you to https://gist.github.com/dfee/6ed3a4b05cfe7a6faf40a2102408d5d8|
IPV4SEG  = r'(?:25[0-5]|(?:2[0-4]|1{0,1}[0-9]){0,1}[0-9])'                #|
IPV4ADDR = r'(?:(?:' + IPV4SEG + r'\.){3,3}' + IPV4SEG + r')'             #|
IPV6SEG  = r'(?:(?:[0-9a-fA-F]){1,4})'                                    #|
IPV6GROUPS = (                                                            #|
    r'(?:' + IPV6SEG + r':){7,7}' + IPV6SEG,                              #|
    r'(?:' + IPV6SEG + r':){1,7}:',                                       #|
    r'(?:' + IPV6SEG + r':){1,6}:' + IPV6SEG,                             #|
    r'(?:' + IPV6SEG + r':){1,5}(?::' + IPV6SEG + r'){1,2}',              #|
    r'(?:' + IPV6SEG + r':){1,4}(?::' + IPV6SEG + r'){1,3}',              #|
    r'(?:' + IPV6SEG + r':){1,3}(?::' + IPV6SEG + r'){1,4}',              #|
    r'(?:' + IPV6SEG + r':){1,2}(?::' + IPV6SEG + r'){1,5}',              #|
    IPV6SEG + r':(?:(?::' + IPV6SEG + r'){1,6})',                         #|
    r':(?:(?::' + IPV6SEG + r'){1,7}|:)',                                 #|
    r'fe80:(?::' + IPV6SEG + r'){0,4}%[0-9a-zA-Z]{1,}',                   #|
    r'::(?:ffff(?::0{1,4}){0,1}:){0,1}[^\s:]' + IPV4ADDR,                 #|
    r'(?:' + IPV6SEG + r':){1,4}:[^\s:]' + IPV4ADDR,                      #|
)                                                                         #|
IPV6ADDR = '|'.join(['(?:{})'.format(g) for g in IPV6GROUPS[::-1]])       #|
#Thank you to https://gist.github.com/dfee/6ed3a4b05cfe7a6faf40a2102408d5d8|
#---------------------------------------------------------------------------

def main(argv):
    usage = "Usage: wgtables <interface> <up | down> <-4 | -6> [ -gw <gateway> ] [-a <address>...] <address>"
    if len(argv) < 5:
        print(usage)
        sys.exit(2)
    addresses = []
    for i, arg in enumerate(argv):
        if arg == "-a":
            addresses.append(argv[i + 1])
    addresses.append(argv[-1])
    result = 0
    if argv[2] == "up":
        for i, address in enumerate(addresses):
            if i == 0:
                result += exec("add", argv, address, doroute=True)
            else:
                result += exec("add", argv, address)
        sys.exit(result)
    elif argv[2] == "down":
        for i, address in enumerate(addresses):
            if i == 0:
                result += exec("del", argv, address, doroute=True)
            else:
                result += exec("del", argv, address)
        sys.exit(result)
    else:
        print(usage)
        sys.exit(2)

def exec(type, argv, address, doroute=False):
    if "-4" in argv and re.match(IPV4ADDR, address):
        iptype = ''
    elif "-6" in argv and re.match(IPV6ADDR, address):
        iptype = '-6 '
    else:
        print("Cannot identify address")
        sys.exit(2)
    if "-gw" in argv:
        ip = argv[argv.index('-gw') + 1]
        gateway4 = ip if re.match(IPV4ADDR, ip) else None
        gateway6 = ip if re.match(IPV6ADDR, ip) else None
    else:
        gateway4 = None
        gateway6 = None
    subnet4 = address.replace(address.split('.')[-1], '0') + address[-3:]
    subnet6 = address.replace(address.split('::')[-1], '') + address[-3:]
    addressraw = re.match((IPV6ADDR if iptype else IPV4ADDR), address).group(0)
    addressend = addressraw + ('/128' if re.match(IPV6ADDR, addressraw) else '/32')
    gateway4 = address.replace(address.split('.')[-1], '1') if gateway4 == None else gateway4
    gateway6 = address.replace(address.split('::')[-1], '1') if gateway6 == None else gateway6
    interface = argv[1]
    command1 = f"ip {iptype}route {type} {subnet6 if iptype else subnet4} dev {interface} src {addressraw} table {interface}"
    command2 = f"ip {iptype}route {type} default via {gateway6 if iptype else gateway4} dev {interface} table {interface}"
    command3 = f"ip {iptype}rule {type} from {addressend} table {interface}"
    command4 = f"ip {iptype}rule {type} to {addressend} table {interface}"
    if doroute:
        print(command1 + '\n' + command2 + '\n' + command3 + '\n' + command4)
        err1 = os.system(command1)
        err2 = os.system(command2)
        err3 = os.system(command3)
        err4 = os.system(command4)
    else:
        print(command3 + '\n' + command4)
        err1 = 0
        err2 = 0
        err3 = os.system(command3)
        err4 = os.system(command4)
    return (err1 + err2 + err3 + err4)

if __name__ == '__main__':
    main(sys.argv)

UPDATE: added gateway flag

allows you to set custom gateway and fixes a few bugs

usage:

wgtables %i up -4 -gw 10.0.0.3 10.0.0.4/32
wgtables %i up -6 -gw 2001:db8::3 2001:db8::4/32

UPDATE: allowed user to use multiple addresses

You can now use multiple addresses and it will use python3 by default

usage:

wgtables %i up -4 -a 10.0.0.4/32 -a 10.0.0.5/32 -a 10.0.0.3/32 10.0.0.2/32
wgtables %i up -6 -a 2001:db8::3/32 -a 2001:db8::4/32 -a 2001:db8::5/32 2001:db8::2/32

Must always have 1 address without -a option at the very end, compatible with -gw option