add support for wildcard and naked domains to PowerDNS module

lemur/plugins/lemur_acme/powerdns.py

@@ -49,16 +49,20 @@ class Record:
         return self._data["name"]
 
     @property
-    def disabled(self):
+    def type(self):
-        return self._data["disabled"]
+        return self._data["type"]
+
+    @property
+    def ttl(self):
+        return self._data["ttl"]
 
     @property
     def content(self):
         return self._data["content"]
 
     @property
-    def ttl(self):
+    def disabled(self):
-        return self._data["ttl"]
+        return self._data["disabled"]
 
 
 def get_zones(account_number):
@@ -92,42 +96,32 @@ def get_zones(account_number):
 def create_txt_record(domain, token, account_number):
     """ Create a TXT record for the given domain and token and return a change_id tuple """
     _check_conf()
-    zone_name = _get_zone_name(domain, account_number)
+
-    server_id = current_app.config.get("ACME_POWERDNS_SERVERID", "localhost")
-    zone_id = zone_name + "."
-    domain_id = domain + "."
-    path = f"/api/v1/servers/{server_id}/zones/{zone_id}"
-    payload = {
-        "rrsets": [
-            {
-                "name": domain_id,
-                "type": "TXT",
-                "ttl": 300,
-                "changetype": "REPLACE",
-                "records": [
-                    {
-                        "content": f"\"{token}\"",
-                        "disabled": False
-                    }
-                ],
-                "comments": []
-            }
-        ]
-    }
     function = sys._getframe().f_code.co_name
     log_data = {
         "function": function,
         "fqdn": domain,
         "token": token,
     }
+
+    # Create new record
+    domain_id = domain + "."
+    records = [Record({'name': domain_id, 'content': f"\"{token}\"", 'disabled': False})]
+
+    # Get current records
+    cur_records = _get_txt_records(domain)
+    for record in cur_records:
+        if record.content != token:
+            records.append(record)
+
     try:
-        _patch(path, payload)
+        _patch_txt_records(domain, account_number, records)
-        log_data["message"] = "TXT record successfully created"
+        log_data["message"] = "TXT record(s) successfully created"
         current_app.logger.debug(log_data)
     except Exception as e:
         sentry.captureException()
         log_data["Exception"] = e
-        log_data["message"] = "Unable to create TXT record"
+        log_data["message"] = "Unable to create TXT record(s)"
         current_app.logger.debug(log_data)
 
     change_id = (domain, token)
@@ -173,43 +167,78 @@ def wait_for_dns_change(change_id, account_number=None):
 def delete_txt_record(change_id, account_number, domain, token):
     """ Delete the TXT record for the given domain and token """
     _check_conf()
-    zone_name = _get_zone_name(domain, account_number)
+
-    server_id = current_app.config.get("ACME_POWERDNS_SERVERID", "localhost")
-    zone_id = zone_name + "."
-    domain_id = domain + "."
-    path = f"/api/v1/servers/{server_id}/zones/{zone_id}"
-    payload = {
-        "rrsets": [
-            {
-                "name": domain_id,
-                "type": "TXT",
-                "ttl": 300,
-                "changetype": "DELETE",
-                "records": [
-                    {
-                        "content": f"\"{token}\"",
-                        "disabled": False
-                    }
-                ],
-                "comments": []
-            }
-        ]
-    }
     function = sys._getframe().f_code.co_name
     log_data = {
         "function": function,
         "fqdn": domain,
-        "token": token
+        "token": token,
     }
-    try:
+
-        _patch(path, payload)
+    # Determine if we can delete whole RRset or just one record
-        log_data["message"] = "TXT record successfully deleted"
+    cur_records = _get_txt_records(domain)
-        current_app.logger.debug(log_data)
+    found = False
-    except Exception as e:
+    new_records = []
-        sentry.captureException()
+    for record in cur_records:
-        log_data["Exception"] = e
+        if record.content == f"\"{token}\"":
-        log_data["message"] = "Unable to delete TXT record"
+            found = True
+        else:
+            new_records.append(record)
+
+    if not found:  # Record not found in DNS
+        log_data["message"] = "Unable to delete TXT record: TXT record not found"
         current_app.logger.debug(log_data)
+        return
+
+    elif new_records:  # Removing Record from RRSet via Patch
+        try:
+            _patch_txt_records(domain, account_number, new_records)
+            log_data["message"] = "TXT record successfully deleted"
+            current_app.logger.debug(log_data)
+        except Exception as e:
+            sentry.captureException()
+            log_data["Exception"] = e
+            log_data["message"] = "Unable to delete TXT record: patching exception"
+            current_app.logger.debug(log_data)
+
+    else:  # Delete current records
+        zone_name = _get_zone_name(domain, account_number)
+        server_id = current_app.config.get("ACME_POWERDNS_SERVERID", "localhost")
+        zone_id = zone_name + "."
+        domain_id = domain + "."
+        path = f"/api/v1/servers/{server_id}/zones/{zone_id}"
+        payload = {
+            "rrsets": [
+                {
+                    "name": domain_id,
+                    "type": "TXT",
+                    "ttl": 300,
+                    "changetype": "DELETE",
+                    "records": [
+                        {
+                            "content": f"\"{token}\"",
+                            "disabled": False
+                        }
+                    ],
+                    "comments": []
+                }
+            ]
+        }
+        function = sys._getframe().f_code.co_name
+        log_data = {
+            "function": function,
+            "fqdn": domain,
+            "token": token
+        }
+        try:
+            _patch(path, payload)
+            log_data["message"] = "TXT record successfully deleted"
+            current_app.logger.debug(log_data)
+        except Exception as e:
+            sentry.captureException()
+            log_data["Exception"] = e
+            log_data["message"] = "Unable to delete TXT record"
+            current_app.logger.debug(log_data)
 
 
 def _check_conf():
@@ -243,6 +272,33 @@ def _get_zone_name(domain, account_number):
     return zone_name
 
 
+def _get_txt_records(domain):
+    """Retrieve TXT records for a given domain and return list of Record Objects"""
+    server_id = current_app.config.get("ACME_POWERDNS_SERVERID", "localhost")
+
+    path = f"/api/v1/servers/{server_id}/search-data?q={domain}&max=100&object_type=record"
+    function = sys._getframe().f_code.co_name
+    log_data = {
+        "function": function
+    }
+    try:
+        records = _get(path)
+        log_data["message"] = "Retrieved TXT Records Successfully"
+        current_app.logger.debug(log_data)
+
+    except Exception as e:
+        sentry.captureException()
+        log_data["message"] = "Failed to Retrieve TXT Records"
+        current_app.logger.debug(log_data)
+        raise
+
+    txt_records = []
+    for record in records:
+        cur_record = Record(record)
+        txt_records.append(cur_record)
+    return txt_records
+
+
 def _get(path, params=None):
     """ Execute a GET request on the given URL (base_uri + path) and return response as JSON object """
     base_uri = current_app.config.get("ACME_POWERDNS_DOMAIN")
@@ -257,6 +313,40 @@ def _get(path, params=None):
     return resp.json()
 
 
+def _patch_txt_records(domain, account_number, records):
+    """Send Patch request to PowerDNS Server"""
+
+    domain_id = domain + "."
+
+    # Create records
+    txt_records = []
+    for record in records:
+        txt_records.append(
+            {'content': record.content, 'disabled': record.disabled}
+        )
+
+    # Create RRSet
+    payload = {
+        "rrsets": [
+            {
+                "name": domain_id,
+                "type": "TXT",
+                "ttl": 300,
+                "changetype": "REPLACE",
+                "records": txt_records,
+                "comments": []
+            }
+        ]
+    }
+
+    # Create Txt Records
+    server_id = current_app.config.get("ACME_POWERDNS_SERVERID", "localhost")
+    zone_name = _get_zone_name(domain, account_number)
+    zone_id = zone_name + "."
+    path = f"/api/v1/servers/{server_id}/zones/{zone_id}"
+    _patch(path, payload)
+
+
 def _patch(path, payload):
     """ Execute a Patch request on the given URL (base_uri + path) with given payload """
     base_uri = current_app.config.get("ACME_POWERDNS_DOMAIN")

lemur/plugins/lemur_acme/tests/test_powerdns.py

@@ -48,13 +48,14 @@ class TestPowerdns(unittest.TestCase):
         self.assertEqual(result, zone)
 
     @patch("lemur.plugins.lemur_acme.powerdns.current_app")
-    def test_create_txt_record(self, mock_current_app):
+    def test_create_txt_record_write_only(self, mock_current_app):
         domain = "_acme_challenge.test.example.com"
         zone = "test.example.com"
         token = "ABCDEFGHIJ"
         account_number = "1234567890"
         change_id = (domain, token)
         powerdns._check_conf = Mock()
+        powerdns._get_txt_records = Mock(return_value=[])
         powerdns._get_zone_name = Mock(return_value=zone)
         mock_current_app.logger.debug = Mock()
         mock_current_app.config.get = Mock(return_value="localhost")
@@ -63,24 +64,74 @@ class TestPowerdns(unittest.TestCase):
             "function": "create_txt_record",
             "fqdn": domain,
             "token": token,
-            "message": "TXT record successfully created"
+            "message": "TXT record(s) successfully created"
         }
         result = powerdns.create_txt_record(domain, token, account_number)
         mock_current_app.logger.debug.assert_called_with(log_data)
         self.assertEqual(result, change_id)
 
+    @patch("lemur.plugins.lemur_acme.powerdns.current_app")
+    def test_create_txt_record_append(self, mock_current_app):
+        domain = "_acme_challenge.test.example.com"
+        zone = "test.example.com"
+        token = "ABCDEFGHIJ"
+        account_number = "1234567890"
+        change_id = (domain, token)
+        powerdns._check_conf = Mock()
+        cur_token = "123456"
+        cur_records = [powerdns.Record({'name': domain, 'content': cur_token, 'disabled': False})]
+        powerdns._get_txt_records = Mock(return_value=cur_records)
+        powerdns._get_zone_name = Mock(return_value=zone)
+        mock_current_app.logger.debug = Mock()
+        mock_current_app.config.get = Mock(return_value="localhost")
+        powerdns._patch = Mock()
+        log_data = {
+            "function": "create_txt_record",
+            "fqdn": domain,
+            "token": token,
+            "message": "TXT record(s) successfully created"
+        }
+        expected_path = f"/api/v1/servers/localhost/zones/test.example.com."
+        expected_payload = {
+            "rrsets": [
+                {
+                    "name": domain + ".",
+                    "type": "TXT",
+                    "ttl": 300,
+                    "changetype": "REPLACE",
+                    "records": [
+                        {
+                            "content": f"\"{token}\"",
+                            "disabled": False
+                        },
+                        {
+                            "content": f"\"{cur_token}\"",
+                            "disabled": False
+                        }
+                    ],
+                    "comments": []
+                }
+            ]
+        }
+
+        result = powerdns.create_txt_record(domain, token, account_number)
+        mock_current_app.logger.debug.assert_called_with(log_data)
+        powerdns._patch.assert_called_with(expected_path, expected_payload)
+        self.assertEqual(result, change_id)
+
     @patch("lemur.plugins.lemur_acme.powerdns.dnsutil")
     @patch("lemur.plugins.lemur_acme.powerdns.current_app")
     @patch("lemur.extensions.metrics")
     @patch("time.sleep")
     def test_wait_for_dns_change(self, mock_sleep, mock_metrics, mock_current_app, mock_dnsutil):
         domain = "_acme-challenge.test.example.com"
-        token = "ABCDEFG"
+        token1 = "ABCDEFG"
+        token2 = "HIJKLMN"
         zone_name = "test.example.com"
         nameserver = "1.1.1.1"
-        change_id = (domain, token)
+        change_id = (domain, token1)
         powerdns._check_conf = Mock()
-        mock_records = (token,)
+        mock_records = (token2, token1)
         mock_current_app.config.get = Mock(return_value=1)
         powerdns._get_zone_name = Mock(return_value=zone_name)
         mock_dnsutil.get_authoritative_nameserver = Mock(return_value=nameserver)
@@ -114,7 +165,7 @@ class TestPowerdns(unittest.TestCase):
             "function": "delete_txt_record",
             "fqdn": domain,
             "token": token,
-            "message": "TXT record successfully deleted"
+            "message": "Unable to delete TXT record: TXT record not found"
         }
         powerdns.delete_txt_record(change_id, account_number, domain, token)
         mock_current_app.logger.debug.assert_called_with(log_data)