dnstools: handle multiple strings in record value

modoboa/dnstools/lib.py

@@ -9,42 +9,36 @@ from modoboa.admin import lib as admin_lib
 from . import constants
 
 
-def get_spf_record(domain):
-    """Return SPF record for domain (if any)."""
-    records = admin_lib.get_dns_records(domain, "TXT")
+def _get_record_type_value(records, rr_type):
     if records is None:
         return None
     for record in records:
-        value = str(record).strip('"')
-        if value.startswith("v=spf1"):
+        # Multiple strings are separated by a space as described in:
+        # https://tools.ietf.org/html/rfc4408#section-3.1.3
+        value = str(record).replace('" "', '').strip('"')
+        if value.startswith('v={}'.format(rr_type)):
             return value
     return None
 
 
+def get_spf_record(domain):
+    """Return SPF record for domain (if any)."""
+    records = admin_lib.get_dns_records(domain, "TXT")
+    return _get_record_type_value(records, 'spf1')
+
+
 def get_dkim_record(domain, selector):
     """Return DKIM records form domain (if any)."""
     name = "{}._domainkey.{}".format(selector, domain)
     records = admin_lib.get_dns_records(name, "TXT")
-    if records is None:
-        return None
-    for record in records:
-        value = str(record).strip('"')
-        if value.startswith("v=DKIM1"):
-            return value
-    return None
+    return _get_record_type_value(records, 'DKIM1')
 
 
 def get_dmarc_record(domain):
     """Return DMARC record for domain (if any)."""
     name = "_dmarc.{}".format(domain)
     records = admin_lib.get_dns_records(name, "TXT")
-    if records is None:
-        return None
-    for record in records:
-        value = str(record).strip('"')
-        if value.startswith("v=DMARC1"):
-            return value
-    return None
+    return _get_record_type_value(records, 'DMARC1')
 
 
 def _get_simple_record(name):

modoboa/dnstools/tests.py

@@ -1,8 +1,12 @@
 """App related tests."""
 
+from unittest import mock
+
 from django.urls import reverse
 from django.test import SimpleTestCase
 
+from dns.rdtypes.ANY.TXT import TXT
+
 from modoboa.lib.tests import ModoTestCase
 
 from . import factories
@@ -78,8 +82,25 @@ BAD_DKIM_RECORDS = [
 ]
 
 
-class TestCase(SimpleTestCase):
-    """Test test library."""
+class LibTestCase(SimpleTestCase):
+    """TestCase for library methods."""
+
+    @mock.patch('modoboa.admin.lib.get_dns_records')
+    def test_get_record_type_value(self, mock_get_dns_records):
+        mock_get_dns_records.return_value = [
+            TXT("IN", "TXT", ["v=spf1 mx -all"]),
+            TXT("IN", "TXT", ["v=DKIM1 p=XXXXX", "YYYYY"]),
+            TXT("IN", "TXT", ["v=DMARC1 p=reject"]),
+        ]
+        self.assertEqual(
+            lib.get_spf_record("example.com"), "v=spf1 mx -all"
+        )
+        self.assertEqual(
+            lib.get_dkim_record("example.com", "mail"), "v=DKIM1 p=XXXXXYYYYY"
+        )
+        self.assertEqual(
+            lib.get_dmarc_record("example.com"), "v=DMARC1 p=reject"
+        )
 
     def test_check_spf_syntax(self):
         for record in BAD_SPF_RECORDS: