wireguard: T3642: improve "set" commands for generated key-pairs

1 files changed, 43 insertions, 49 deletions

diff --git a/src/op_mode/pki.py b/src/op_mode/pki.py
index 55330cbc2..d28cee5d0 100755
--- a/src/op_mode/pki.py
+++ b/src/op_mode/pki.py
@@ -202,30 +202,31 @@ def install_keypair(name, key_type, private_key=None, public_key=None, passphras
 
         if install_private_key:
             install_private_pem = "".join(private_key_pem.strip().split("\n")[1:-1])
-            print("set pki key-pair %s private key '%s'" % (name, install_private_pem))
+            print(f"set pki key-pair {name} private key '{install_private_pem}'")
             if passphrase:
-                print("set pki key-pair %s private password-protected" % (name,))
+                print(f"set pki key-pair {name} private password-protected")
         else:
             print("Private key:")
             print(private_key_pem)
 
-def install_wireguard_key(name, private_key, public_key):
+def install_wireguard_key(interface, private_key, public_key):
     # Show conf commands for installing wireguard key pairs
-    is_interface = re.match(r'^wg[\d]+$', name)
-
-    print("Configure mode commands to install key:")
-    if is_interface:
-        print("set interfaces wireguard %s private-key '%s'" % (name, private_key))
-        print("")
-        print("Public key for use on peer configuration: " + public_key)
-    else:
-        print("set interfaces wireguard [INTERFACE] peer %s public-key '%s'" % (name, public_key))
-        print("")
-        print("Private key for use on peer configuration: " + private_key)
-
-def install_wireguard_psk(name, psk):
+    from vyos.ifconfig import Section
+    if Section.section(interface) != 'wireguard':
+        print(f'"{interface}" is not a WireGuard interface name!')
+        exit(1)
+
+    print("Configure mode commands to install key:", end="\n\n")
+    print(f"set interfaces wireguard {interface} private-key '{private_key}'", end="\n\n")
+    print(f"Public key to use on peer system: '{public_key}'")
+
+def install_wireguard_psk(interface, peer, psk):
+    from vyos.ifconfig import Section
+    if Section.section(interface) != 'wireguard':
+        print(f'"{interface}" is not a WireGuard interface name!')
+        exit(1)
     # Show conf commands for installing wireguard psk
-    print("set interfaces wireguard [INTERFACE] peer %s preshared-key '%s'" % (name, psk))
+    print(f"set interfaces wireguard {interface} peer {peer} preshared-key '{psk}'")
 
 def ask_passphrase():
     passphrase = None
@@ -632,48 +633,37 @@ def generate_openvpn_key(name, install=False, file=False):
         key_data = "".join(key_lines[1:-1]) # Remove wrapper tags and line endings
         key_version = '1'
 
+        import re
         version_search = re.search(r'BEGIN OpenVPN Static key V(\d+)', result) # Future-proofing (hopefully)
         if version_search:
             key_version = version_search[1]
 
+        base = f"set pki openvpn shared-secret {name}"
         print("Configure mode commands to install OpenVPN key:")
-        print("set pki openvpn shared-secret %s key '%s'" % (name, key_data))
-        print("set pki openvpn shared-secret %s version '%s'" % (name, key_version))
+        print(f"{base} key '{key_data}'")
+        print(f"{base} version '{key_version}'")
 
     if file:
         write_file(f'{name}.key', result)
 
-def generate_wireguard_key(name, install=False, file=False):
+def generate_wireguard_key(interface=None, install=False):
     private_key = cmd('wg genkey')
     public_key = cmd('wg pubkey', input=private_key)
 
-    if not install:
-        print("Private key: " + private_key)
-        print("Public key: " + public_key)
-
-    if install:
-        install_wireguard_key(name, private_key, public_key)
-
-    if file:
-        write_file(f'{name}_public.key', public_key)
-        write_file(f'{name}_private.key', private_key)
+    if interface and install:
+        install_wireguard_key(interface, private_key, public_key)
+    else:
+        print(f'Private key: {private_key}')
+        print(f'Public key: {public_key}', end='\n\n')
 
-def generate_wireguard_psk(name, install=False, file=False):
+def generate_wireguard_psk(interface=None, peer=None, install=False):
     psk = cmd('wg genpsk')
-
-    if not install and not file:
-        print("Pre-shared key:")
-        print(psk)
-        return None
-
-    if install:
-        install_wireguard_psk(name, psk)
-
-    if file:
-        write_file(f'{name}.key', psk)
+    if interface and peer and install:
+        install_wireguard_psk(interface, peer, psk)
+    else:
+        print(f'Pre-shared key: {psk}')
 
 # Show functions
-
 def show_certificate_authority(name=None):
     headers = ['Name', 'Subject', 'Issuer CN', 'Issued', 'Expiry', 'Private Key', 'Parent']
     data = []
@@ -790,10 +780,13 @@ if __name__ == '__main__':
     # OpenVPN
     parser.add_argument('--openvpn', help='OpenVPN TLS key', required=False)
 
-    # Wireguard
+    # WireGuard
     parser.add_argument('--wireguard', help='Wireguard', action='store_true')
-    parser.add_argument('--key', help='Wireguard key pair', required=False)
-    parser.add_argument('--psk', help='Wireguard pre shared key', required=False)
+    group = parser.add_mutually_exclusive_group()
+    group.add_argument('--key', help='Wireguard key pair', action='store_true', required=False)
+    group.add_argument('--psk', help='Wireguard pre shared key', action='store_true', required=False)
+    parser.add_argument('--interface', help='Install generated keys into running-config for named interface', action='store')
+    parser.add_argument('--peer', help='Install generated keys into running-config for peer', action='store')
 
     # Global
     parser.add_argument('--file', help='Write generated keys into specified filename', action='store_true')
@@ -833,9 +826,10 @@ if __name__ == '__main__':
 
             elif args.wireguard:
                 if args.key:
-                    generate_wireguard_key(args.key, install=args.install, file=args.file)
-                elif args.psk:
-                    generate_wireguard_psk(args.psk, install=args.install, file=args.file)
+                    generate_wireguard_key(args.interface, install=args.install)
+                if args.psk:
+                    generate_wireguard_psk(args.interface, peer=args.peer, install=args.install)
+
         elif args.action == 'show':
             if args.ca:
                 show_certificate_authority(None if args.ca == 'all' else args.ca)