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

2 files changed, 87 insertions, 87 deletions

diff --git a/op-mode-definitions/pki.xml.in b/op-mode-definitions/pki.xml.in
index a11814c8a..6b9b0d3f6 100644
--- a/op-mode-definitions/pki.xml.in
+++ b/op-mode-definitions/pki.xml.in
@@ -282,60 +282,66 @@
           </node>
           <node name="wireguard">
             <properties>
-              <help>Generate Wireguard keys</help>
+              <help>Generate WireGuard keys</help>
             </properties>
             <children>
               <node name="key-pair">
                 <properties>
-                  <help>Generate Wireguard key pair for use with server or peer</help>
+                  <help>Generate WireGuard public/private key-pair</help>
                 </properties>
                 <children>
-                  <tagNode name="file">
+                  <node name="install">
                     <properties>
-                      <help>Write generated Wireguard keys into the specified filename</help>
-                      <completionHelp>
-                        <list>&lt;filename&gt;</list>
-                      </completionHelp>
+                      <help>Generate CLI commands to install WireGuard key to configuration</help>
                     </properties>
-                    <command>sudo ${vyos_op_scripts_dir}/pki.py --action generate --wireguard --key "$6" --file</command>
-                  </tagNode>
-                  <tagNode name="install">
-                    <properties>
-                      <help>Commands for installing generated Wireguard key into running configuration</help>
-                      <completionHelp>
-                        <list>&lt;interface&gt; &lt;peer&gt;</list>
-                      </completionHelp>
-                    </properties>
-                    <command>sudo ${vyos_op_scripts_dir}/pki.py --action generate --wireguard --key "$6" --install</command>
-                  </tagNode>
+                    <children>
+                      <tagNode name="interface">
+                        <properties>
+                          <help>WireGuard Interface used in install command</help>
+                          <completionHelp>
+                            <path>interfaces wireguard</path>
+                          </completionHelp>
+                        </properties>
+                        <command>sudo ${vyos_op_scripts_dir}/pki.py --action generate --wireguard --key --interface "$7" --install</command>
+                      </tagNode>
+                    </children>
+                  </node>
                 </children>
-                <command>sudo ${vyos_op_scripts_dir}/pki.py --action generate --wireguard --key "noname"</command>
+                <command>sudo ${vyos_op_scripts_dir}/pki.py --action generate --wireguard --key</command>
               </node>
-              <node name="pre-shared-key">
+              <node name="preshared-key">
                 <properties>
-                  <help>Generate pre-shared key for use with a Wireguard peer</help>
+                  <help>Generate WireGuard pre-shared key</help>
                 </properties>
                 <children>
-                  <tagNode name="file">
+                  <node name="install">
                     <properties>
-                      <help>Write generated Wireguard PSK into the specified filename</help>
-                      <completionHelp>
-                        <list>&lt;filename&gt;</list>
-                      </completionHelp>
+                      <help>Generate CLI commands to install WireGuard key to configuration</help>
                     </properties>
-                    <command>sudo ${vyos_op_scripts_dir}/pki.py --action generate --wireguard --psk "$6" --file</command>
-                  </tagNode>
-                  <tagNode name="install">
-                    <properties>
-                      <help>Commands for installing generated Wireguard PSK on specified peer into running configuration</help>
-                      <completionHelp>
-                        <list>&lt;peer&gt;</list>
-                      </completionHelp>
-                    </properties>
-                    <command>sudo ${vyos_op_scripts_dir}/pki.py --action generate --wireguard --psk "$6" --install</command>
-                  </tagNode>
+                    <children>
+                      <tagNode name="interface">
+                        <properties>
+                          <help>WireGuard Interface used in install command</help>
+                          <completionHelp>
+                            <path>interfaces wireguard</path>
+                          </completionHelp>
+                        </properties>
+                        <children>
+                          <tagNode name="peer">
+                            <properties>
+                              <help>Interface used for install command</help>
+                              <completionHelp>
+                                <path>interfaces wireguard ${COMP_WORDS[COMP_CWORD-2]} peer</path>
+                              </completionHelp>
+                            </properties>
+                            <command>sudo ${vyos_op_scripts_dir}/pki.py --action generate --wireguard --psk --interface "$7" --peer "$9" --install</command>
+                          </tagNode>
+                        </children>
+                      </tagNode>
+                    </children>
+                  </node>
                 </children>
-                <command>sudo ${vyos_op_scripts_dir}/pki.py --action generate --wireguard --psk "noname"</command>
+                <command>sudo ${vyos_op_scripts_dir}/pki.py --action generate --wireguard --psk</command>
               </node>
             </children>
           </node>
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)