Merge pull request #4468 from sarthurdev/T5707

wireguard: T7387: Optimise wireguard peer handling

2 files changed, 115 insertions, 60 deletions

diff --git a/python/vyos/ifconfig/wireguard.py b/python/vyos/ifconfig/wireguard.py
index f5217aecb..3a28723b3 100644
--- a/python/vyos/ifconfig/wireguard.py
+++ b/python/vyos/ifconfig/wireguard.py
@@ -22,12 +22,13 @@ from tempfile import NamedTemporaryFile
 from hurry.filesize import size
 from hurry.filesize import alternative
 
+from vyos.base import Warning
 from vyos.configquery import ConfigTreeQuery
 from vyos.ifconfig import Interface
 from vyos.ifconfig import Operational
 from vyos.template import is_ipv6
 from vyos.template import is_ipv4
-
+from vyos.utils.network import get_wireguard_peers
 class WireGuardOperational(Operational):
     def _dump(self):
         """Dump wireguard data in a python friendly way."""
@@ -251,92 +252,131 @@ class WireGuardIf(Interface):
         """Get a synthetic MAC address."""
         return self.get_mac_synthetic()
 
+    def get_peer_public_keys(self, config, disabled=False):
+        """Get list of configured peer public keys"""
+        if 'peer' not in config:
+            return []
+
+        public_keys = []
+
+        for _, peer_config in config['peer'].items():
+            if disabled == ('disable' in peer_config):
+                public_keys.append(peer_config['public_key'])
+
+        return public_keys
+
     def update(self, config):
         """General helper function which works on a dictionary retrived by
         get_config_dict(). It's main intention is to consolidate the scattered
         interface setup code and provide a single point of entry when workin
         on any interface."""
-        tmp_file = NamedTemporaryFile('w')
-        tmp_file.write(config['private_key'])
-        tmp_file.flush()
 
         # Wireguard base command is identical for every peer
         base_cmd = f'wg set {self.ifname}'
+
         interface_cmd = base_cmd
         if 'port' in config:
             interface_cmd += ' listen-port {port}'
         if 'fwmark' in config:
             interface_cmd += ' fwmark {fwmark}'
 
-        interface_cmd += f' private-key {tmp_file.name}'
-        interface_cmd = interface_cmd.format(**config)
-        # T6490: execute command to ensure interface configured
-        self._cmd(interface_cmd)
+        with NamedTemporaryFile('w') as tmp_file:
+            tmp_file.write(config['private_key'])
+            tmp_file.flush()
 
-        # If no PSK is given remove it by using /dev/null - passing keys via
-        # the shell (usually bash) is considered insecure, thus we use a file
-        no_psk_file = '/dev/null'
+            interface_cmd += f' private-key {tmp_file.name}'
+            interface_cmd = interface_cmd.format(**config)
+            # T6490: execute command to ensure interface configured
+            self._cmd(interface_cmd)
+
+        current_peer_public_keys = get_wireguard_peers(self.ifname)
+
+        if 'rebuild_required' in config:
+            # Remove all existing peers that no longer exist in config
+            current_public_keys = self.get_peer_public_keys(config)
+            cmd_remove_peers = [f' peer {public_key} remove'
+                                for public_key in current_peer_public_keys
+                                if public_key not in current_public_keys]
+            if cmd_remove_peers:
+                self._cmd(base_cmd + ''.join(cmd_remove_peers))
 
         if 'peer' in config:
+            # Group removal of disabled peers in one command
+            current_disabled_peers = self.get_peer_public_keys(config, disabled=True)
+            cmd_disabled_peers = [f' peer {public_key} remove'
+                                  for public_key in current_disabled_peers]
+            if cmd_disabled_peers:
+                self._cmd(base_cmd + ''.join(cmd_disabled_peers))
+
+            peer_cmds = []
+            peer_domain_cmds = []
+            peer_psk_files = []
+
             for peer, peer_config in config['peer'].items():
                 # T4702: No need to configure this peer when it was explicitly
                 # marked as disabled - also active sessions are terminated as
                 # the public key was already removed when entering this method!
                 if 'disable' in peer_config:
-                    # remove peer if disabled, no error report even if peer not exists
-                    cmd = base_cmd + ' peer {public_key} remove'
-                    self._cmd(cmd.format(**peer_config))
                     continue
 
-                psk_file = no_psk_file
-
                 # start of with a fresh 'wg' command
-                peer_cmd = base_cmd + ' peer {public_key}'
+                peer_cmd = ' peer {public_key}'
 
-                try:
-                    cmd = peer_cmd
-
-                    if 'preshared_key' in peer_config:
-                        psk_file = '/tmp/tmp.wireguard.psk'
-                        with open(psk_file, 'w') as f:
-                            f.write(peer_config['preshared_key'])
-                    cmd += f' preshared-key {psk_file}'
-
-                    # Persistent keepalive is optional
-                    if 'persistent_keepalive' in peer_config:
-                        cmd += ' persistent-keepalive {persistent_keepalive}'
-
-                    # Multiple allowed-ip ranges can be defined - ensure we are always
-                    # dealing with a list
-                    if isinstance(peer_config['allowed_ips'], str):
-                        peer_config['allowed_ips'] = [peer_config['allowed_ips']]
-                    cmd += ' allowed-ips ' + ','.join(peer_config['allowed_ips'])
-
-                    self._cmd(cmd.format(**peer_config))
-
-                    cmd = peer_cmd
-
-                    # Ensure peer is created even if dns not working
-                    if {'address', 'port'} <= set(peer_config):
-                        if is_ipv6(peer_config['address']):
-                            cmd += ' endpoint [{address}]:{port}'
-                        elif is_ipv4(peer_config['address']):
-                            cmd += ' endpoint {address}:{port}'
-                        else:
-                            # don't set endpoint if address uses domain name
-                            continue
-                    elif {'host_name', 'port'} <= set(peer_config):
-                        cmd += ' endpoint {host_name}:{port}'
-
-                    self._cmd(cmd.format(**peer_config), env={
+                cmd = peer_cmd
+
+                if 'preshared_key' in peer_config:
+                    with NamedTemporaryFile(mode='w', delete=False) as tmp_file:
+                        tmp_file.write(peer_config['preshared_key'])
+                        tmp_file.flush()
+                        cmd += f' preshared-key {tmp_file.name}'
+                        peer_psk_files.append(tmp_file.name)
+                else:
+                    # If no PSK is given remove it by using /dev/null - passing keys via
+                    # the shell (usually bash) is considered insecure, thus we use a file
+                    cmd += f' preshared-key /dev/null'
+
+                # Persistent keepalive is optional
+                if 'persistent_keepalive' in peer_config:
+                    cmd += ' persistent-keepalive {persistent_keepalive}'
+
+                # Multiple allowed-ip ranges can be defined - ensure we are always
+                # dealing with a list
+                if isinstance(peer_config['allowed_ips'], str):
+                    peer_config['allowed_ips'] = [peer_config['allowed_ips']]
+                cmd += ' allowed-ips ' + ','.join(peer_config['allowed_ips'])
+
+                peer_cmds.append(cmd.format(**peer_config))
+
+                cmd = peer_cmd
+
+                # Ensure peer is created even if dns not working
+                if {'address', 'port'} <= set(peer_config):
+                    if is_ipv6(peer_config['address']):
+                        cmd += ' endpoint [{address}]:{port}'
+                    elif is_ipv4(peer_config['address']):
+                        cmd += ' endpoint {address}:{port}'
+                    else:
+                        # don't set endpoint if address uses domain name
+                        continue
+                elif {'host_name', 'port'} <= set(peer_config):
+                    cmd += ' endpoint {host_name}:{port}'
+                else:
+                    continue
+
+                peer_domain_cmds.append(cmd.format(**peer_config))
+
+            try:
+                if peer_cmds:
+                    self._cmd(base_cmd + ''.join(peer_cmds))
+
+                if peer_domain_cmds:
+                    self._cmd(base_cmd + ''.join(peer_domain_cmds), env={
                         'WG_ENDPOINT_RESOLUTION_RETRIES': config['max_dns_retry']})
-                except:
-                    # todo: logging
-                    pass
-                finally:
-                    # PSK key file is not required to be stored persistently as its backed by CLI
-                    if psk_file != no_psk_file and os.path.exists(psk_file):
-                        os.remove(psk_file)
+            except Exception as e:
+                Warning(f'Failed to apply Wireguard peers on {self.ifname}: {e}')
+            finally:
+                for tmp in peer_psk_files:
+                    os.unlink(tmp)
 
         # call base class
         super().update(config)
diff --git a/python/vyos/utils/network.py b/python/vyos/utils/network.py
index 20b6a3c9e..0a84be478 100644
--- a/python/vyos/utils/network.py
+++ b/python/vyos/utils/network.py
@@ -416,6 +416,21 @@ def is_wireguard_key_pair(private_key: str, public_key:str) -> bool:
     else:
         return False
 
+def get_wireguard_peers(ifname: str) -> list:
+    """
+    Return list of configured Wireguard peers for interface
+    :param ifname: Interface name
+    :type ifname: str
+    :return: list of public keys
+    :rtype: list
+    """
+    if not interface_exists(ifname):
+        return []
+
+    from vyos.utils.process import cmd
+    peers = cmd(f'wg show {ifname} peers')
+    return peers.splitlines()
+
 def is_subnet_connected(subnet, primary=False):
     """
     Verify is the given IPv4/IPv6 subnet is connected to any interface on this