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