Commit 3a764e63 authored by Petr Špaček's avatar Petr Špaček

qprep: remove unused cruft from makedirs, rename it to qprep

qprep reads list of queries in text format and writes them into
'queries' sub-database of LMDB environment specified by argv[1]
parent e8b25669
......@@ -19,3 +19,8 @@ def key_stream(lenv, db):
while cont:
yield cur.key()
cont = cur.next()
def qid2key(qid):
"""Encode query ID to database key"""
return str(qid).encode('ascii')
import itertools
import multiprocessing.pool as pool
import errno
from pprint import pprint
import os
import sys
import lmdb
import dbhelper
import makeq
import qtext
def read_text():
......@@ -18,44 +14,35 @@ def read_text():
if line:
i += 1
yield (i, line)
if i % 10000 == 0:
print(i)
#if i > 770000:
# print(line)
def gen_q(qtext):
def gen_q(qstr):
try:
qry = makeq.qfromtext(qtext.split())
qry = qtext.qfromtext(qstr.split())
except BaseException:
print('line malformed: %s' % qtext)
print('line malformed: %s' % qstr)
return
if makeq.is_blacklisted(qry):
if qtext.is_blacklisted(qry):
return
return qry.to_wire()
def write_file(qid, wire):
dirname = '%09d' % qid
qfilename = '%s/q.dns' % dirname
try:
os.mkdir(dirname)
except OSError as ex:
if not ex.errno == errno.EEXIST:
raise
with open(qfilename, 'wb') as qfile:
qfile.write(wire)
def write_lmdb(qid, wire):
"""
Worker: write query wire format into database
"""
global env
global db
key = str(qid).encode('ascii')
key = dbhelper.qid2key(qid)
with env.begin(db, write=True) as txn:
txn.put(key, wire)
def lmdb_init(envdir):
"""
Worker: initialize LMDB env and open 'queries' database
"""
global env
global db
......@@ -68,25 +55,17 @@ def lmdb_init(envdir):
env = lmdb.Environment(**config)
db = env.open_db(key=b'queries', **dbhelper.db_open)
def gen_wrapper_lmdb(args):
qid, qtext = args
wire = gen_q(qtext)
qid, qstr = args
wire = gen_q(qstr)
write_lmdb(qid, wire)
def main():
qstream = itertools.islice(read_text(), 1000000)
#qstream = read_text()
#for i in map(gen_qfile, qstream):
# pass
with pool.Pool(initializer=lmdb_init, initargs=(sys.argv[1],)) as p:
for i in p.imap_unordered(gen_wrapper_lmdb, qstream, chunksize=1000):
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):
pass
# LMDB specifics
lmdb_init(sys.argv[1])
pprint(env.info())
env.close()
if __name__ == '__main__':
main()
......@@ -23,9 +23,6 @@ def qfromtext(*args):
args = qparser.parse_args(arglist)
return dns.message.make_query(args.qname, args.qtype, args.qclass, want_dnssec=True, payload=4096)
def qsfrompcap(pcapname):
pass
def is_blacklisted(msg):
if len(msg.question) >= 1:
......@@ -33,6 +30,7 @@ def is_blacklisted(msg):
return True
return False
def main():
qry = qfromtext(sys.argv[1:])
if is_blacklisted(qry):
......
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