From 27bfe741862ae8b170e02251324bf4c553d204ed Mon Sep 17 00:00:00 2001
From: Benjamin Collet <benjamin@collet.eu>
Date: Sun, 11 Apr 2021 10:24:41 +0200
Subject: [PATCH] Add support for per-zone template

---
 config.yml.dist                               | 25 +++++++++++--------
 exporter.py                                   | 12 +++++----
 .../zonefile.jj2.dist                         |  0
 3 files changed, 21 insertions(+), 16 deletions(-)
 rename zonefile.jj2.dist => templates/zonefile.jj2.dist (100%)

diff --git a/config.yml.dist b/config.yml.dist
index 103e5cb..b67549a 100644
--- a/config.yml.dist
+++ b/config.yml.dist
@@ -1,13 +1,16 @@
 netbox:
-  url:    'https://netbox.local'
-  token:  ''
+  url:        'https://netbox.local'
+  token:      ''
 zones:
-  - type: 'reverse'
-    name: '192.0.2.0/24'
-    file: '/path/to/file/db.192.0.2'
-  - type: 'reverse'
-    name: '2001:db8::/32'
-    file: '/path/to/file/db'2001:db8'
-  - type: 'forward'
-    name: 'example.com'
-    file: '/path/to/file/example.com'
+  - type:     'reverse'
+    name:     '192.0.2.0/24'
+    file:     '/path/to/file/db.192.0.2'
+    template: 'zonefile.jj2'
+  - type:     'reverse'
+    name:     '2001:db8::/32'
+    file:     '/path/to/file/db'2001:db8'
+    template: 'zonefile.jj2'
+  - type:     'forward'
+    name:     'example.com'
+    file:     '/path/to/file/example.com'
+    template: 'zonefile.jj2'
diff --git a/exporter.py b/exporter.py
index 0040e4d..bcab05a 100755
--- a/exporter.py
+++ b/exporter.py
@@ -15,7 +15,7 @@ config = os.path.join(os.path.dirname(os.path.realpath(__file__)),'config.yml')
 with open(config, 'r') as ymlfile:
     cfg = yaml.load(ymlfile, Loader=yaml.FullLoader)
 
-templates_dir = os.path.dirname(os.path.realpath(__file__))
+templates_dir = os.path.join(os.path.dirname(os.path.realpath(__file__)),'templates')
 
 def ptr(nb, args):
     device_primary = {}
@@ -81,7 +81,7 @@ def ptr(nb, args):
 
     file_loader = FileSystemLoader(templates_dir)
     env = Environment(loader=file_loader)
-    template = env.get_template("zonefile.jj2")
+    template = env.get_template(args.template)
     output = template.render(serial=serial,records=records)
 
     if not hasattr(args, 'output'):
@@ -166,7 +166,7 @@ def dns(nb, args):
 
     file_loader = FileSystemLoader(templates_dir)
     env = Environment(loader=file_loader)
-    template = env.get_template("zonefile.jj2")
+    template = env.get_template(args.template)
     output = template.render(serial=serial,records=records)
 
     if not hasattr(args, 'output'):
@@ -182,9 +182,9 @@ def autodns(nb, args):
 
     for zone in cfg['zones']:
         if zone['type'] == 'reverse':
-            ptr(nb, argparse.Namespace(prefix=zone['name'],output=zone['file']))
+            ptr(nb, argparse.Namespace(prefix=zone['name'],template=zone['template'],output=zone['file']))
         elif zone['type'] == 'forward':
-            dns(nb, argparse.Namespace(domain=zone['name'],output=zone['file']))
+            dns(nb, argparse.Namespace(domain=zone['name'],template=zone['template'],output=zone['file']))
 
 
 if __name__ == '__main__':
@@ -200,9 +200,11 @@ if __name__ == '__main__':
 
     subparser = subparsers.add_parser('ptr', help='Generate reverse DNS zone file for prefix')
     subparser.add_argument('prefix', type=str, help='Prefix')
+    subparser.add_argument('template', type=str, help='template')
 
     subparser = subparsers.add_parser('dns', help='Generate DNS zone file for domain')
     subparser.add_argument('domain', type=str, help='Domain')
+    subparser.add_argument('template', type=str, help='template')
 
     args = parser.parse_args()
     globals()[args.action](nb, args)
diff --git a/zonefile.jj2.dist b/templates/zonefile.jj2.dist
similarity index 100%
rename from zonefile.jj2.dist
rename to templates/zonefile.jj2.dist