#!/usr/bin/env python

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

import os, base64, argparse, jinja2

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

templates = dict(

    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 %}
''',

    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 %}
''')

ap = argparse.ArgumentParser(description = __doc__)
ap.add_argument("-a", "--algorithm", choices = algorithm_choices, default = algorithm_choices[0], help = "TSIG algorithm")
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 = [], help = "address(es) of primary nameserver(s)")
ap.add_argument("-z", "--zones",     nargs = "+", default = [], 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_bits[args.algorithm])).decode("ascii")

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