Workflow: Can now import DnsMass output
Well, in a specific format but DnsMass nonetheless
This commit is contained in:
parent
189deeb559
commit
ddceed3d25
GPG key ID:
D8A7ECA00A8CD3DD
5 changed files with 152 additions and 345 deletions
185
feed_dns.py
185
feed_dns.py
|
|
@ -7,23 +7,24 @@ import logging
|
|||
import sys
|
||||
import typing
|
||||
import multiprocessing
|
||||
import enum
|
||||
|
||||
NUMBER_THREADS = 2
|
||||
BLOCK_SIZE = 100
|
||||
RecordType = enum.Enum('RecordType', 'A AAAA CNAME PTR')
|
||||
Record = typing.Tuple[RecordType, int, str, str]
|
||||
|
||||
# select, confirm, write
|
||||
FUNCTION_MAP: typing.Any = {
|
||||
'a': (
|
||||
RecordType.A: (
|
||||
database.Database.get_ip4,
|
||||
database.Database.get_domain_in_zone,
|
||||
database.Database.set_hostname,
|
||||
),
|
||||
'cname': (
|
||||
RecordType.CNAME: (
|
||||
database.Database.get_domain,
|
||||
database.Database.get_domain_in_zone,
|
||||
database.Database.set_hostname,
|
||||
),
|
||||
'ptr': (
|
||||
RecordType.PTR: (
|
||||
database.Database.get_domain,
|
||||
database.Database.get_ip4_in_network,
|
||||
database.Database.set_ip4address,
|
||||
|
|
@ -33,12 +34,12 @@ FUNCTION_MAP: typing.Any = {
|
|||
|
||||
class Reader(multiprocessing.Process):
|
||||
def __init__(self,
|
||||
lines_queue: multiprocessing.Queue,
|
||||
recs_queue: multiprocessing.Queue,
|
||||
write_queue: multiprocessing.Queue,
|
||||
index: int = 0):
|
||||
super(Reader, self).__init__()
|
||||
self.log = logging.getLogger(f'rd{index:03d}')
|
||||
self.lines_queue = lines_queue
|
||||
self.recs_queue = recs_queue
|
||||
self.write_queue = write_queue
|
||||
self.index = index
|
||||
|
||||
|
|
@ -48,15 +49,19 @@ class Reader(multiprocessing.Process):
|
|||
self.db.enter_step('line_wait')
|
||||
block: typing.List[str]
|
||||
try:
|
||||
for block in iter(self.lines_queue.get, None):
|
||||
for line in block:
|
||||
dtype, updated, name, value = line
|
||||
for block in iter(self.recs_queue.get, None):
|
||||
record: Record
|
||||
for record in block:
|
||||
# print(55, record)
|
||||
dtype, updated, name, value = record
|
||||
self.db.enter_step('feed_switch')
|
||||
select, confirm, write = FUNCTION_MAP[dtype]
|
||||
for rule in select(self.db, value):
|
||||
# print(60, rule, list(confirm(self.db, name)))
|
||||
if not any(confirm(self.db, name)):
|
||||
# print(62, write, name, updated, rule)
|
||||
self.db.enter_step('wait_put')
|
||||
self.write_queue.put((write, name, updated))
|
||||
self.write_queue.put((write, name, updated, rule))
|
||||
self.db.enter_step('line_wait')
|
||||
except KeyboardInterrupt:
|
||||
self.log.error('Interrupted')
|
||||
|
|
@ -82,9 +87,10 @@ class Writer(multiprocessing.Process):
|
|||
fun: typing.Callable
|
||||
name: str
|
||||
updated: int
|
||||
for fun, name, updated in iter(self.write_queue.get, None):
|
||||
source: int
|
||||
for fun, name, updated, source in iter(self.write_queue.get, None):
|
||||
self.db.enter_step('exec')
|
||||
fun(self.db, name, updated)
|
||||
fun(self.db, name, updated, source=source)
|
||||
self.db.enter_step('line_wait')
|
||||
except KeyboardInterrupt:
|
||||
self.log.error('Interrupted')
|
||||
|
|
@ -93,29 +99,142 @@ class Writer(multiprocessing.Process):
|
|||
self.db.close()
|
||||
|
||||
|
||||
class Parser():
|
||||
def __init__(self,
|
||||
buf: typing.Any,
|
||||
recs_queue: multiprocessing.Queue,
|
||||
block_size: int,
|
||||
):
|
||||
super(Parser, self).__init__()
|
||||
self.buf = buf
|
||||
self.log = logging.getLogger('pr ')
|
||||
self.recs_queue = recs_queue
|
||||
self.block: typing.List[Record] = list()
|
||||
self.block_size = block_size
|
||||
self.db = database.Database() # Just for timing
|
||||
self.db.log = logging.getLogger('pr ')
|
||||
|
||||
def register(self, record: Record) -> None:
|
||||
self.db.enter_step('register')
|
||||
self.block.append(record)
|
||||
if len(self.block) >= self.block_size:
|
||||
self.db.enter_step('put_block')
|
||||
self.recs_queue.put(self.block)
|
||||
self.block = list()
|
||||
|
||||
def run(self) -> None:
|
||||
self.consume()
|
||||
self.recs_queue.put(self.block)
|
||||
self.db.close()
|
||||
|
||||
def consume(self) -> None:
|
||||
raise NotImplementedError
|
||||
|
||||
|
||||
class Rapid7Parser(Parser):
|
||||
TYPES = {
|
||||
'a': RecordType.A,
|
||||
'aaaa': RecordType.AAAA,
|
||||
'cname': RecordType.CNAME,
|
||||
'ptr': RecordType.PTR,
|
||||
}
|
||||
|
||||
def consume(self) -> None:
|
||||
for line in self.buf:
|
||||
self.db.enter_step('parse_rapid7')
|
||||
try:
|
||||
data = json.loads(line)
|
||||
except json.decoder.JSONDecodeError:
|
||||
continue
|
||||
record = (
|
||||
Rapid7Parser.TYPES[data['type']],
|
||||
int(data['timestamp']),
|
||||
data['name'],
|
||||
data['value']
|
||||
)
|
||||
self.register(record)
|
||||
|
||||
|
||||
class DnsMassParser(Parser):
|
||||
# dnsmass --output Snrql
|
||||
# --retry REFUSED,SERVFAIL --resolvers nameservers-ipv4
|
||||
TYPES = {
|
||||
'A': (RecordType.A, -1, None),
|
||||
'AAAA': (RecordType.AAAA, -1, None),
|
||||
'CNAME': (RecordType.CNAME, -1, -1),
|
||||
}
|
||||
|
||||
def consume(self) -> None:
|
||||
self.db.enter_step('parse_dnsmass')
|
||||
timestamp = 0
|
||||
header = True
|
||||
for line in self.buf:
|
||||
line = line[:-1]
|
||||
if not line:
|
||||
header = True
|
||||
continue
|
||||
|
||||
split = line.split(' ')
|
||||
try:
|
||||
if header:
|
||||
timestamp = int(split[1])
|
||||
header = False
|
||||
else:
|
||||
dtype, name_offset, value_offset = \
|
||||
DnsMassParser.TYPES[split[1]]
|
||||
record = (
|
||||
dtype,
|
||||
timestamp,
|
||||
split[0][:name_offset],
|
||||
split[2][:value_offset],
|
||||
)
|
||||
self.register(record)
|
||||
self.db.enter_step('parse_dnsmass')
|
||||
except KeyError:
|
||||
continue
|
||||
|
||||
|
||||
PARSERS = {
|
||||
'rapid7': Rapid7Parser,
|
||||
'dnsmass': DnsMassParser,
|
||||
}
|
||||
|
||||
if __name__ == '__main__':
|
||||
|
||||
# Parsing arguments
|
||||
log = logging.getLogger('feed_dns')
|
||||
parser = argparse.ArgumentParser(
|
||||
args_parser = argparse.ArgumentParser(
|
||||
description="TODO")
|
||||
parser.add_argument(
|
||||
# '-i', '--input', type=argparse.FileType('rb'), default=sys.stdin.buffer,
|
||||
args_parser.add_argument(
|
||||
'parser',
|
||||
choices=PARSERS.keys(),
|
||||
help="TODO")
|
||||
args_parser.add_argument(
|
||||
'-i', '--input', type=argparse.FileType('r'), default=sys.stdin,
|
||||
help="TODO")
|
||||
args = parser.parse_args()
|
||||
args_parser.add_argument(
|
||||
'-j', '--workers', type=int, default=4,
|
||||
help="TODO")
|
||||
args_parser.add_argument(
|
||||
'-b', '--block-size', type=int, default=100,
|
||||
help="TODO")
|
||||
args = args_parser.parse_args()
|
||||
|
||||
DB = database.Database(write=False) # Not needed, just for timing
|
||||
DB.log = logging.getLogger('db ')
|
||||
|
||||
lines_queue: multiprocessing.Queue = multiprocessing.Queue(maxsize=100)
|
||||
write_queue: multiprocessing.Queue = multiprocessing.Queue(maxsize=100)
|
||||
recs_queue: multiprocessing.Queue = multiprocessing.Queue(
|
||||
maxsize=10*args.workers)
|
||||
write_queue: multiprocessing.Queue = multiprocessing.Queue(
|
||||
maxsize=10*args.workers)
|
||||
|
||||
DB.enter_step('proc_create')
|
||||
readers: typing.List[Reader] = list()
|
||||
for w in range(NUMBER_THREADS):
|
||||
readers.append(Reader(lines_queue, write_queue, w))
|
||||
for w in range(args.workers):
|
||||
readers.append(Reader(recs_queue, write_queue, w))
|
||||
writer = Writer(write_queue)
|
||||
parser = PARSERS[args.parser](
|
||||
args.input, recs_queue, args.block_size)
|
||||
|
||||
DB.enter_step('proc_start')
|
||||
for reader in readers:
|
||||
|
|
@ -123,28 +242,12 @@ if __name__ == '__main__':
|
|||
writer.start()
|
||||
|
||||
try:
|
||||
block: typing.List[str] = list()
|
||||
DB.enter_step('iowait')
|
||||
for line in args.input:
|
||||
DB.enter_step('block_append')
|
||||
DB.enter_step('feed_json_parse')
|
||||
data = json.loads(line)
|
||||
line = (data['type'],
|
||||
int(data['timestamp']),
|
||||
data['name'],
|
||||
data['value'])
|
||||
block.append(line)
|
||||
if len(block) >= BLOCK_SIZE:
|
||||
DB.enter_step('wait_put')
|
||||
lines_queue.put(block)
|
||||
block = list()
|
||||
DB.enter_step('iowait')
|
||||
DB.enter_step('wait_put')
|
||||
lines_queue.put(block)
|
||||
DB.enter_step('parser_run')
|
||||
parser.run()
|
||||
|
||||
DB.enter_step('end_put')
|
||||
for _ in range(NUMBER_THREADS):
|
||||
lines_queue.put(None)
|
||||
for _ in range(args.workers):
|
||||
recs_queue.put(None)
|
||||
write_queue.put(None)
|
||||
|
||||
DB.enter_step('proc_join')
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue