wlb: T114: Add firewall group support for WAN load balancer

3 files changed, 54 insertions, 1 deletions

diff --git a/src/conf_mode/load-balancing_wan.py b/src/conf_mode/load-balancing_wan.py
index dc7f1c838..8efb33f47 100755
--- a/src/conf_mode/load-balancing_wan.py
+++ b/src/conf_mode/load-balancing_wan.py
@@ -18,6 +18,7 @@ from sys import exit
 
 from vyos.config import Config
 from vyos.configdep import set_dependents, call_dependents
+from vyos.utils.dict import dict_search_args
 from vyos.utils.process import cmd
 from vyos import ConfigError
 from vyos import airbag
@@ -25,6 +26,13 @@ airbag.enable()
 
 service = 'vyos-wan-load-balance.service'
 
+valid_groups = [
+    'address_group',
+    'domain_group',
+    'network_group',
+    'port_group'
+]
+
 def get_config(config=None):
     if config:
         conf = config
@@ -38,6 +46,10 @@ def get_config(config=None):
                               get_first_key=True,
                               with_recursive_defaults=True)
 
+    if lb:
+        lb['firewall_group'] = conf.get_config_dict(['firewall', 'group'], key_mangling=('-', '_'), get_first_key=True,
+                                        no_tag_node_value_mangle=True)
+
     # prune limit key if not set by user
     for rule in lb.get('rule', []):
         if lb.from_defaults(['rule', rule, 'limit']):
@@ -89,6 +101,43 @@ def verify(lb):
 
             for direction in ['source', 'destination']:
                 if direction in rule_conf:
+                    side_conf = rule_conf[direction]
+
+                    if 'group' in side_conf:
+                        if len({'address_group', 'network_group', 'domain_group'} & set(side_conf['group'])) > 1:
+                            raise ConfigError('Only one address-group, network-group or domain-group can be specified')
+
+                        for group in valid_groups:
+                            if group in side_conf['group']:
+                                group_name = side_conf['group'][group]
+                                error_group = group.replace("_", "-")
+
+                                if group in ['address_group', 'network_group', 'domain_group']:
+                                    if 'address' in side_conf:
+                                        raise ConfigError(f'{error_group} and address cannot both be defined')
+
+                                if group in ['port_group']:
+                                    if 'port' in side_conf:
+                                        raise ConfigError(f'{error_group} and port cannot both be defined')
+
+                                if group_name and group_name[0] == '!':
+                                    group_name = group_name[1:]
+
+                                group_obj = dict_search_args(lb['firewall_group'], group, group_name)
+
+                                if group_obj is None:
+                                    raise ConfigError(f'Invalid {error_group} "{group_name}" on load-balancing wan rule')
+
+                                if not group_obj:
+                                    Warning(f'{error_group} "{group_name}" has no members!')
+
+                    if dict_search_args(side_conf, 'group', 'port_group'):
+                        if 'protocol' not in rule_conf:
+                            raise ConfigError('Protocol must be defined if specifying a port-group')
+
+                        if rule_conf['protocol'] not in ['tcp', 'udp', 'tcp_udp']:
+                            raise ConfigError('Protocol must be tcp, udp, or tcp_udp when specifying a port-group')
+
                     if 'port' in rule_conf[direction]:
                         if 'protocol' not in rule_conf:
                             raise ConfigError(f'Protocol required to specify port on load-balancing wan rule {rule_id}')
diff --git a/src/helpers/vyos-load-balancer.py b/src/helpers/vyos-load-balancer.py
index 6b260bf2c..5df83b31e 100755
--- a/src/helpers/vyos-load-balancer.py
+++ b/src/helpers/vyos-load-balancer.py
@@ -160,6 +160,9 @@ def get_config():
     lb = conf.get_config_dict(base, key_mangling=('-', '_'),
                             get_first_key=True, with_recursive_defaults=True)
 
+    lb['firewall_group'] = conf.get_config_dict(['firewall', 'group'], key_mangling=('-', '_'), get_first_key=True,
+                                    no_tag_node_value_mangle=True)
+
     # prune limit key if not set by user
     for rule in lb.get('rule', []):
         if lb.from_defaults(['rule', rule, 'limit']):
@@ -279,7 +282,7 @@ if __name__ == '__main__':
                         if state_changed and state['failure_count'] >= int(health_conf['failure_count']):
                             state['state'] = False
                             state['state_changed'] = True
-                    
+
                     #Force state changed to trigger the first write
                     if init == True:
                         state['state_changed'] = True
diff --git a/src/services/vyos-domain-resolver b/src/services/vyos-domain-resolver
index 17dae38e0..e1a52c93a 100755
--- a/src/services/vyos-domain-resolver
+++ b/src/services/vyos-domain-resolver
@@ -48,6 +48,7 @@ ipv4_tables = {
     'ip vyos_mangle',
     'ip vyos_filter',
     'ip vyos_nat',
+    'ip vyos_wanloadbalance',
     'ip raw'
 }