From 37da9126c62103a760f949e571d6d17a6209746e Mon Sep 17 00:00:00 2001
From: Knut Ahlers <knut@ahlers.me>
Date: Fri, 14 Dec 2018 19:18:10 +0100
Subject: [PATCH] Add consul support

Signed-off-by: Knut Ahlers <knut@ahlers.me>
---
 checkZonefile.py     | 11 +++++++++--
 consul.py            | 46 ++++++++++++++++++++++++++++++++++++++++++++
 generateZonefiles.py |  7 +++++++
 requirements.txt     |  1 +
 4 files changed, 63 insertions(+), 2 deletions(-)
 create mode 100644 consul.py

diff --git a/checkZonefile.py b/checkZonefile.py
index 5d4bb50..0d85cc3 100644
--- a/checkZonefile.py
+++ b/checkZonefile.py
@@ -136,18 +136,25 @@ def check_nameserver(nameservers):
 
 
 def check_zone(name, config):
-    expected_zone = ['mailserver', 'mailserver_set', 'entries', 'default_ttl']
+    expected_zone = ['mailserver', 'mailserver_set',
+                     'entries', 'default_ttl', 'from_consul']
 
     for k in config.keys():
         if k not in expected_zone:
             warn('Unexpected entry in zone {} found: {}'.format(name, k))
 
-    if 'mailserver' not in config and 'mailserver_set' not in config and ('entries' not in config or len(config['entries']) == 0):
+    if 'mailserver' not in config and 'mailserver_set' not in config and ('entries' not in config or len(config['entries']) == 0) and 'from_consul' not in config:
         warn('Zone {} has no mailservers and no entries'.format(name))
 
     if 'mailserver' in config and 'mailserver_set' in config:
         error('Zone {} contains mailserver and mailserver_set'.format(name))
 
+    if 'from_consul' in config and config['from_consul'] and 'entries' in config:
+        error('Zone {} contains entries and from_consul flag'.format(name))
+
+    if 'from_consul' in config and config['from_consul'] and 'mailserver' in config:
+        warn('Zone {} contains mailserver and from_consul flag'.format(name))
+
     if 'mailserver' in config:
         for mx, weight in config['mailserver'].items():
             if not str(weight).isdigit():
diff --git a/consul.py b/consul.py
new file mode 100644
index 0000000..eca7da5
--- /dev/null
+++ b/consul.py
@@ -0,0 +1,46 @@
+import base64
+import json
+import os
+import requests
+
+
+def query_zone_entries(zone):
+    if os.getenv('CONSUL_HTTP_ADDR') == '' or os.getenv('CONSUL_HTTP_TOKEN') == '':
+        raise Exception(
+            'Consul query does not work with CONSUL_HTTP_ADDR or CONSUL_HTTP_TOKEN unset')
+
+    return parse_raw_consul(zone)
+
+
+def read_raw_from_consul(zone):
+    resp = requests.get('{}/v1/kv/dns/{}?recurse=true'.format(
+        os.getenv('CONSUL_HTTP_ADDR'),
+        zone.rstrip('.'),
+    ),
+        headers={
+        'X-Consul-Token': os.getenv('CONSUL_HTTP_TOKEN'),
+    })
+
+    if resp.status_code == 404:
+        return []
+
+    return resp.json()
+
+
+def parse_raw_consul(zone):
+    entries = []
+
+    for raw_entry in read_raw_from_consul(zone):
+        sub_entries = json.loads(base64.b64decode(raw_entry['Value']))
+
+        # Key consists of at least 2 elements: dns/ahlers.me/subdomain OR dns/ahlers.me
+        key = raw_entry['Key'].split('/')[2:]
+        name = ''
+        if len(key) > 0 and key[0] != '@':
+            name = key[0]
+
+        for entry in sub_entries:
+            entry['name'] = name
+            entries.append(entry)
+
+    return entries
diff --git a/generateZonefiles.py b/generateZonefiles.py
index fbdd597..72b4044 100644
--- a/generateZonefiles.py
+++ b/generateZonefiles.py
@@ -8,6 +8,9 @@ import os.path
 import sys
 import time
 
+# Custom modules
+import consul
+
 # Third-party imports
 import dns.resolver
 import dns.rdatatype
@@ -138,6 +141,10 @@ def main():
         for entry in default(config, 'entries', []):
             entries.extend(sanitize(entry))
 
+        if default(config, 'from_consul', False):
+            for entry in consul.query_zone_entries(zone):
+                entries.extend(sanitize(entry))
+
         mailserver = default(config, 'mailserver', {})
         if 'mailserver_set' in config and config['mailserver_set'] in zone_data['mailserver_sets']:
             mailserver = zone_data['mailserver_sets'][config['mailserver_set']]
diff --git a/requirements.txt b/requirements.txt
index bed4974..3e39ac9 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,3 +1,4 @@
 jinja2
 PyYAML
 dnspython
+requests