1#!/usr/bin/python
2
3#
4# $Id: gen-data-queryperf.py,v 1.2 2008/06/13 18:17:08 jinmei Exp $
5#
6# Contributed by Stephane Bortzmeyer <bortzmeyer@nic.fr>
7#
8# "A small tool which may be useful with contrib/queryperf. This script
9#  can generate files of queries, both with random names (to test the
10#  behaviour with NXdomain) and with domains from a real zone file."
11#
12
13import sys
14import getopt
15import random
16import re
17
18ldh = []
19# Letters
20for i in range(97, 122):
21    ldh.append(chr(i))
22# Digits
23for i in range(48, 57):
24    ldh.append(chr(i))
25# Hyphen
26ldh.append('-')
27
28maxsize=10
29tld='org'
30num=4
31percent_random = 0.3
32gen = None
33zone_file = None
34domains = {}
35domain_ns = r'^([a-z0-9-\.]+)((\s+\d+)?(\s+IN)?|(\s+IN)(\s+\d+)?)\s+NS'
36domain_ns_re = re.compile(domain_ns, re.IGNORECASE)
37
38def remove_tld(label, tld):
39    if label.endswith('.' + tld + '.'):
40        return label[0:-(1+ len(tld) + 1)]
41    else:
42        return label
43
44def gen_random_label():
45    label = ""
46    for i in range(gen.randint(1, maxsize)):
47        label = label + gen.choice(ldh)
48    return label
49
50def make_domain(label):
51    return "www." + label + "." + tld + "     A"
52
53def usage():
54    sys.stdout.write("Usage: " + sys.argv[0] + " [-n number] " + \
55                     "[-p percent-random] [-t TLD]\n")
56    sys.stdout.write("       [-m MAXSIZE] [-f zone-file]\n")
57
58try:
59    optlist, args = getopt.getopt(sys.argv[1:], "hp:f:n:t:m:",
60                                  ["help", "percentrandom=", "zonefile=",
61                                   "number=", "tld=",
62                                   "maxsize="])
63    for option, value in optlist:
64        if option == "--help" or option == "-h":
65            usage()
66            sys.exit(0)
67        elif option == "--number" or option == "-n":
68            num = int(value)
69        elif option == "--maxsize" or option == "-m":
70            maxsize = int(value)
71        elif option == "--percentrandom" or option == "-p":
72            percent_random = float(value)
73        elif option == "--tld" or option == "-t":
74            tld = str(value)
75        elif option == "--zonefile" or option == "-f":
76            zone_file = str(value)
77        else:
78            error("Unknown option " + option)
79except getopt.error, reason:
80    sys.stderr.write(sys.argv[0] + ": " + str(reason) + "\n")
81    usage()
82    sys.exit(1)
83
84if len(args) <> 0:
85    usage()
86    sys.exit(1)
87
88gen = random.Random()
89if zone_file:
90    file = open(zone_file)
91    line = file.readline()
92    while line:
93        domain_line = domain_ns_re.match(line)
94        if domain_line:
95            print domain_line.group(1)
96            domain = remove_tld(domain_line.group(1), tld)
97            domains[domain] = 1
98        line = file.readline()
99    file.close()
100if zone_file:
101    domains = domains.keys()
102    if len(domains) == 0:
103        sys.stderr.write("No domains found in '%s'\n" % zone_file)
104        sys.exit(1)
105for i in range(num):
106    if zone_file:
107        if gen.random() < percent_random:
108            sys.stdout.write(make_domain(gen_random_label()))
109        else:
110            sys.stdout.write(make_domain(gen.choice(domains)))
111    else:
112        sys.stdout.write(make_domain(gen_random_label()))
113    sys.stdout.write("\n")
114