Small optimisation of feed_switch

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],