OIption --file, to read domain names from a file. Closes #4

homer.py

@@ -13,12 +13,14 @@ import getopt
 import urllib.parse
 import time
 
+# Values that can be changed from the command line
 post = False
 verbose = False
 insecure = False
 head = False
 rtype = 'AAAA'
 tests = 1 # Number of repeated tests
+ifile = None # Input file
 
 def error(msg=None):
     if msg is None:
@@ -31,9 +33,12 @@ def usage(msg=None):
         print(msg,file=sys.stderr)
     print("Usage: %s [-P] [-k] url domain-name [DNS type]" % sys.argv[0], file=sys.stderr)
 
+# Main program
+name = None
+message = None
 try:
-    optlist, args = getopt.getopt (sys.argv[1:], "hvPker:",
-                                   ["help", "verbose", "head", "insecure", "POST", "repeat="])
+    optlist, args = getopt.getopt (sys.argv[1:], "hvPker:f:",
+                                   ["help", "verbose", "head", "insecure", "POST", "repeat=", "file="])
     for option, value in optlist:
         if option == "--help" or option == "-h":
             usage()
@@ -50,6 +55,8 @@ try:
             tests = int(value)
             if tests <= 1:
                 error("--repeat needs a value > 1")
+        elif option == "--file" or option == "-f":
+            ifile = value
         else:
             error("Unknown option %s" % option)
 except getopt.error as reason:
@@ -58,24 +65,31 @@ except getopt.error as reason:
 if post and head:
     usage("POST or HEAD but not both")
     sys.exit(1)
-if len(args) != 2 and len(args) != 3:
+if ifile is None and (len(args) != 2 and len(args) != 3):
     usage("Wrong number of arguments")
     sys.exit(1)
+if ifile is not None and len(args) != 1:
+    usage("Wrong number of arguments (if --file is used, do not indicate the domain name)")
+    sys.exit(1)
 url = args[0]
-name = args[1]
-if len(args) == 3:
+if ifile is None:
+    name = args[1]
+    if len(args) == 3:
         rtype = args[2]
 c = pycurl.Curl()
-message = dns.message.make_query(name, dns.rdatatype.from_text(rtype))
-message.id = 0 # DoH requests that
+if ifile is None:
+    message = dns.message.make_query(name, dns.rdatatype.from_text(rtype))
+    message.id = 0 # DoH requests that
 if head:
     c.setopt(pycurl.NOBODY, True)
 if post:
     c.setopt(c.URL, url)
-    data = message.to_wire()
     c.setopt(pycurl.POST, True)
+    if ifile is None:
+        data = message.to_wire()
         c.setopt(pycurl.POSTFIELDS, data)
 else:
+    if ifile is None:
         dns_req = base64.urlsafe_b64encode(message.to_wire()).decode('UTF8').rstrip('=')
         c.setopt(c.URL, url + ("?dns=%s" % dns_req))
 c.setopt(pycurl.HTTPHEADER, ["Content-type: application/dns-message"])
@@ -90,8 +104,28 @@ if insecure:
 c.setopt(pycurl.HTTP_VERSION, pycurl.CURL_HTTP_VERSION_2)
 ok = True
 start = time.time()
+if ifile is not None:
+    input = open(ifile)
 for i in range (0, tests):
-    print("Test %i" % i)
+    if tests > 1:
+        print("\nTest %i" % i)
+    if ifile is not None:
+        line = input.readline()
+        if line[:-1] == "":
+            error("Not enough data in %s for the %i tests" % (ifile, tests))
+        if line.find(' ') == -1:
+            name = line[:-1]
+            rtype = 'AAAA'
+        else:
+            (name, rtype) = line.split()
+        message = dns.message.make_query(name, dns.rdatatype.from_text(rtype))
+        message.id = 0 # DoH requests that
+        if post:
+            data = message.to_wire()
+            c.setopt(pycurl.POSTFIELDS, data)
+        else:
+            dns_req = base64.urlsafe_b64encode(message.to_wire()).decode('UTF8').rstrip('=')
+            c.setopt(c.URL, url + ("?dns=%s" % dns_req))
     buffer = io.BytesIO()
     c.setopt(c.WRITEDATA, buffer)
     c.perform()
@@ -114,9 +148,17 @@ for i in range (0, tests):
         print("HTTP error %i: %s" % (rcode, body[0:1000].decode()), file=sys.stderr)
         ok = False
     buffer.close()
+    if tests > 1 and i == 0:
+        start2 = time.time()
 c.close()
 stop = time.time()
-print("Total elapsed time: %.2f seconds (%.2f ms/request)" % (stop-start, (stop-start)*1000/tests))
+if tests > 1:
+    extra = ", %.2f ms/request if we ignore the first one" % ((stop-start2)*1000/(tests-1))
+else:
+    extra = ""
+print("\nTotal elapsed time: %.2f seconds (%.2f ms/request %s)" % (stop-start, (stop-start)*1000/tests, extra))
+if ifile is not None:
+    input.close()
 if ok:
     sys.exit(0)
 else: