Commit eee27a71 authored by Petr Špaček's avatar Petr Špaček

qprep: cleanup and preparation for user interface implementation

parent 3a764e63
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
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__':
......
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()
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment