From 48831659e5896e4354f4affec71f8bdf5fe082bd Mon Sep 17 00:00:00 2001
From: Ross McFarland <rwmcfa1@gmail.com>
Date: Wed, 14 Sep 2022 06:46:13 -0700
Subject: [PATCH] Add allowed_ttls support to TtlRestrictionFilter

---
 octodns/processor/restrict.py            | 17 +++++++---
 tests/test_octodns_processor_restrict.py | 41 +++++++++++++-----------
 2 files changed, 36 insertions(+), 22 deletions(-)

diff --git a/octodns/processor/restrict.py b/octodns/processor/restrict.py
index 8afc0a5..1903995 100644
--- a/octodns/processor/restrict.py
+++ b/octodns/processor/restrict.py
@@ -18,9 +18,12 @@ class RestrictionException(ProcessorException):
 
 class TtlRestrictionFilter(BaseProcessor):
     '''
-    Ensure that configured TTLs are between a configured minimum and maximum.
+    Ensure that configured TTLs are between a configured minimum and maximum or
+    in an allowed set of values.
+
     The default minimum is 1 (the behavior of 0 is undefined spec-wise) and the
-    default maximum is 604800 (seven days.)
+    default maximum is 604800 (seven days.) allowed_ttls is only used when
+    explicitly configured and min and max are ignored in that case.
 
     Example usage:
 
@@ -29,6 +32,7 @@ class TtlRestrictionFilter(BaseProcessor):
         class: octodns.processor.restrict.TtlRestrictionFilter
         min_ttl: 60
         max_ttl: 3600
+        # allowed_ttls: [300, 900, 3600]
 
     zones:
       exxampled.com.:
@@ -54,16 +58,21 @@ class TtlRestrictionFilter(BaseProcessor):
 
     SEVEN_DAYS = 60 * 60 * 24 * 7
 
-    def __init__(self, name, min_ttl=1, max_ttl=SEVEN_DAYS):
+    def __init__(self, name, min_ttl=1, max_ttl=SEVEN_DAYS, allowed_ttls=None):
         super().__init__(name)
         self.min_ttl = min_ttl
         self.max_ttl = max_ttl
+        self.allowed_ttls = set(allowed_ttls) if allowed_ttls else None
 
     def process_source_zone(self, zone, *args, **kwargs):
         for record in zone.records:
             if record._octodns.get('lenient'):
                 continue
-            if record.ttl < self.min_ttl:
+            if self.allowed_ttls and record.ttl not in self.allowed_ttls:
+                raise RestrictionException(
+                    f'{record.fqdn} ttl={record.ttl} not an allowed value, allowed_ttls={self.allowed_ttls}'
+                )
+            elif record.ttl < self.min_ttl:
                 raise RestrictionException(
                     f'{record.fqdn} ttl={record.ttl} too low, min_ttl={self.min_ttl}'
                 )
diff --git a/tests/test_octodns_processor_restrict.py b/tests/test_octodns_processor_restrict.py
index 9aae744..4848ae6 100644
--- a/tests/test_octodns_processor_restrict.py
+++ b/tests/test_octodns_processor_restrict.py
@@ -9,24 +9,6 @@ from octodns.zone import Zone
 
 
 class TestTtlRestrictionFilter(TestCase):
-    zone = Zone('unit.tests.', [])
-    matches = Record.new(
-        zone, 'matches', {'type': 'A', 'ttl': 42, 'value': '1.2.3.4'}
-    )
-    zone.add_record(matches)
-    doesnt = Record.new(
-        zone, 'doesnt', {'type': 'A', 'ttl': 42, 'value': '2.3.4.5'}
-    )
-    zone.add_record(doesnt)
-    matchable1 = Record.new(
-        zone, 'start-f43ad96-end', {'type': 'A', 'ttl': 42, 'value': '3.4.5.6'}
-    )
-    zone.add_record(matchable1)
-    matchable2 = Record.new(
-        zone, 'start-a3b444c-end', {'type': 'A', 'ttl': 42, 'value': '4.5.6.7'}
-    )
-    zone.add_record(matchable2)
-
     def test_restrict_ttl(self):
         # configured values
         restrictor = TtlRestrictionFilter('test', min_ttl=32, max_ttl=1024)
@@ -106,3 +88,26 @@ class TestTtlRestrictionFilter(TestCase):
             'high.unit.tests. ttl=999999 too high, max_ttl=604800',
             str(ctx.exception),
         )
+
+        # allowed_ttls
+        restrictor = TtlRestrictionFilter('test', allowed_ttls=[42, 300])
+
+        # add 300 (42 is already there)
+        another = Record.new(
+            zone, 'another', {'type': 'A', 'ttl': 300, 'value': '4.5.6.7'}
+        )
+        zone.add_record(another)
+
+        # 42 and 300 are allowed through
+        restricted = restrictor.process_source_zone(zone)
+        self.assertEqual(zone.records, restricted.records)
+
+        # 16 is not
+        copy = zone.copy()
+        copy.add_record(low)
+        with self.assertRaises(RestrictionException) as ctx:
+            restrictor.process_source_zone(copy)
+        self.assertEqual(
+            'low.unit.tests. ttl=0 not an allowed value, allowed_ttls={42, 300}',
+            str(ctx.exception),
+        )