from dnslib import RR, QTYPE, A


class DatabaseLookupResolver:
    def __init__(self, storage, zone, keymaker=None, ttl=60):
        self.storage = storage
        self.zone = zone
        self.keymaker = keymaker
        self.ttl = ttl

    def resolve(self, request, handler):
        reply = request.reply()
        key = self.keymaker(request) if self.keymaker else str(request.q.qname)
        address = self.storage.get(key)
        if address is not None:
            reply.add_answer(
                RR(self.zone, QTYPE.A, rdata=A(address), ttl=self.ttl)
            )
        return reply


def main():
    import argparse
    import json
    import time

    from dnslib.server import DNSLogger, DNSServer

    from labDNS.storages import RedisStorage, DictStorage

    parser = argparse.ArgumentParser(
        description="Database lookup based DNS resolver"
    )

    parser.add_argument(
        "--storage", "-s", default='dict', choices=['redis', 'dict']
    )
    parser.add_argument("--config", "-c", type=json.loads, default=dict())
    parser.add_argument("--zone", "-z")
    parser.add_argument("--ttl", "-t", default=10, type=int)
    parser.add_argument("--log", "-l", default="request,reply,truncated,error")
    parser.add_argument("--port", "-p", default=53, type=int)
    parser.add_argument("--address", "-a", default="localhost")

    args = parser.parse_args()

    config = args.config

    if args.storage == 'redis':
        storage = RedisStorage(config)
    elif args.storage == 'dict':
        storage = DictStorage(config)

    resolver = DatabaseLookupResolver(storage, args.zone, ttl=args.ttl)
    logger = DNSLogger(args.log)

    server = DNSServer(
        resolver,
        port=args.port,
        address=args.address,
        logger=logger,
    )

    server.start_thread()

    while server.isAlive():
        time.sleep(1)

if __name__ == '__main__':
    main()

@@ -318,16 +318,30 @@
    key =
+ (%0A           
  self.ke
@@ -354,16 +354,25 @@
 (request
+, handler
 ) if sel
@@ -381,16 +381,28 @@
 keymaker
+%0A           
  else st
@@ -420,16 +420,26 @@
 .qname)%0A
+        )%0A
         
@@ -703,16 +703,56 @@
 ort time
+%0A    from importlib import import_module
 %0A%0A    fr
@@ -1439,16 +1439,74 @@
 alhost%22)
+%0A    parser.add_argument(%22--keymaker%22, %22-k%22, default=None)
 %0A%0A    ar
@@ -1696,24 +1696,105 @@
 ge(config)%0A%0A
+    keymaker = import_module(args.keymaker).keymaker if args.keymaker else None%0A%0A
     resolver
@@ -1819,16 +1819,25 @@
 esolver(
+%0A        
 storage,
@@ -1856,24 +1856,48 @@
 ttl=args.ttl
+, keymaker=keymaker%0A    
 )%0A    logger

allow specifying keymaker as a module