dns-generator/main.py

import argparse
import string
import itertools
from entry import entry
from namebench import namebench
from bind import Bind
from pdns import Pdns
from nsd import Nsd
from knot import Knot
from yadifa import Yadifa


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Generate fake DNS zonefiles and corresponding namebench file')
    parser.add_argument('--subdomains', type=bool, default=True, help='Generate subdomains too')
    parser.add_argument('--postfix', type=str, default='.test', help='Domain postfix eg .test or .com')
    parser.add_argument('--no-namebench', type=bool, default=False, help='Skip generation of namebench file')
    parser.add_argument('--no-zonefile', type=bool, default=False, help='Skip generation of zonefiles')
    parser.add_argument('--dnssec', type=bool, default=False, help='Prepare generation for DNSSEC')
    parser.add_argument('program', metavar='P', type=str, help='Program to generate (BIND, PDNS, NSD, KNOT, YADIFA, *ALL*)')
    parser.add_argument('domains', metavar='N', type=int, help='Number of domains to generate')

    args = parser.parse_args()

    if args.dnssec is not False:
        args.dnssec = True
    else:
        args.dnssec = False

    print(args)

    chars = list(string.ascii_lowercase)

    names = list(itertools.product(chars, repeat=4))

    entries = list()

    for i in range(0, args.domains):
        domain = ''.join(names[i]) + args.postfix
        record = entry(domain)

        if(args.subdomains):
            record.add_subdomain(entry('www', 'A', '127.0.0.1'))
            record.add_subdomain(entry('sub', 'A', '127.0.0.2'))

        entries.append(record)

    print('Generated {0} DNS domains'.format(len(entries)))

    if not args.no_namebench:
        print('Creating namebench file \'namebench.txt\'')
        nb = namebench()
        for record in entries:
            nb.generate_to_buffer(record)
        nb.write_buffer()

    if not args.no_zonefile:
        programs = {
            'BIND': Bind(entries),
            'PDNS': Pdns(entries),
            'NSD': Nsd(entries),
            'KNOT': Knot(entries),
            'YADIFA': Yadifa(entries),
        }

        if args.program == 'ALL':
            for i in programs:
                print('Creating zones for {0}'.format(i))
                programs[i].named(args.dnssec)
                programs[i].zones()
        else:
            p = programs.get(args.program, 'UNKNOWN')

            if p != 'UNKNOWN':
                p.named(args.dnssec)
                p.zones()
            else:
                print('Unknown program {0}'.format(args.program))

    print('Done')