#!/usr/bin/env python

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

import os, base64, argparse

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

bind9_template = '''
key {name} {{
        algorithm {algorithm};
        secret "{secret}";
}};
'''

nsd_template = '''
key:
        name: "{name}"
        algorithm: {algorithm}
        secret: "{secret}"
'''

ap = argparse.ArgumentParser(description = __doc__)
ap.add_argument("-a", "--algorithm",
                choices = tuple(sorted(algorithm_bits)),
                default = sorted(algorithm_bits)[0])
ap.add_argument("-f", "--format",
                choices = ("bind9", "nsd"))
ap.add_argument("-n", "--name",
                default = "tsig.example.org")
ap.add_argument("-o", "--output",
                default = "-", type = argparse.FileType("w"))
ap.add_argument("-s", "--servers", nargs = "+")
ap.add_argument("-z", "--zones",   nargs = "+")
args = ap.parse_args()

params = dict(
    name = args.name,
    algorithm = args.algorithm,
    secret = base64.b64encode(os.urandom(algorithm_bits[args.algorithm])).decode("ascii"),
)

if args.format is None or args.format == "bind9":
    args.output.write(bind9_template.format(**params))

if args.format is None or args.format == "nsd":
    args.output.write(nsd_template.format(**params))