#!/usr/bin/env python

"Pure Python TSIG key generator with multiple output formats."

import os, base64, argparse, jinja2

algorithm_bytes = dict(("hmac-sha{}".format(bits), bits // 8) for bits in (256, 384, 512))
algorithm_choices = tuple(sorted(algorithm_bytes))
templates = {}

templates["bind9"] = '''
key {{ args.name }} {
    algorithm {{ args.algorithm }};
    secret "{{ args.key }}";
};

{% if args.servers %}
{%- for server in args.servers -%}
server {{ server }} { keys { {{ args.name }}; }; };
{% endfor -%}
{% for zone in args.zones %}
zone "{{ zone }}" {
    file "{{args.directory }}/{{ zone }}"; 
    type slave; 
    masters { {%- for server in args.servers %} {{ server }}; {% endfor %} };
};
{% endfor %}
{% endif %}
'''

templates["nsd"] = '''
key:
    name: "{{ args.name }}"
    algorithm: {{ args.algorithm }}
    secret: "{{ args.key }}"

{% if args.servers -%}
pattern:
    name: "secondary.{{ args.name }}"
    zonefile: "{{ args.directory }}/%s"
{%- for server in args.servers %}
    allow-notify: {{ server }} {{ args.name }}
    request-xfr:  {{ server }} {{ args.name }}
{%- endfor %}
{%- for zone in args.zones %}

zone:
    name: "{{ zone }}"
    include-pattern: "secondary.{{ args.name }}"
{%- endfor %}
{% endif %}
'''

templates["knot"] = '''
### WARNING: KNOT CONFIGURATION NOT YET TESTED ###

key:
  - id: {{ args.name }}
    algorithm: {{ args.algorithm }}
    secret: {{ args.key }}

{% if args.servers -%}
remote:
  - id: {{ args.name }}
    key: {{ args.name }}
{%- for server in args.servers %}
    address: {{ server }}@53
{%- endfor %}

acl:
  - id: notify_from_{{ args.name }}
    action: notify
    key: {{ args.name }}
{%- for server in args.servers %}
    address: {{ server }}@53
{%- endfor %}
{# #}
{%- for zone in args.zones %}
zone:
  - domain: {{ zone }}
    storage: {{ args.directory }}/{{ zone }}
    master: {{ args.name }}
    acl: notify_from_{{ args.name }}
{% endfor %}
{% endif %}
'''

ap = argparse.ArgumentParser(description = __doc__)
ap.add_argument("-a", "--algorithm", choices = algorithm_choices, default = algorithm_choices[0],
                help = "TSIG algorithm (default: " + algorithm_choices[0] + ")")
ap.add_argument("-d", "--directory", default = "secondary", help = "where to store secondary zone files")
ap.add_argument("-f", "--format",    choices = tuple(templates), default = "nsd", help = "output format (default: nsd)")
ap.add_argument("-k", "--key", help = "TSIG shared secret (default: generate new secret)")
ap.add_argument("-n", "--name",      default = "tsig.example.org", help = "DNS name for TSIG shared secret")
ap.add_argument("-o", "--output",    default = "-", type = argparse.FileType("w"), help = "output file")
ap.add_argument("-s", "--servers",   nargs = "+", default = [], metavar = "SERVER", help = "address(es) of primary nameserver(s)")
ap.add_argument("-z", "--zones",     nargs = "+", default = [], metavar = "ZONE", help = "zone(s) to xfr from specified nameserver(s)")
args = ap.parse_args()

if args.key is None:
    args.key = base64.b64encode(os.urandom(algorithm_bytes[args.algorithm])).decode("ascii")

args.output.write(jinja2.Template(templates[args.format]).render(args = args))