qprep: cleanup and preparation for user interface implementation

respdiff/blacklist.py

+import dns.rdatatype
+
+def obj_blacklisted(msg):
+    """
+    Detect blacklisted DNS message objects.
+    """
+    if len(msg.question) >= 1:
+        if msg.question[0].rdtype == dns.rdatatype.ANY:
+            return True
+    return False

respdiff/qprep.py

 import multiprocessing.pool as pool
 import sys
 
+import dns.message
+import dns.rdatatype
 import lmdb
 
+import blacklist
 import dbhelper
-import qtext
 
 
-def read_text():
+def read_lines(instream):
+    """
+    Yield (line number, stripped line text). Skip empty lines.
+    """
     i = 0
-    for line in sys.stdin:
+    for line in instream:
         line = line.strip()
         if line:
             i += 1
             yield (i, line)
 
 
-def gen_q(qstr):
+def int_or_fromtext(value, fromtext):
     try:
-        qry = qtext.qfromtext(qstr.split())
-    except BaseException:
-        print('line malformed: %s' % qstr)
-        return
-    if qtext.is_blacklisted(qry):
-        return
-    return qry.to_wire()
+        return int(value)
+    except ValueError:
+        return fromtext(value)
 
 
-def write_lmdb(qid, wire):
-    """
-    Worker: write query wire format into database
+def q_fromtext(line):
     """
-    global env
-    global db
+    Convert line from <qname> <RR type> to DNS query in IN class.
 
-    key = dbhelper.qid2key(qid)
-    with env.begin(db, write=True) as txn:
-        txn.put(key, wire)
+    Returns: DNS message object
+    Raises: ValueError or dns.exception.Exception on invalid input
+    """
+    qname, qtype = line.rsplit(None, 1)
+    qname = dns.name.from_text(qname)
+    qtype = int_or_fromtext(qtype, dns.rdatatype.from_text)
+    return dns.message.make_query(qname, qtype, dns.rdataclass.IN,
+                                  want_dnssec=True, payload=4096)
 
 
-def lmdb_init(envdir):
+def wrk_lmdb_init(envdir):
     """
     Worker: initialize LMDB env and open 'queries' database
     """
@@ -56,15 +59,34 @@ def lmdb_init(envdir):
     db = env.open_db(key=b'queries', **dbhelper.db_open)
 
 
-def gen_wrapper_lmdb(args):
-    qid, qstr = args
-    wire = gen_q(qstr)
-    write_lmdb(qid, wire)
+def wrk_lmdb_write(qid, wire):
+    """
+    Worker: write query wire format into database
+    """
+    global env
+    global db
+
+    key = dbhelper.qid2key(qid)
+    with env.begin(db, write=True) as txn:
+        txn.put(key, wire)
+
+
+def wrk_process_line(args):
+    """
+    Worker: parse input line and write (qid, wire formt) to LMDB queries DB
+
+    Skips over empty lines, raises for malformed inputs.
+    """
+    qid, line = args
+    msg = q_fromtext(line)
+    if not blacklist.obj_blacklisted(msg):
+        wrk_lmdb_write(qid, msg.to_wire())
+
 
 def main():
-    qstream = read_text()
-    with pool.Pool(initializer=lmdb_init, initargs=(sys.argv[1],)) as workers:
-        for _ in workers.imap_unordered(gen_wrapper_lmdb, qstream, chunksize=1000):
+    qstream = read_lines(sys.stdin)
+    with pool.Pool(initializer=wrk_lmdb_init, initargs=(sys.argv[1],)) as workers:
+        for _ in workers.imap_unordered(wrk_process_line, qstream, chunksize=1000):
             pass
 
 if __name__ == '__main__':

respdiff/qtext.py

-import argparse
-import sys
-
-import dns.name
-import dns.message
-import dns.rdataclass
-import dns.rdatatype
-
-
-qparser = argparse.ArgumentParser(description='Generate DNS message with query')
-qparser.add_argument('qname', type=lambda x: int_or_fromtext(x, dns.name.from_text))
-qparser.add_argument('qclass', type=lambda x: int_or_fromtext(x, dns.rdataclass.from_text), nargs='?', default='IN')
-qparser.add_argument('qtype', type=lambda x: int_or_fromtext(x, dns.rdatatype.from_text))
-
-def int_or_fromtext(value, fromtext):
-    try:
-        return int(value)
-    except ValueError:
-        return fromtext(value)
-
-def qfromtext(*args):
-    arglist = ['--'] + args[0]
-    args = qparser.parse_args(arglist)
-    return dns.message.make_query(args.qname, args.qtype, args.qclass, want_dnssec=True, payload=4096)
-
-
-def is_blacklisted(msg):
-    if len(msg.question) >= 1:
-        if msg.question[0].rdtype == dns.rdatatype.ANY:
-            return True
-    return False
-
-
-def main():
-    qry = qfromtext(sys.argv[1:])
-    if is_blacklisted(qry):
-        sys.exit('query blacklisted')
-    sys.stdout.write(qry.to_wire())
-
-
-if __name__ == "__main__":
-    main()