1 files changed, 118 insertions, 70 deletions
diff --git a/src/op_mode/wireguard.py b/src/op_mode/wireguard.py
index 66622c04c..e48da2e40 100755
--- a/src/op_mode/wireguard.py
+++ b/src/op_mode/wireguard.py
@@ -19,91 +19,139 @@
 import argparse
 import os
 import sys
+import shutil
 import subprocess
 import syslog as sl
 
+
 from vyos import ConfigError
 
 dir = r'/config/auth/wireguard'
-pk  = dir + '/private.key'
-pub = dir + '/public.key'
 psk = dir + '/preshared.key'
 
+
 def check_kmod():
-  """ check if kmod is loaded, if not load it """
-  if not os.path.exists('/sys/module/wireguard'):
-    sl.syslog(sl.LOG_NOTICE, "loading wirguard kmod") 
-    if  os.system('sudo modprobe wireguard') != 0:
-      sl.syslog(sl.LOG_ERR, "modprobe wireguard failed")
-      raise ConfigError("modprobe wireguard failed")
-
-def generate_keypair():
-  """ generates a keypair which is stored in /config/auth/wireguard """
-  ret = subprocess.call(['wg genkey | tee ' + pk + '|wg pubkey > ' + pub], shell=True)
-  if ret != 0:
-    raise ConfigError("wireguard key-pair generation failed")
-  else:
-    sl.syslog(sl.LOG_NOTICE, "new keypair wireguard key generated in " + dir)
-
-def genkey():
-  """ helper function to check, regenerate the keypair """
-  old_umask = os.umask(0o077)
-  if os.path.exists(pk) and os.path.exists(pub):
-    try:
-      choice = input("You already have a wireguard key-pair already, do you want to re-generate? [y/n] ")
-      if choice == 'y' or choice == 'Y':
-        generate_keypair()
-    except KeyboardInterrupt:
-        sys.exit(0)
-  else:
-    """ if keypair is bing executed from a running iso """
-    if not os.path.exists(dir):
-      os.umask(old_umask)
-      subprocess.call(['sudo mkdir -p ' + dir], shell=True)
-      subprocess.call(['sudo chgrp vyattacfg ' + dir], shell=True)
-      subprocess.call(['sudo chmod 770 ' + dir], shell=True)
-    generate_keypair()
-  os.umask(old_umask)
+    """ check if kmod is loaded, if not load it """
+    if not os.path.exists('/sys/module/wireguard'):
+        sl.syslog(sl.LOG_NOTICE, "loading wirguard kmod")
+        if os.system('sudo modprobe wireguard') != 0:
+            sl.syslog(sl.LOG_ERR, "modprobe wireguard failed")
+            raise ConfigError("modprobe wireguard failed")
 
-def showkey(key):
-  """ helper function to show privkey or pubkey """
-  if key == "pub":
-    if os.path.exists(pub):
-      print ( open(pub).read().strip() )
+
+def generate_keypair(pk, pub):
+    """ generates a keypair which is stored in /config/auth/wireguard """
+    old_umask = os.umask(0o027)
+    ret = subprocess.call(
+        ['wg genkey | tee ' + pk + '|wg pubkey > ' + pub], shell=True)
+    if ret != 0:
+        raise ConfigError("wireguard key-pair generation failed")
     else:
-      print("no public key found")
+        sl.syslog(
+            sl.LOG_NOTICE, "new keypair wireguard key generated in " + dir)
+    os.umask(old_umask)
 
-  if key == "pk":
-    if os.path.exists(pk):
-      print ( open(pk).read().strip() )
+
+def genkey(location):
+    """ helper function to check, regenerate the keypair """
+    pk = "{}/private.key".format(location)
+    pub = "{}/public.key".format(location)
+    old_umask = os.umask(0o027)
+    if os.path.exists(pk) and os.path.exists(pub):
+        try:
+            choice = input(
+                "You already have a wireguard key-pair, do you want to re-generate? [y/n] ")
+            if choice == 'y' or choice == 'Y':
+                generate_keypair(pk, pub)
+        except KeyboardInterrupt:
+            sys.exit(0)
+    else:
+        """ if keypair is bing executed from a running iso """
+        if not os.path.exists(location):
+            subprocess.call(['sudo mkdir -p ' + location], shell=True)
+            subprocess.call(['sudo chgrp vyattacfg ' + location], shell=True)
+            subprocess.call(['sudo chmod 750 ' + location], shell=True)
+        generate_keypair(pk, pub)
+    os.umask(old_umask)
+
+
+def showkey(key):
+    """ helper function to show privkey or pubkey """
+    if os.path.exists(key):
+        print (open(key).read().strip())
     else:
-      print("no private key found")
+        print ("{} not found".format(key))
+
 
 def genpsk():
-  """ generates a preshared key and shows it on stdout, it's stroed only in the config """
-  subprocess.call(['wg genpsk'], shell=True)
+    """
+        generates a preshared key and shows it on stdout,
+        it's stored only in the cli config
+    """
+
+    subprocess.call(['wg genpsk'], shell=True)
+
+
+def list_key_dirs():
+    """ lists all dirs under /config/auth/wireguard """
+    if os.path.exists(dir):
+        nks = next(os.walk(dir))[1]
+        for nk in nks:
+            print (nk)
+
+
+def del_key_dir(kname):
+    """ deletes /config/auth/wireguard/<kname> """
+    kdir = "{0}/{1}".format(dir, kname)
+    if not os.path.isdir(kdir):
+        print ("named keypair {} not found".format(kname))
+        return 1
+    shutil.rmtree(kdir)
+
 
 if __name__ == '__main__':
-  check_kmod()
-
-  parser = argparse.ArgumentParser(description='wireguard key management')
-  parser.add_argument('--genkey', action="store_true", help='generate key-pair')
-  parser.add_argument('--showpub', action="store_true", help='shows public key')
-  parser.add_argument('--showpriv', action="store_true", help='shows private key')
-  parser.add_argument('--genpsk', action="store_true", help='generates preshared-key')
-  args = parser.parse_args()
-
-  try:
-    if args.genkey:
-      genkey()
-    if args.showpub:
-      showkey("pub")
-    if args.showpriv:
-      showkey("pk")
-    if args.genpsk:
-      genpsk()
-
-  except ConfigError as e:
-    print(e)
-    sys.exit(1)
+    check_kmod()
+    parser = argparse.ArgumentParser(description='wireguard key management')
+    parser.add_argument(
+        '--genkey', action="store_true", help='generate key-pair')
+    parser.add_argument(
+        '--showpub', action="store_true", help='shows public key')
+    parser.add_argument(
+        '--showpriv', action="store_true", help='shows private key')
+    parser.add_argument(
+        '--genpsk', action="store_true", help='generates preshared-key')
+    parser.add_argument(
+        '--location', action="store", help='key location within {}'.format(dir))
+    parser.add_argument(
+        '--listkdir', action="store_true", help='lists named keydirectories')
+    parser.add_argument(
+        '--delkdir', action="store_true", help='removes named keydirectories')
+    args = parser.parse_args()
+
+    try:
+        if args.genkey:
+            if args.location:
+                genkey("{0}/{1}".format(dir, args.location))
+            else:
+                genkey(dir)
+
+        if args.showpub:
+            if args.location:
+                showkey("{0}/{1}/public.key".format(dir, args.location))
+            else:
+                showkey("{}/public.key".format(dir))
+        if args.showpriv:
+            if args.location:
+                showkey("{0}/{1}/private.key".format(dir, args.location))
+            else:
+                showkey("{}/private".format(dir))
+        if args.genpsk:
+            genpsk()
+        if args.listkdir:
+            list_key_dirs()
+        if args.delkdir:
+            del_key_dir(args.location)
 
+    except ConfigError as e:
+        print(e)
+        sys.exit(1)