home *** CD-ROM | disk | FTP | other *** search
/ OS/2 Shareware BBS: 10 Tools / 10-Tools.zip / pyos2bin.zip / Demo / dns / dnslib.py < prev    next >
Text File  |  1996-07-30  |  15KB  |  589 lines

  1. # Domain Name Server (DNS) interface
  2. #
  3. # See RFC 1035:
  4. # ------------------------------------------------------------------------
  5. # Network Working Group                                     P. Mockapetris
  6. # Request for Comments: 1035                                           ISI
  7. #                                                            November 1987
  8. # Obsoletes: RFCs 882, 883, 973
  9. #             DOMAIN NAMES - IMPLEMENTATION AND SPECIFICATION
  10. # ------------------------------------------------------------------------
  11.  
  12.  
  13. import string
  14.  
  15. import dnstype
  16. import dnsclass
  17. import dnsopcode
  18.  
  19.  
  20. # Low-level 16 and 32 bit integer packing and unpacking
  21.  
  22. def pack16bit(n):
  23.     return chr((n>>8)&0xFF) + chr(n&0xFF)
  24.  
  25. def pack32bit(n):
  26.     return chr((n>>24)&0xFF) + chr((n>>16)&0xFF) \
  27.           + chr((n>>8)&0xFF) + chr(n&0xFF)
  28.  
  29. def unpack16bit(s):
  30.     return (ord(s[0])<<8) | ord(s[1])
  31.  
  32. def unpack32bit(s):
  33.     return (ord(s[0])<<24) | (ord(s[1])<<16) \
  34.           | (ord(s[2])<<8) | ord(s[3])
  35.  
  36. def addr2bin(addr):
  37.     if type(addr) == type(0):
  38.         return addr
  39.     bytes = string.splitfields(addr, '.')
  40.     if len(bytes) != 4: raise ValueError, 'bad IP address'
  41.     n = 0
  42.     for byte in bytes: n = n<<8 | string.atoi(byte)
  43.     return n
  44.  
  45. def bin2addr(n):
  46.     return '%d.%d.%d.%d' % ((n>>24)&0xFF, (n>>16)&0xFF,
  47.           (n>>8)&0xFF, n&0xFF)
  48.  
  49.  
  50. # Packing class
  51.  
  52. class Packer:
  53.     def __init__(self):
  54.         self.buf = ''
  55.         self.index = {}
  56.     def getbuf(self):
  57.         return self.buf
  58.     def addbyte(self, c):
  59.         if len(c) != 1: raise TypeError, 'one character expected'
  60.         self.buf = self.buf + c
  61.     def addbytes(self, bytes):
  62.         self.buf = self.buf + bytes
  63.     def add16bit(self, n):
  64.         self.buf = self.buf + pack16bit(n)
  65.     def add32bit(self, n):
  66.         self.buf = self.buf + pack32bit(n)
  67.     def addaddr(self, addr):
  68.         n = addr2bin(addr)
  69.         self.buf = self.buf + pack32bit(n)
  70.     def addstring(self, s):
  71.         self.addbyte(chr(len(s)))
  72.         self.addbytes(s)
  73.     def addname(self, name):
  74.         # Domain name packing (section 4.1.4)
  75.         # Add a domain name to the buffer, possibly using pointers.
  76.         # The case of the first occurrence of a name is preserved.
  77.         # Redundant dots are ignored.
  78.         list = []
  79.         for label in string.splitfields(name, '.'):
  80.             if label:
  81.                 if len(label) > 63:
  82.                     raise PackError, 'label too long'
  83.                 list.append(label)
  84.         keys = []
  85.         for i in range(len(list)):
  86.             key = string.upper(string.joinfields(list[i:], '.'))
  87.             keys.append(key)
  88.             if self.index.has_key(key):
  89.                 pointer = self.index[key]
  90.                 break
  91.         else:
  92.             i = len(list)
  93.             pointer = None
  94.         # Do it into temporaries first so exceptions don't
  95.         # mess up self.index and self.buf
  96.         buf = ''
  97.         offset = len(self.buf)
  98.         index = []
  99.         for j in range(i):
  100.             label = list[j]
  101.             n = len(label)
  102.             if offset + len(buf) < 0x3FFF:
  103.                 index.append(keys[j], offset + len(buf))
  104.             else:
  105.                 print 'dnslib.Packer.addname:',
  106.                 print 'warning: pointer too big'
  107.             buf = buf + (chr(n) + label)
  108.         if pointer:
  109.             buf = buf + pack16bit(pointer | 0xC000)
  110.         else:
  111.             buf = buf + '\0'
  112.         self.buf = self.buf + buf
  113.         for key, value in index:
  114.             self.index[key] = value
  115.     def dump(self):
  116.         keys = self.index.keys()
  117.         keys.sort()
  118.         print '-'*40
  119.         for key in keys:
  120.             print '%20s %3d' % (key, self.index[key])
  121.         print '-'*40
  122.         space = 1
  123.         for i in range(0, len(self.buf)+1, 2):
  124.             if self.buf[i:i+2] == '**':
  125.                 if not space: print
  126.                 space = 1
  127.                 continue
  128.             space = 0
  129.             print '%4d' % i,
  130.             for c in self.buf[i:i+2]:
  131.                 if ' ' < c < '\177':
  132.                     print ' %c' % c,
  133.                 else:
  134.                     print '%2d' % ord(c),
  135.             print
  136.         print '-'*40
  137.  
  138.  
  139. # Unpacking class
  140.  
  141. UnpackError = 'dnslib.UnpackError'    # Exception
  142.  
  143. class Unpacker:
  144.     def __init__(self, buf):
  145.         self.buf = buf
  146.         self.offset = 0
  147.     def getbyte(self):
  148.         c = self.buf[self.offset]
  149.         self.offset = self.offset + 1
  150.         return c
  151.     def getbytes(self, n):
  152.         s = self.buf[self.offset : self.offset + n]
  153.         if len(s) != n: raise UnpackError, 'not enough data left'
  154.         self.offset = self.offset + n
  155.         return s
  156.     def get16bit(self):
  157.         return unpack16bit(self.getbytes(2))
  158.     def get32bit(self):
  159.         return unpack32bit(self.getbytes(4))
  160.     def getaddr(self):
  161.         return bin2addr(self.get32bit())
  162.     def getstring(self):
  163.         return self.getbytes(ord(self.getbyte()))
  164.     def getname(self):
  165.         # Domain name unpacking (section 4.1.4)
  166.         c = self.getbyte()
  167.         i = ord(c)
  168.         if i & 0xC0 == 0xC0:
  169.             d = self.getbyte()
  170.             j = ord(d)
  171.             pointer = ((i<<8) | j) & ~0xC000
  172.             save_offset = self.offset
  173.             try:
  174.                 self.offset = pointer
  175.                 domain = self.getname()
  176.             finally:
  177.                 self.offset = save_offset
  178.             return domain
  179.         if i == 0:
  180.             return ''
  181.         domain = self.getbytes(i)
  182.         remains = self.getname()
  183.         if not remains:
  184.             return domain
  185.         else:
  186.             return domain + '.' + remains
  187.  
  188.  
  189. # Test program for packin/unpacking (section 4.1.4)
  190.  
  191. def testpacker():
  192.     N = 25
  193.     R = range(N)
  194.     import timing
  195.     # See section 4.1.4 of RFC 1035
  196.     timing.start()
  197.     for i in R:
  198.         p = Packer()
  199.         p.addbytes('*' * 20)
  200.         p.addname('f.ISI.ARPA')
  201.         p.addbytes('*' * 8)
  202.         p.addname('Foo.F.isi.arpa')
  203.         p.addbytes('*' * 18)
  204.         p.addname('arpa')
  205.         p.addbytes('*' * 26)
  206.         p.addname('')
  207.     timing.finish()
  208.     print round(timing.milli() * 0.001 / N, 3), 'seconds per packing'
  209.     p.dump()
  210.     u = Unpacker(p.buf)
  211.     u.getbytes(20)
  212.     u.getname()
  213.     u.getbytes(8)
  214.     u.getname()
  215.     u.getbytes(18)
  216.     u.getname()
  217.     u.getbytes(26)
  218.     u.getname()
  219.     timing.start()
  220.     for i in R:
  221.         u = Unpacker(p.buf)
  222.         res = (u.getbytes(20),
  223.                u.getname(),
  224.                u.getbytes(8),
  225.                u.getname(),
  226.                u.getbytes(18),
  227.                u.getname(),
  228.                u.getbytes(26),
  229.                u.getname())
  230.     timing.finish()
  231.     print round(timing.milli() * 0.001 / N, 3), 'seconds per unpacking'
  232.     for item in res: print item
  233.  
  234.  
  235. # Pack/unpack RR toplevel format (section 3.2.1)
  236.  
  237. class RRpacker(Packer):
  238.     def __init__(self):
  239.         Packer.__init__(self)
  240.         self.rdstart = None
  241.     def addRRheader(self, name, type, klass, ttl, *rest):
  242.         self.addname(name)
  243.         self.add16bit(type)
  244.         self.add16bit(klass)
  245.         self.add32bit(ttl)
  246.         if rest:
  247.             if res[1:]: raise TypeError, 'too many args'
  248.             rdlength = rest[0]
  249.         else:
  250.             rdlength = 0
  251.         self.add16bit(rdlength)
  252.         self.rdstart = len(self.buf)
  253.     def patchrdlength(self):
  254.         rdlength = unpack16bit(self.buf[self.rdstart-2:self.rdstart])
  255.         if rdlength == len(self.buf) - self.rdstart:
  256.             return
  257.         rdata = self.buf[self.rdstart:]
  258.         save_buf = self.buf
  259.         ok = 0
  260.         try:
  261.             self.buf = self.buf[:self.rdstart-2]
  262.             self.add16bit(len(rdata))
  263.             self.buf = self.buf + rdata
  264.             ok = 1
  265.         finally:
  266.             if not ok: self.buf = save_buf
  267.     def endRR(self):
  268.         if self.rdstart is not None:
  269.             self.patchrdlength()
  270.         self.rdstart = None
  271.     def getbuf(self):
  272.         if self.rdstart is not None: self.patchrdlenth()
  273.         return Packer.getbuf(self)
  274.     # Standard RRs (section 3.3)
  275.     def addCNAME(self, name, klass, ttl, cname):
  276.         self.addRRheader(name, dnstype.CNAME, klass, ttl)
  277.         self.addname(cname)
  278.         self.endRR()
  279.     def addHINFO(self, name, klass, ttl, cpu, os):
  280.         self.addRRheader(name, dnstype.HINFO, klass, ttl)
  281.         self.addstring(cpu)
  282.         self.addstring(os)
  283.         self.endRR()
  284.     def addMX(self, name, klass, ttl, preference, exchange):
  285.         self.addRRheader(name, dnstype.MX, klass, ttl)
  286.         self.add16bit(preference)
  287.         self.addname(exchange)
  288.         self.endRR()
  289.     def addNS(self, name, klass, ttl, nsdname):
  290.         self.addRRheader(name, dnstype.NS, klass, ttl)
  291.         self.addname(nsdname)
  292.         self.endRR()
  293.     def addPTR(self, name, klass, ttl, ptrdname):
  294.         self.addRRheader(name, dnstype.PTR, klass, ttl)
  295.         self.addname(ptrdname)
  296.         self.endRR()
  297.     def addSOA(self, name, klass, ttl,
  298.           mname, rname, serial, refresh, retry, expire, minimum):
  299.         self.addRRheader(name, dnstype.SOA, klass, ttl)
  300.         self.addname(mname)
  301.         self.addname(rname)
  302.         self.add32bit(serial)
  303.         self.add32bit(refresh)
  304.         self.add32bit(retry)
  305.         self.add32bit(expire)
  306.         self.add32bit(minimum)
  307.         self.endRR()
  308.     def addTXT(self, name, klass, ttl, list):
  309.         self.addRRheader(name, dnstype.TXT, klass, ttl)
  310.         for txtdata in list:
  311.             self.addstring(txtdata)
  312.         self.endRR()
  313.     # Internet specific RRs (section 3.4) -- class = IN
  314.     def addA(self, name, ttl, address):
  315.         self.addRRheader(name, dnstype.A, dnsclass.IN, ttl)
  316.         self.addaddr(address)
  317.         self.endRR()
  318.     def addWKS(self, name, ttl, address, protocol, bitmap):
  319.         self.addRRheader(name, dnstype.WKS, dnsclass.IN, ttl)
  320.         self.addaddr(address)
  321.         self.addbyte(chr(protocol))
  322.         self.addbytes(bitmap)
  323.         self.endRR()
  324.  
  325.  
  326. class RRunpacker(Unpacker):
  327.     def __init__(self, buf):
  328.         Unpacker.__init__(self, buf)
  329.         self.rdend = None
  330.     def getRRheader(self):
  331.         name = self.getname()
  332.         type = self.get16bit()
  333.         klass = self.get16bit()
  334.         ttl = self.get32bit()
  335.         rdlength = self.get16bit()
  336.         self.rdend = self.offset + rdlength
  337.         return (name, type, klass, ttl, rdlength)
  338.     def endRR(self):
  339.         if self.offset != self.rdend:
  340.             raise UnpackError, 'end of RR not reached'
  341.     def getCNAMEdata(self):
  342.         return self.getname()
  343.     def getHINFOdata(self):
  344.         return self.getstring(), self.getstring()
  345.     def getMXdata(self):
  346.         return self.get16bit(), self.getname()
  347.     def getNSdata(self):
  348.         return self.getname()
  349.     def getPTRdata(self):
  350.         return self.getname()
  351.     def getSOAdata(self):
  352.         return self.getname(), \
  353.                self.getname(), \
  354.                self.get32bit(), \
  355.                self.get32bit(), \
  356.                self.get32bit(), \
  357.                self.get32bit(), \
  358.                self.get32bit()
  359.     def getTXTdata(self):
  360.         list = []
  361.         while self.offset != self.rdend:
  362.             list.append(self.getstring())
  363.         return list
  364.     def getAdata(self):
  365.         return self.getaddr()
  366.     def getWKSdata(self):
  367.         address = self.getaddr()
  368.         protocol = ord(self.getbyte())
  369.         bitmap = self.getbytes(self.rdend - self.offset)
  370.         return address, protocol, bitmap
  371.  
  372.  
  373. # Pack/unpack Message Header (section 4.1)
  374.  
  375. class Hpacker(Packer):
  376.     def addHeader(self, id, qr, opcode, aa, tc, rd, ra, z, rcode,
  377.           qdcount, ancount, nscount, arcount):
  378.         self.add16bit(id)
  379.         self.add16bit((qr&1)<<15 | (opcode*0xF)<<11 | (aa&1)<<10
  380.               | (tc&1)<<9 | (rd&1)<<8 | (ra&1)<<7
  381.               | (z&7)<<4 | (rcode&0xF))
  382.         self.add16bit(qdcount)
  383.         self.add16bit(ancount)
  384.         self.add16bit(nscount)
  385.         self.add16bit(arcount)
  386.  
  387. class Hunpacker(Unpacker):
  388.     def getHeader(self):
  389.         id = self.get16bit()
  390.         flags = self.get16bit()
  391.         qr, opcode, aa, tc, rd, ra, z, rcode = (
  392.               (flags>>15)&1,
  393.               (flags>>11)&0xF,
  394.               (flags>>10)&1,
  395.               (flags>>9)&1,
  396.               (flags>>8)&1,
  397.               (flags>>7)&1,
  398.               (flags>>4)&7,
  399.               (flags>>0)&0xF)
  400.         qdcount = self.get16bit()
  401.         ancount = self.get16bit()
  402.         nscount = self.get16bit()
  403.         arcount = self.get16bit()
  404.         return (id, qr, opcode, aa, tc, rd, ra, z, rcode,
  405.               qdcount, ancount, nscount, arcount)
  406.  
  407.  
  408. # Pack/unpack Question (section 4.1.2)
  409.  
  410. class Qpacker(Packer):
  411.     def addQuestion(self, qname, qtype, qclass):
  412.         self.addname(qname)
  413.         self.add16bit(qtype)
  414.         self.add16bit(qclass)
  415.  
  416. class Qunpacker(Unpacker):
  417.     def getQuestion(self):
  418.         return self.getname(), self.get16bit(), self.get16bit()
  419.  
  420.  
  421. # Pack/unpack Message(section 4)
  422. # NB the order of the base classes is important for __init__()!
  423.  
  424. class Mpacker(RRpacker, Qpacker, Hpacker):
  425.     pass
  426.  
  427. class Munpacker(RRunpacker, Qunpacker, Hunpacker):
  428.     pass
  429.  
  430.  
  431. # Routines to print an unpacker to stdout, for debugging.
  432. # These affect the unpacker's current position!
  433.  
  434. def dumpM(u):
  435.     print 'HEADER:',
  436.     (id, qr, opcode, aa, tc, rd, ra, z, rcode,
  437.           qdcount, ancount, nscount, arcount) = u.getHeader()
  438.     print 'id=%d,' % id,
  439.     print 'qr=%d, opcode=%d, aa=%d, tc=%d, rd=%d, ra=%d, z=%d, rcode=%d,' \
  440.           % (qr, opcode, aa, tc, rd, ra, z, rcode)
  441.     if tc: print '*** response truncated! ***'
  442.     if rcode: print '*** nonzero error code! (%d) ***' % rcode
  443.     print '  qdcount=%d, ancount=%d, nscount=%d, arcount=%d' \
  444.           % (qdcount, ancount, nscount, arcount)
  445.     for i in range(qdcount):
  446.         print 'QUESTION %d:' % i,
  447.         dumpQ(u)
  448.     for i in range(ancount):
  449.         print 'ANSWER %d:' % i,
  450.         dumpRR(u)
  451.     for i in range(nscount):
  452.         print 'AUTHORITY RECORD %d:' % i,
  453.         dumpRR(u)
  454.     for i in range(arcount):
  455.         print 'ADDITIONAL RECORD %d:' % i,
  456.         dumpRR(u)
  457.  
  458. def dumpQ(u):
  459.     qname, qtype, qclass = u.getQuestion()
  460.     print 'qname=%s, qtype=%d(%s), qclass=%d(%s)' \
  461.           % (qname,
  462.              qtype, dnstype.typestr(qtype),
  463.              qclass, dnsclass.classstr(qclass))
  464.  
  465. def dumpRR(u):
  466.     name, type, klass, ttl, rdlength = u.getRRheader()
  467.     typename = dnstype.typestr(type)
  468.     print 'name=%s, type=%d(%s), class=%d(%s), ttl=%d' \
  469.           % (name,
  470.              type, typename,
  471.              klass, dnsclass.classstr(klass),
  472.              ttl)
  473.     mname = 'get%sdata' % typename
  474.     if hasattr(u, mname):
  475.         print '  formatted rdata:', getattr(u, mname)()
  476.     else:
  477.         print '  binary rdata:', u.getbytes(rdlength)
  478.  
  479.  
  480. # Test program
  481.  
  482. def test():
  483.     import sys
  484.     import getopt
  485.     import socket
  486.     protocol = 'udp'
  487.     server = 'cnri.reston.va.us' # XXX adapt this to your local 
  488.     port = 53
  489.     opcode = dnsopcode.QUERY
  490.     rd = 0
  491.     qtype = dnstype.MX
  492.     qname = 'cwi.nl'
  493.     try:
  494.         opts, args = getopt.getopt(sys.argv[1:], 'Trs:tu')
  495.         if len(args) > 2: raise getopt.error, 'too many arguments'
  496.     except getopt.error, msg:
  497.         print msg
  498.         print 'Usage: python dnslib.py',
  499.         print '[-T] [-r] [-s server] [-t] [-u]',
  500.         print '[qtype [qname]]'
  501.         print '-T:        run testpacker() and exit'
  502.         print '-r:        recursion desired (default not)'
  503.         print '-s server: use server (default %s)' % server
  504.         print '-t:        use TCP protocol'
  505.         print '-u:        use UDP protocol (default)'
  506.         print 'qtype:     query type (default %s)' % \
  507.               dnstype.typestr(qtype)
  508.         print 'qname:     query name (default %s)' % qname
  509.         print 'Recognized qtype values:'
  510.         qtypes = dnstype.typemap.keys()
  511.         qtypes.sort()
  512.         n = 0
  513.         for qtype in qtypes:
  514.             n = n+1
  515.             if n >= 8: n = 1; print
  516.             print '%s = %d' % (dnstype.typemap[qtype], qtype),
  517.         print
  518.         sys.exit(2)
  519.     for o, a in opts:
  520.         if o == '-T': testpacker(); return
  521.         if o == '-t': protocol = 'tcp'
  522.         if o == '-u': protocol = 'udp'
  523.         if o == '-s': server = a
  524.         if o == '-r': rd = 1
  525.     if args[0:]:
  526.         try:
  527.             qtype = eval(string.upper(args[0]), dnstype.__dict__)
  528.         except (NameError, SyntaxError):
  529.             print 'bad query type:', `args[0]`
  530.             sys.exit(2)
  531.     if args[1:]:
  532.         qname = args[1]
  533.     if qtype == dnstype.AXFR:
  534.         print 'Query type AXFR, protocol forced to TCP'
  535.         protocol = 'tcp'
  536.     print 'QTYPE %d(%s)' % (qtype, dnstype.typestr(qtype))
  537.     m = Mpacker()
  538.     m.addHeader(0,
  539.           0, opcode, 0, 0, rd, 0, 0, 0,
  540.           1, 0, 0, 0)
  541.     m.addQuestion(qname, qtype, dnsclass.IN)
  542.     request = m.getbuf()
  543.     if protocol == 'udp':
  544.         s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
  545.         s.connect((server, port))
  546.         s.send(request)
  547.         reply = s.recv(1024)
  548.     else:
  549.         s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
  550.         s.connect((server, port))
  551.         s.send(pack16bit(len(request)) + request)
  552.         s.shutdown(1)
  553.         f = s.makefile('r')
  554.         header = f.read(2)
  555.         if len(header) < 2:
  556.             print '*** EOF ***'
  557.             return
  558.         count = unpack16bit(header)
  559.         reply = f.read(count)
  560.         if len(reply) != count:
  561.             print '*** Incomplete reply ***'
  562.             return
  563.     u = Munpacker(reply)
  564.     dumpM(u)
  565.     if protocol == 'tcp' and qtype == dnstype.AXFR:
  566.         while 1:
  567.             header = f.read(2)
  568.             if len(header) < 2:
  569.                 print '========== EOF =========='
  570.                 break
  571.             count = unpack16bit(header)
  572.             if not count:
  573.                 print '========== ZERO COUNT =========='
  574.                 break
  575.             print '========== NEXT =========='
  576.             reply = f.read(count)
  577.             if len(reply) != count:
  578.                 print '*** Incomplete reply ***'
  579.                 break
  580.             u = Munpacker(reply)
  581.             dumpM(u)
  582.  
  583.  
  584. # Run test program when called as a script
  585.  
  586. if __name__ == '__main__':
  587.     test()
  588.