From d1bdf2b9d80d2e34b7370823d6f684102d7c9f4e Mon Sep 17 00:00:00 2001
From: sarthurdev <965089+sarthurdev@users.noreply.github.com>
Date: Sun, 5 Jun 2022 10:59:47 +0200
Subject: firewall: T970: Maintain a domain state to fallback if resolution
 fails

---
 python/vyos/firewall.py                  |  9 ++++-----
 src/conf_mode/firewall.py                |  3 ++-
 src/helpers/vyos-domain-group-resolve.py | 24 +++++++++++++++---------
 3 files changed, 21 insertions(+), 15 deletions(-)

diff --git a/python/vyos/firewall.py b/python/vyos/firewall.py
index 8a1237ca9..b962c4f18 100644
--- a/python/vyos/firewall.py
+++ b/python/vyos/firewall.py
@@ -26,21 +26,20 @@ def get_ips_domains_dict(list_domains):
     """
     Get list of IPv4 addresses by list of domains
     Ex: get_ips_domains_dict(['ex1.com', 'ex2.com'])
-        ['192.0.2.1', '192.0.2.2', '192.0.2.3']
+        {'ex1.com': ['192.0.2.1'], 'ex2.com': ['192.0.2.2', '192.0.2.3']}
     """
     from socket import gethostbyname_ex
     from socket import gaierror
 
-    ip_list = []
+    ip_dict = {}
     for domain in list_domains:
         try:
             _, _, ips = gethostbyname_ex(domain)
-            for entry in ips:
-                ip_list.append(entry)
+            ip_dict[domain] = ips
         except gaierror:
             pass
 
-    return ip_list
+    return ip_dict
 
 def nft_init_set(group_name, table="filter", family="ip"):
     """
diff --git a/src/conf_mode/firewall.py b/src/conf_mode/firewall.py
index 3c6aff386..335098bf1 100755
--- a/src/conf_mode/firewall.py
+++ b/src/conf_mode/firewall.py
@@ -427,7 +427,8 @@ def apply(firewall):
                     domains.append(address)
                 # Add elements to domain-group, try to resolve domain => ip
                 # and add elements to nft set
-                elements = get_ips_domains_dict(domains)
+                ip_dict = get_ips_domains_dict(domains)
+                elements = sum(ip_dict.values(), [])
                 nft_init_set(group)
                 nft_add_set_elements(group, elements)
         else:
diff --git a/src/helpers/vyos-domain-group-resolve.py b/src/helpers/vyos-domain-group-resolve.py
index ebb2057ec..e8501cfc6 100755
--- a/src/helpers/vyos-domain-group-resolve.py
+++ b/src/helpers/vyos-domain-group-resolve.py
@@ -28,10 +28,11 @@ from vyos.util import call
 
 base = ['firewall', 'group', 'domain-group']
 check_required = True
-count_failed = 0
+# count_failed = 0
 # Timeout in sec between checks
 timeout = 300
 
+domain_state = {}
 
 if __name__ == '__main__':
 
@@ -41,14 +42,19 @@ if __name__ == '__main__':
             domain_groups = config.get_config_dict(base, key_mangling=('-', '_'), get_first_key=True)
             for set_name, domain_config in domain_groups.items():
                 list_domains = domain_config['address']
-                elements = get_ips_domains_dict(list_domains)
+                elements = []
+                ip_dict = get_ips_domains_dict(list_domains)
+
+                for domain in list_domains:
+                    # Resolution succeeded, update domain state
+                    if domain in ip_dict:
+                        domain_state[domain] = ip_dict[domain]
+                        elements += ip_dict[domain]
+                    # Resolution failed, use previous domain state
+                    elif domain in domain_state:
+                        elements += domain_state[domain]
+
                 # Resolve successful
-                if bool(elements):
+                if elements:
                     nft_update_set_elements(set_name, elements)
-                    count_failed = 0
-                else:
-                    count_failed += 1
-                    # Domains not resolved 3 times by timeout
-                    if count_failed >= timeout * 3:
-                        nft_flush_set(set_name)
         time.sleep(timeout)
-- 
cgit v1.2.3