diff --git a/feed_dns.py b/feed_dns.py index c2438d8..d72dc49 100755 --- a/feed_dns.py +++ b/feed_dns.py @@ -8,20 +8,19 @@ import typing import multiprocessing import enum -RecordType = enum.Enum('RecordType', 'A AAAA CNAME PTR') -Record = typing.Tuple[RecordType, int, str, str] +Record = typing.Tuple[typing.Callable, typing.Callable, int, str, str] # select, write FUNCTION_MAP: typing.Any = { - RecordType.A: ( + 'a': ( database.Database.get_ip4, database.Database.set_hostname, ), - RecordType.CNAME: ( + 'cname': ( database.Database.get_domain, database.Database.set_hostname, ), - RecordType.PTR: ( + 'ptr': ( database.Database.get_domain, database.Database.set_ip4address, ), @@ -47,10 +46,9 @@ class Writer(multiprocessing.Process): record: Record for record in block: - rtype, updated, name, value = record + select, write, updated, name, value = record self.db.enter_step('feed_switch') - select, write = FUNCTION_MAP[rtype] for source in select(self.db, value): # write(self.db, name, updated, source=source) write(self.db, name, updated) @@ -94,13 +92,6 @@ class Parser(): class Rapid7Parser(Parser): - TYPES = { - 'a': RecordType.A, - 'aaaa': RecordType.AAAA, - 'cname': RecordType.CNAME, - 'ptr': RecordType.PTR, - } - def consume(self) -> None: data = dict() for line in self.buf: @@ -112,8 +103,10 @@ class Rapid7Parser(Parser): val = split[k+2] data[key] = val + select, writer = FUNCTION_MAP[data['type']] record = ( - Rapid7Parser.TYPES[data['type']], + select, + writer, int(data['timestamp']), data['name'], data['value'] @@ -125,9 +118,9 @@ 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), + 'A': (FUNCTION_MAP['a'][0], FUNCTION_MAP['a'][1], -1, None), + # 'AAAA': (FUNCTION_MAP['aaaa'][0], FUNCTION_MAP['aaaa'][1], -1, None), + 'CNAME': (FUNCTION_MAP['cname'][0], FUNCTION_MAP['cname'][1], -1, -1), } def consume(self) -> None: @@ -146,10 +139,11 @@ class DnsMassParser(Parser): timestamp = int(split[1]) header = False else: - dtype, name_offset, value_offset = \ + select, write, name_offset, value_offset = \ DnsMassParser.TYPES[split[1]] record = ( - dtype, + select, + write, timestamp, split[0][:name_offset], split[2][:value_offset],