From dfc0760adf8a966c11e8f87ff8e10faf87787065 Mon Sep 17 00:00:00 2001
From: Samuel Parkinson <samuel.parkinson@ft.com>
Date: Wed, 15 Feb 2023 16:45:50 +0000
Subject: [PATCH] Count extra lookups for the `include` mechanism

Co-authored-by: Jon Nangle <jon.nangle@ft.com>
---
 octodns/processor/spf.py            | 70 +++++++++++++++++++----------
 tests/test_octodns_processor_spf.py | 28 ++++++++++--
 2 files changed, 71 insertions(+), 27 deletions(-)

diff --git a/octodns/processor/spf.py b/octodns/processor/spf.py
index 258ad76..947d651 100644
--- a/octodns/processor/spf.py
+++ b/octodns/processor/spf.py
@@ -4,6 +4,10 @@
 
 from logging import getLogger
 
+import dns.resolver
+
+from octodns.record.base import Record
+
 from .base import BaseProcessor, ProcessorException
 
 
@@ -22,37 +26,55 @@ class SpfDnsLookupProcessor(BaseProcessor):
         self.log.debug(f"SpfDnsLookupProcessor: {name}")
         super().__init__(name)
 
-    def process_source_zone(self, zone, *args, **kwargs):
-        for record in zone.records:
-            if record._type != 'TXT':
-                continue
+    def _lookup(
+        self, record: Record, values: list[str], lookups: int = 0
+    ) -> int:
+        # SPF values must begin with 'v=spf1 '
+        spf = [value for value in values if value.startswith('v=spf1 ')]
 
-            if record._octodns.get('lenient'):
-                continue
+        if len(spf) == 0:
+            return lookups
 
-            # SPF values must begin with 'v=spf1 '
-            values = [
-                value for value in record.values if value.startswith('v=spf1 ')
-            ]
+        if len(spf) > 1:
+            raise SpfValueException(
+                f"{record.fqdn} has more than one SPF value"
+            )
 
-            if len(values) == 0:
-                continue
+        spf = spf[0]
 
-            if len(values) > 1:
-                raise SpfValueException(
-                    f"{record.fqdn} has more than one SPF value"
+        terms = spf.removeprefix('v=spf1 ').split(' ')
+
+        for term in terms:
+            if lookups > 10:
+                raise SpfDnsLookupException(
+                    f"{record.fqdn} exceeds the 10 DNS lookup limit in the SPF record"
                 )
 
-            lookups = 0
-            terms = values[0].removeprefix('v=spf1 ').split(' ')
+            # These mechanisms cost one DNS lookup each
+            if term.startswith(
+                ('a', 'mx', 'exists:', 'redirect', 'include:', 'ptr')
+            ):
+                lookups += 1
 
-            for term in terms:
-                if lookups > 10:
-                    raise SpfDnsLookupException(
-                        f"{record.fqdn} has too many SPF DNS lookups"
-                    )
+            # The include mechanism can result in further lookups after resolving the DNS record
+            if term.startswith('include:'):
+                answer = dns.resolver.resolve(
+                    term.removeprefix('include:'), 'TXT'
+                )
+                lookups = self._lookup(
+                    record, [value.to_text()[1:-1] for value in answer], lookups
+                )
+
+        return lookups
+
+    def process_source_zone(self, zone, *args, **kwargs):
+        for record in zone.records:
+            if record._type != 'TXT':
+                continue
+
+            if record._octodns.get('lenient'):
+                continue
 
-                if term in ['a', 'mx', 'exists', 'redirect']:
-                    lookups += 1
+            self._lookup(record, record.values)
 
         return zone
diff --git a/tests/test_octodns_processor_spf.py b/tests/test_octodns_processor_spf.py
index b4cd3a1..2388242 100644
--- a/tests/test_octodns_processor_spf.py
+++ b/tests/test_octodns_processor_spf.py
@@ -23,7 +23,10 @@ class TestSpfDnsLookupProcessor(TestCase):
                 {
                     'type': 'TXT',
                     'ttl': 86400,
-                    'values': ['v=spf1 a ~all', 'v=DMARC1\; p=reject\;'],
+                    'values': [
+                        'v=spf1 a include:_spf.google.com ~all',
+                        'v=DMARC1\; p=reject\;',
+                    ],
                 },
             )
         )
@@ -38,7 +41,7 @@ class TestSpfDnsLookupProcessor(TestCase):
                     'type': 'TXT',
                     'ttl': 86400,
                     'values': [
-                        'v=spf1 a a a a a a a a a a -all',
+                        'v=spf1 a ip4:1.2.3.4 ip4:1.2.3.4 ip4:1.2.3.4 ip4:1.2.3.4 ip4:1.2.3.4 ip4:1.2.3.4 ip4:1.2.3.4 ip4:1.2.3.4 ip4:1.2.3.4 ip4:1.2.3.4 ip4:1.2.3.4 -all',
                         'v=DMARC1\; p=reject\;',
                     ],
                 },
@@ -56,7 +59,26 @@ class TestSpfDnsLookupProcessor(TestCase):
                     'type': 'TXT',
                     'ttl': 86400,
                     'values': [
-                        'v=spf1 a mx exists redirect a a a a a a a ~all',
+                        'v=spf1 a mx exists:example.com a a a a a a a a ~all',
+                        'v=DMARC1\; p=reject\;',
+                    ],
+                },
+            )
+        )
+
+        with self.assertRaises(SpfDnsLookupException):
+            processor.process_source_zone(zone)
+
+        zone = Zone('unit.tests.', [])
+        zone.add_record(
+            Record.new(
+                zone,
+                '',
+                {
+                    'type': 'TXT',
+                    'ttl': 86400,
+                    'values': [
+                        'v=spf1 include:example.com include:_spf.google.com include:_spf.google.com include:_spf.google.com ~all',
                         'v=DMARC1\; p=reject\;',
                     ],
                 },