web/collect/client/DNS/Base.py

   1 """
   2 $Id: Base.py,v 1.12.2.15 2011/03/19 22:15:01 customdesigned Exp $
   3
   4 This file is part of the pydns project.
   5 Homepage: http://pydns.sourceforge.net
   6
   7 This code is covered by the standard Python License.  See LICENSE for details.
   8
   9     Base functionality. Request and Response classes, that sort of thing.
  10 """
  11
  12 import socket, string, types, time, select
  13 import Type,Class,Opcode
  14 import asyncore
  15 #
  16 # This random generator is used for transaction ids and port selection.  This
  17 # is important to prevent spurious results from lost packets, and malicious
  18 # cache poisoning.  This doesn't matter if you are behind a caching nameserver
  19 # or your app is a primary DNS server only. To install your own generator,
  20 # replace DNS.Base.random.  SystemRandom uses /dev/urandom or similar source.
  21 #
  22 try:
  23   from random import SystemRandom
  24   random = SystemRandom()
  25 except:
  26   import random
  27
  28 class DNSError(Exception): pass
  29
  30 # Lib uses DNSError, so import after defining.
  31 import Lib
  32
  33 defaults= { 'protocol':'udp', 'port':53, 'opcode':Opcode.QUERY,
  34             'qtype':Type.A, 'rd':1, 'timing':1, 'timeout': 30,
  35             'server_rotate': 0 }
  36
  37 defaults['server']=[]
  38
  39 def ParseResolvConf(resolv_path="/etc/resolv.conf"):
  40     "parses the /etc/resolv.conf file and sets defaults for name servers"
  41     global defaults
  42     lines=open(resolv_path).readlines()
  43     for line in lines:
  44         line = string.strip(line)
  45         if not line or line[0]==';' or line[0]=='#':
  46             continue
  47         fields=string.split(line)
  48         if len(fields) < 2:
  49             continue
  50         if fields[0]=='domain' and len(fields) > 1:
  51             defaults['domain']=fields[1]
  52         if fields[0]=='search':
  53             pass
  54         if fields[0]=='options':
  55             pass
  56         if fields[0]=='sortlist':
  57             pass
  58         if fields[0]=='nameserver':
  59             defaults['server'].append(fields[1])
  60
  61 def DiscoverNameServers():
  62     import sys
  63     if sys.platform in ('win32', 'nt'):
  64         import win32dns
  65         defaults['server']=win32dns.RegistryResolve()
  66     else:
  67         return ParseResolvConf()
  68
  69 class DnsRequest:
  70     """ high level Request object """
  71     def __init__(self,*name,**args):
  72         self.donefunc=None
  73         self.async=None
  74         self.defaults = {}
  75         self.argparse(name,args)
  76         self.defaults = self.args
  77         self.tid = 0
  78
  79     def argparse(self,name,args):
  80         if not name and self.defaults.has_key('name'):
  81             args['name'] = self.defaults['name']
  82         if type(name) is types.StringType:
  83             args['name']=name
  84         else:
  85             if len(name) == 1:
  86                 if name[0]:
  87                     args['name']=name[0]
  88         if defaults['server_rotate'] and \
  89                 type(defaults['server']) == types.ListType:
  90             defaults['server'] = defaults['server'][1:]+defaults['server'][:1]
  91         for i in defaults.keys():
  92             if not args.has_key(i):
  93                 if self.defaults.has_key(i):
  94                     args[i]=self.defaults[i]
  95                 else:
  96                     args[i]=defaults[i]
  97         if type(args['server']) == types.StringType:
  98             args['server'] = [args['server']]
  99         self.args=args
 100
 101     def socketInit(self,a,b):
 102         self.s = socket.socket(a,b)
 103
 104     def processUDPReply(self):
 105         if self.timeout > 0:
 106             r,w,e = select.select([self.s],[],[],self.timeout)
 107             if not len(r):
 108                 raise DNSError, 'Timeout'
 109         (self.reply, self.from_address) = self.s.recvfrom(65535)
 110         self.time_finish=time.time()
 111         self.args['server']=self.ns
 112         return self.processReply()
 113
 114     def _readall(self,f,count):
 115       res = f.read(count)
 116       while len(res) < count:
 117         if self.timeout > 0:
 118             # should we restart timeout everytime we get a dribble of data?
 119             rem = self.time_start + self.timeout - time.time()
 120             if rem <= 0: raise DNSError,'Timeout'
 121             self.s.settimeout(rem)
 122         buf = f.read(count - len(res))
 123         if not buf:
 124           raise DNSError,'incomplete reply - %d of %d read' % (len(res),count)
 125         res += buf
 126       return res
 127
 128     def processTCPReply(self):
 129         if self.timeout > 0:
 130             self.s.settimeout(self.timeout)
 131         else:
 132             self.s.settimeout(None)
 133         f = self.s.makefile('r')
 134         header = self._readall(f,2)
 135         count = Lib.unpack16bit(header)
 136         self.reply = self._readall(f,count)
 137         self.time_finish=time.time()
 138         self.args['server']=self.ns
 139         return self.processReply()
 140
 141     def processReply(self):
 142         self.args['elapsed']=(self.time_finish-self.time_start)*1000
 143         u = Lib.Munpacker(self.reply)
 144         r=Lib.DnsResult(u,self.args)
 145         r.args=self.args
 146         #self.args=None  # mark this DnsRequest object as used.
 147         return r
 148         #### TODO TODO TODO ####
 149 #        if protocol == 'tcp' and qtype == Type.AXFR:
 150 #            while 1:
 151 #                header = f.read(2)
 152 #                if len(header) < 2:
 153 #                    print '========== EOF =========='
 154 #                    break
 155 #                count = Lib.unpack16bit(header)
 156 #                if not count:
 157 #                    print '========== ZERO COUNT =========='
 158 #                    break
 159 #                print '========== NEXT =========='
 160 #                reply = f.read(count)
 161 #                if len(reply) != count:
 162 #                    print '*** Incomplete reply ***'
 163 #                    break
 164 #                u = Lib.Munpacker(reply)
 165 #                Lib.dumpM(u)
 166
 167     def getSource(self):
 168         "Pick random source port to avoid DNS cache poisoning attack."
 169         while True:
 170             try:
 171                 source_port = random.randint(1024,65535)
 172                 self.s.bind(('', source_port))
 173                 break
 174             except socket.error, msg:
 175                 # Error 98, 'Address already in use'
 176                 if msg[0] != 98: raise
 177
 178     def conn(self):
 179         self.getSource()
 180         self.s.connect((self.ns,self.port))
 181
 182     def req(self,*name,**args):
 183         " needs a refactoring "
 184         self.argparse(name,args)
 185         #if not self.args:
 186         #    raise DNSError,'reinitialize request before reuse'
 187         protocol = self.args['protocol']
 188         self.port = self.args['port']
 189         self.tid = random.randint(0,65535)
 190         self.timeout = self.args['timeout'];
 191         opcode = self.args['opcode']
 192         rd = self.args['rd']
 193         server=self.args['server']
 194         if type(self.args['qtype']) == types.StringType:
 195             try:
 196                 qtype = getattr(Type, string.upper(self.args['qtype']))
 197             except AttributeError:
 198                 raise DNSError,'unknown query type'
 199         else:
 200             qtype=self.args['qtype']
 201         if not self.args.has_key('name'):
 202             print self.args
 203             raise DNSError,'nothing to lookup'
 204         qname = self.args['name']
 205         if qtype == Type.AXFR:
 206             print 'Query type AXFR, protocol forced to TCP'
 207             protocol = 'tcp'
 208         #print 'QTYPE %d(%s)' % (qtype, Type.typestr(qtype))
 209         m = Lib.Mpacker()
 210         # jesus. keywords and default args would be good. TODO.
 211         m.addHeader(self.tid,
 212               0, opcode, 0, 0, rd, 0, 0, 0,
 213               1, 0, 0, 0)
 214         m.addQuestion(qname, qtype, Class.IN)
 215         self.request = m.getbuf()
 216         try:
 217             if protocol == 'udp':
 218                 self.sendUDPRequest(server)
 219             else:
 220                 self.sendTCPRequest(server)
 221         except socket.error, reason:
 222             raise DNSError, reason
 223         if self.async:
 224             return None
 225         else:
 226             if not self.response:
 227                 raise DNSError,'no working nameservers found'
 228             return self.response
 229
 230     def sendUDPRequest(self, server):
 231         "refactor me"
 232         self.response=None
 233         for self.ns in server:
 234             #print "trying udp",self.ns
 235             try:
 236                 if self.ns.count(':'):
 237                     if hasattr(socket,'has_ipv6') and socket.has_ipv6:
 238                         self.socketInit(socket.AF_INET6, socket.SOCK_DGRAM)
 239                     else: continue
 240                 else:
 241                     self.socketInit(socket.AF_INET, socket.SOCK_DGRAM)
 242                 try:
 243                     # TODO. Handle timeouts &c correctly (RFC)
 244                     self.time_start=time.time()
 245                     self.conn()
 246                     if not self.async:
 247                         self.s.send(self.request)
 248                         r=self.processUDPReply()
 249                         # Since we bind to the source port and connect to the
 250                         # destination port, we don't need to check that here,
 251                         # but do make sure it's actually a DNS request that the
 252                         # packet is in reply to.
 253                         while r.header['id'] != self.tid        \
 254                                 or self.from_address[1] != self.port:
 255                             r=self.processUDPReply()
 256                         self.response = r
 257                         # FIXME: check waiting async queries
 258                 finally:
 259                     if not self.async:
 260                         self.s.close()
 261             except socket.error:
 262                 continue
 263             break
 264
 265     def sendTCPRequest(self, server):
 266         " do the work of sending a TCP request "
 267         self.response=None
 268         for self.ns in server:
 269             #print "trying tcp",self.ns
 270             try:
 271                 if self.ns.count(':'):
 272                     if hasattr(socket,'has_ipv6') and socket.has_ipv6:
 273                         self.socketInit(socket.AF_INET6, socket.SOCK_STREAM)
 274                     else: continue
 275                 else:
 276                     self.socketInit(socket.AF_INET, socket.SOCK_STREAM)
 277                 try:
 278                     # TODO. Handle timeouts &c correctly (RFC)
 279                     self.time_start=time.time()
 280                     self.conn()
 281                     buf = Lib.pack16bit(len(self.request))+self.request
 282                     # Keep server from making sendall hang
 283                     self.s.setblocking(0)
 284                     # FIXME: throws WOULDBLOCK if request too large to fit in
 285                     # system buffer
 286                     self.s.sendall(buf)
 287                     # SHUT_WR breaks blocking IO with google DNS (8.8.8.8)
 288                     #self.s.shutdown(socket.SHUT_WR)
 289                     r=self.processTCPReply()
 290                     if r.header['id'] == self.tid:
 291                         self.response = r
 292                         break
 293                 finally:
 294                     self.s.close()
 295             except socket.error:
 296                 continue
 297
 298 #class DnsAsyncRequest(DnsRequest):
 299 class DnsAsyncRequest(DnsRequest,asyncore.dispatcher_with_send):
 300     " an asynchronous request object. out of date, probably broken "
 301     def __init__(self,*name,**args):
 302         DnsRequest.__init__(self, *name, **args)
 303         # XXX todo
 304         if args.has_key('done') and args['done']:
 305             self.donefunc=args['done']
 306         else:
 307             self.donefunc=self.showResult
 308         #self.realinit(name,args) # XXX todo
 309         self.async=1
 310     def conn(self):
 311         self.getSource()
 312         self.connect((self.ns,self.port))
 313         self.time_start=time.time()
 314         if self.args.has_key('start') and self.args['start']:
 315             asyncore.dispatcher.go(self)
 316     def socketInit(self,a,b):
 317         self.create_socket(a,b)
 318         asyncore.dispatcher.__init__(self)
 319         self.s=self
 320     def handle_read(self):
 321         if self.args['protocol'] == 'udp':
 322             self.response=self.processUDPReply()
 323             if self.donefunc:
 324                 apply(self.donefunc,(self,))
 325     def handle_connect(self):
 326         self.send(self.request)
 327     def handle_write(self):
 328         pass
 329     def showResult(self,*s):
 330         self.response.show()
 331
 332 #
 333 # $Log: Base.py,v $
 334 # Revision 1.12.2.15  2011/03/19 22:15:01  customdesigned
 335 # Added rotation of name servers - SF Patch ID: 2795929
 336 #
 337 # Revision 1.12.2.14  2011/03/17 03:46:03  customdesigned
 338 # Simple test for google DNS with tcp
 339 #
 340 # Revision 1.12.2.13  2011/03/17 03:08:03  customdesigned
 341 # Use blocking IO with timeout for TCP replies.
 342 #
 343 # Revision 1.12.2.12  2011/03/16 17:50:00  customdesigned
 344 # Fix non-blocking TCP replies.  (untested)
 345 #
 346 # Revision 1.12.2.11  2010/01/02 16:31:23  customdesigned
 347 # Handle large TCP replies (untested).
 348 #
 349 # Revision 1.12.2.10  2008/08/01 03:58:03  customdesigned
 350 # Don't try to close socket when never opened.
 351 #
 352 # Revision 1.12.2.9  2008/08/01 03:48:31  customdesigned
 353 # Fix more breakage from port randomization patch.  Support Ipv6 queries.
 354 #
 355 # Revision 1.12.2.8  2008/07/31 18:22:59  customdesigned
 356 # Wait until tcp response at least starts coming in.
 357 #
 358 # Revision 1.12.2.7  2008/07/28 01:27:00  customdesigned
 359 # Check configured port.
 360 #
 361 # Revision 1.12.2.6  2008/07/28 00:17:10  customdesigned
 362 # Randomize source ports.
 363 #
 364 # Revision 1.12.2.5  2008/07/24 20:10:55  customdesigned
 365 # Randomize tid in requests, and check in response.
 366 #
 367 # Revision 1.12.2.4  2007/05/22 20:28:31  customdesigned
 368 # Missing import Lib
 369 #
 370 # Revision 1.12.2.3  2007/05/22 20:25:52  customdesigned
 371 # Use socket.inetntoa,inetaton.
 372 #
 373 # Revision 1.12.2.2  2007/05/22 20:21:46  customdesigned
 374 # Trap socket error
 375 #
 376 # Revision 1.12.2.1  2007/05/22 20:19:35  customdesigned
 377 # Skip bogus but non-empty lines in resolv.conf
 378 #
 379 # Revision 1.12  2002/04/23 06:04:27  anthonybaxter
 380 # attempt to refactor the DNSRequest.req method a little. after doing a bit
 381 # of this, I've decided to bite the bullet and just rewrite the puppy. will
 382 # be checkin in some design notes, then unit tests and then writing the sod.
 383 #
 384 # Revision 1.11  2002/03/19 13:05:02  anthonybaxter
 385 # converted to class based exceptions (there goes the python1.4 compatibility :)
 386 #
 387 # removed a quite gross use of 'eval()'.
 388 #
 389 # Revision 1.10  2002/03/19 12:41:33  anthonybaxter
 390 # tabnannied and reindented everything. 4 space indent, no tabs.
 391 # yay.
 392 #
 393 # Revision 1.9  2002/03/19 12:26:13  anthonybaxter
 394 # death to leading tabs.
 395 #
 396 # Revision 1.8  2002/03/19 10:30:33  anthonybaxter
 397 # first round of major bits and pieces. The major stuff here (summarised
 398 # from my local, off-net CVS server :/ this will cause some oddities with
 399 # the
 400 #
 401 # tests/testPackers.py:
 402 #   a large slab of unit tests for the packer and unpacker code in DNS.Lib
 403 #
 404 # DNS/Lib.py:
 405 #   placeholder for addSRV.
 406 #   added 'klass' to addA, make it the same as the other A* records.
 407 #   made addTXT check for being passed a string, turn it into a length 1 list.
 408 #   explicitly check for adding a string of length > 255 (prohibited).
 409 #   a bunch of cleanups from a first pass with pychecker
 410 #   new code for pack/unpack. the bitwise stuff uses struct, for a smallish
 411 #     (disappointly small, actually) improvement, while addr2bin is much
 412 #     much faster now.
 413 #
 414 # DNS/Base.py:
 415 #   added DiscoverNameServers. This automatically does the right thing
 416 #     on unix/ win32. No idea how MacOS handles this.  *sigh*
 417 #     Incompatible change: Don't use ParseResolvConf on non-unix, use this
 418 #     function, instead!
 419 #   a bunch of cleanups from a first pass with pychecker
 420 #
 421 # Revision 1.5  2001/08/09 09:22:28  anthonybaxter
 422 # added what I hope is win32 resolver lookup support. I'll need to try
 423 # and figure out how to get the CVS checkout onto my windows machine to
 424 # make sure it works (wow, doing something other than games on the
 425 # windows machine :)
 426 #
 427 # Code from Wolfgang.Strobl@gmd.de
 428 # win32dns.py from
 429 # http://aspn.activestate.com/ASPN/Cookbook/Python/Recipe/66260
 430 #
 431 # Really, ParseResolvConf() should be renamed "FindNameServers" or
 432 # some such.
 433 #
 434 # Revision 1.4  2001/08/09 09:08:55  anthonybaxter
 435 # added identifying header to top of each file
 436 #
 437 # Revision 1.3  2001/07/19 07:20:12  anthony
 438 # Handle blank resolv.conf lines.
 439 # Patch from Bastian Kleineidam
 440 #
 441 # Revision 1.2  2001/07/19 06:57:07  anthony
 442 # cvs keywords added
 443 #
 444 #