From eaf268aea9b2a55ffcceddecb5b80fa8c5c267c3 Mon Sep 17 00:00:00 2001
From: Franziska Kunsmann <hi@kunsmann.eu>
Date: Mon, 9 Nov 2020 18:46:37 +0100
Subject: [PATCH] libs/tools: change resolve_identifier() to return ipv4 and
 ipv6 separately

---
 bundles/powerdns/metadata.py | 34 +++++++++++-----------------------
 bundles/rspamd/metadata.py   |  7 ++++---
 libs/tools.py                | 15 +++++++++++++--
 3 files changed, 28 insertions(+), 28 deletions(-)

diff --git a/bundles/powerdns/metadata.py b/bundles/powerdns/metadata.py
index 006b0cf..4d4f92f 100644
--- a/bundles/powerdns/metadata.py
+++ b/bundles/powerdns/metadata.py
@@ -35,9 +35,8 @@ def get_ips_of_secondary_nameservers(metadata):
     ips = set()
     for rnode in repo.nodes_in_group('dns'):
         if rnode.metadata.get('powerdns/is_secondary', False):
-            ips.update({
-                str(ip) for ip in repo.libs.tools.resolve_identifier(repo, rnode.name)
-            })
+            for identifier, found_ips in repo.libs.tools.resolve_identifier(repo, rnode.name).items():
+                ips.update({str(ip) for ip in found_ips})
 
     return {
         'powerdns': {
@@ -53,9 +52,8 @@ def get_ips_of_primary_nameservers(metadata):
     ips = set()
     for rnode in repo.nodes_in_group('dns'):
         if not rnode.metadata.get('powerdns/is_secondary', False):
-            ips.update({
-                str(ip) for ip in repo.libs.tools.resolve_identifier(repo, rnode.name)
-            })
+            for identifier, found_ips in repo.libs.tools.resolve_identifier(repo, rnode.name).items():
+                ips.update({str(ip) for ip in found_ips})
 
     return {
         'powerdns': {
@@ -75,29 +73,19 @@ def generate_dns_entries_for_nodes(metadata):
         ip4 = None
         ip6 = None
 
-        ips = repo.libs.tools.resolve_identifier(repo, rnode.name)
-        for ip in ips:
-            if (
-                not ip4 and
-                not ip.is_private and
-                '.' in str(ip) # poor-mans 'is this ipv4' detection
-            ):
+        found_ips = repo.libs.tools.resolve_identifier(repo, rnode.name)
+        for ip in sorted(found_ips['ipv4']):
+            if not ip4 and not ip.is_private:
                 ip4 = ip
 
-            if (
-                not ip6 and
-                not ip.is_private and
-                ':' in str(ip)
-            ):
+        for ip in sorted(found_ips['ipv6']):
+            if not ip6 and not ip.is_private:
                 ip6 = ip
 
         # We're doing this once again to get the nodes which only have
         # private ips.
-        if not ip4:
-            for ip in ips:
-                if '.' in str(ip):
-                    ip4 = ip
-                    break
+        if not ip4 and len(found_ips['ipv4']):
+            ip4 = sorted(found_ips['ipv4'])[0]
 
         if ip4:
             results.add('{} IN A {}'.format(dns_name, ip4))
diff --git a/bundles/rspamd/metadata.py b/bundles/rspamd/metadata.py
index 65cf06a..4a0e9b5 100644
--- a/bundles/rspamd/metadata.py
+++ b/bundles/rspamd/metadata.py
@@ -36,9 +36,10 @@ def populate_permitted_ips_list_with_ips_from_repo(metadata):
     ips = set()
 
     for rnode in repo.nodes:
-        for ip in repo.libs.tools.resolve_identifier(repo, rnode.name):
-            if not ip.is_private:
-                ips.add(str(ip))
+        for identifier, found_ips in repo.libs.tools.resolve_identifier(repo, rnode.name).items():
+            for ip in found_ips:
+                if not ip.is_private:
+                    ips.add(str(ip))
 
     return {
         'rspamd': {
diff --git a/libs/tools.py b/libs/tools.py
index 3f42c74..6bfec8c 100644
--- a/libs/tools.py
+++ b/libs/tools.py
@@ -1,5 +1,5 @@
 from bundlewrap.exceptions import NoSuchGroup, NoSuchNode
-from ipaddress import ip_address
+from ipaddress import ip_address, IPv4Address
 
 def resolve_identifier(repo, identifier):
     """
@@ -29,4 +29,15 @@ def resolve_identifier(repo, identifier):
         if node.metadata.get('external_ipv4', None):
             found_ips.add(ip_address(node.metadata.get('external_ipv4')))
 
-    return found_ips
+    ip_dict = {
+        'ipv4': set(),
+        'ipv6': set(),
+    }
+
+    for ip in found_ips:
+        if isinstance(ip, IPv4Address):
+            ip_dict['ipv4'].add(ip)
+        else:
+            ip_dict['ipv6'].add(ip)
+
+    return ip_dict