MoinQ: DNS/実装/python/dnslib/sample-resolver/proxy

   1 # -*- coding: utf-8 -*-
   2 
   3 from __future__ import print_function
   4 
   5 import binascii,socket,struct
   6 
   7 from dnslib import DNSRecord,RCODE
   8 from dnslib.server import DNSServer,DNSHandler,BaseResolver,DNSLogger
   9 
  10 class ProxyResolver(BaseResolver):
  11     """
  12         Proxy resolver - passes all requests to upstream DNS server and
  13         returns response
  14 
  15         Note that the request/response will be each be decoded/re-encoded 
  16         twice:
  17 
  18         a) Request packet received by DNSHandler and parsed into DNSRecord 
  19         b) DNSRecord passed to ProxyResolver, serialised back into packet 
  20            and sent to upstream DNS server
  21         c) Upstream DNS server returns response packet which is parsed into
  22            DNSRecord
  23         d) ProxyResolver returns DNSRecord to DNSHandler which re-serialises
  24            this into packet and returns to client
  25 
  26         In practice this is actually fairly useful for testing but for a
  27         'real' transparent proxy option the DNSHandler logic needs to be
  28         modified (see PassthroughDNSHandler)
  29 
  30     """
  31 
  32     def __init__(self,address,port,timeout=0):
  33         self.address = address
  34         self.port = port
  35         self.timeout = timeout
  36 
  37     def resolve(self,request,handler):
  38         try:
  39             if handler.protocol == 'udp':
  40                 proxy_r = request.send(self.address,self.port,timeout=self.timeout)
  41             else:
  42                 proxy_r = request.send(self.address,self.port,tcp=True,timeout=self.timeout)
  43             reply = DNSRecord.parse(proxy_r)
  44         except socket.timeout:
  45             reply = request.reply()
  46             reply.header.rcode = getattr(RCODE,'NXDOMAIN')
  47 
  48         return reply
  49 
  50 class PassthroughDNSHandler(DNSHandler):
  51     """
  52         Modify DNSHandler logic (get_reply method) to send directly to 
  53         upstream DNS server rather then decoding/encoding packet and
  54         passing to Resolver (The request/response packets are still
  55         parsed and logged but this is not inline)
  56     """
  57     def get_reply(self,data):
  58         host = self.server.resolver.address
  59         port = self.server.resolver.port
  60 
  61         request = DNSRecord.parse(data)
  62         self.log_request(request)
  63 
  64         if self.protocol == 'tcp':
  65             data = struct.pack("!H",len(data)) + data
  66             response = send_tcp(data,host,port)
  67             response = response[2:]
  68         else:
  69             response = send_udp(data,host,port)
  70 
  71         reply = DNSRecord.parse(response)
  72         self.log_reply(reply)
  73 
  74         return response
  75 
  76 def send_tcp(data,host,port):
  77     """
  78         Helper function to send/receive DNS TCP request
  79         (in/out packets will have prepended TCP length header)
  80     """
  81     sock = socket.socket(socket.AF_INET,socket.SOCK_STREAM)
  82     sock.connect((host,port))
  83     sock.sendall(data)
  84     response = sock.recv(8192)
  85     length = struct.unpack("!H",bytes(response[:2]))[0]
  86     while len(response) - 2 < length:
  87         response += sock.recv(8192)
  88     sock.close()
  89     return response
  90 
  91 def send_udp(data,host,port):
  92     """
  93         Helper function to send/receive DNS UDP request
  94     """
  95     sock = socket.socket(socket.AF_INET,socket.SOCK_DGRAM)
  96     sock.sendto(data,(host,port))
  97     response,server = sock.recvfrom(8192)
  98     sock.close()
  99     return response
 100 
 101 if __name__ == '__main__':
 102 
 103     import argparse,sys,time
 104 
 105     p = argparse.ArgumentParser(description="DNS Proxy")
 106     p.add_argument("--port","-p",type=int,default=8053, metavar="<port>",
 107                     help="Local proxy port (default:8053)")
 108     p.add_argument("--address","-a",default="", metavar="<address>",
 109                     help="Local proxy listen address (default:all)")
 110     p.add_argument("--upstream","-u",default="8.8.8.8:53", metavar="<dns server:port>",
 111                     help="Upstream DNS server:port (default:8.8.8.8:53)")
 112     p.add_argument("--tcp",action='store_true',default=False,
 113                     help="TCP proxy (default: UDP only)")
 114     p.add_argument("--timeout","-o",type=float,default=5, metavar="<timeout>",
 115                     help="Upstream timeout (default: 5s)")
 116     p.add_argument("--passthrough",action='store_true',default=False,
 117                     help="Dont decode/re-encode request/response (default: off)")
 118     p.add_argument("--log",default="request,reply,truncated,error",
 119                     help="Log hooks to enable (default: +request,+reply,+truncated,+error,-recv,-send,-data)")
 120     p.add_argument("--log-prefix",action='store_true',default=False,
 121                     help="Log prefix (timestamp/handler/resolver) (default: False)")
 122     args = p.parse_args()
 123 
 124     args.dns,_,args.dns_port = args.upstream.partition(':')
 125     args.dns_port = int(args.dns_port or 53)
 126 
 127     print("Starting Proxy Resolver (%s:%d -> %s:%d) [%s]" % (
 128                         args.address or "*",args.port,
 129                         args.dns,args.dns_port,
 130                         "UDP/TCP" if args.tcp else "UDP"))
 131 
 132     resolver = ProxyResolver(args.dns,args.dns_port,args.timeout)
 133     handler = PassthroughDNSHandler if args.passthrough else DNSHandler
 134     logger = DNSLogger(args.log,args.log_prefix)
 135     udp_server = DNSServer(resolver,
 136                            port=args.port,
 137                            address=args.address,
 138                            logger=logger,
 139                            handler=handler)
 140     udp_server.start_thread()
 141 
 142     if args.tcp:
 143         tcp_server = DNSServer(resolver,
 144                                port=args.port,
 145                                address=args.address,
 146                                tcp=True,
 147                                logger=logger,
 148                                handler=handler)
 149         tcp_server.start_thread()
 150 
 151     while udp_server.isAlive():
 152         time.sleep(1)