Merge branch 'Drewsif-master'

Author: Mraoul

modules/dns.py

@@ -21,53 +21,70 @@
 # OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF
 # SUCH DAMAGE.
 
-from dnslib import DNSRecord, QR, OPCODE, RCODE, QTYPE, CLASS, DNSError
+from dnslib import (DNSRecord,
+                    QR,
+                    OPCODE,
+                    RCODE,
+                    QTYPE,
+                    CLASS,
+                    DNSError)
 
 from ChopProtocol import ChopProtocol
 
 moduleName = "dns"
 moduleVersion = '0.1'
 minimumChopLib = '4.0'
 
+
 def module_info():
     return "Parse DNS"
 
+
 def init(module_data):
-    module_options = { 'proto': [ { 'udp': 'dns' } ] }
+    module_options = {'proto': [{'udp': 'dns'}]}
     return module_options
 
+
 def handleDatagram(udp):
     ((src, sport), (dst, dport)) = udp.addr
     if sport != 53 and dport != 53:
-        #chop.tsprnt("STOP: %s:%s->%s:%s (%i:%i)" % (src, sport, dst, dport, len(udp.data), len(udp.ip)))
         udp.stop()
         return
 
     try:
         o = DNSRecord.parse(udp.data)
-    except KeyError, e:
+    except KeyError as e:
         chop.prnt("Key error: %s" % str(e))
         return
-    except DNSError, e:
+    except DNSError as e:
         chop.prnt("dnslib error: %s" % str(e))
         return
+    except Exception as e:
+        chop.prnt("Unexpeced exception: %s" % str(e))
+        return
 
     chopp = ChopProtocol('dns')
 
     # Create the dictionary...
-    f = [ o.header.aa and 'AA',
-          o.header.tc and 'TC',
-          o.header.rd and 'RD',
-          o.header.ra and 'RA' ]
-    d = { 'header': {
-                      'id': o.header.id,
-                      'type': QR[o.header.get_qr()],
-                      'opcode': OPCODE[o.header.get_opcode()],
-                      'flags': ",".join(filter(None, f)),
-                      'rcode': RCODE[o.header.rcode],
-                    },
-          'questions': o.questions
-        }
+    f = [o.header.aa and 'AA',
+         o.header.tc and 'TC',
+         o.header.rd and 'RD',
+         o.header.ra and 'RA']
+
+    try:
+        d = {'header': {'id': o.header.id,
+                        'type': QR[o.header.get_qr()],
+                        'opcode': OPCODE[o.header.get_opcode()],
+                        'flags': ",".join(filter(None, f)),
+                        'rcode': RCODE[o.header.rcode]},
+             'questions': o.questions}
+    except DNSError as e:
+        chop.prnt("dnslib error: %s" % str(e))
+        return
+    except Exception as e:
+        chop.prnt("Unexpeted exception: %s" % str(e))
+        return
+
     if OPCODE[o.header.opcode] == 'UPDATE':
         f1 = 'zo'
         f2 = 'pr'
@@ -78,23 +95,32 @@ def handleDatagram(udp):
         f2 = 'a'
         f3 = 'ns'
         f4 = 'ar'
+
     dhdr = d['header']
     dhdr[f1] = o.header.q
     dhdr[f2] = o.header.a
     dhdr[f3] = o.header.auth
-    dhdr[f4]= o.header.ar
+    dhdr[f4] = o.header.ar
+
     d['questions'] = []
     for q in o.questions:
         qname = str(q.get_qname())
         # Strip trailing dot.
         if qname.endswith('.'):
             qname = qname[:-1]
-        dq = {
-              'qname': qname,
-              'qtype': QTYPE[q.qtype],
-              'qclass': CLASS[q.qclass]
-            }
-        d['questions'].append(dq)
+
+        try:
+            dq = {'qname': qname,
+                  'qtype': QTYPE[q.qtype],
+                  'qclass': CLASS[q.qclass]}
+            d['questions'].append(dq)
+        except DNSError as e:
+            chop.prnt("dnslib error: %s" % str(e))
+            return
+        except Exception as e:
+            chop.prnt("Unexpected exception: %s" % str(e))
+            return
+
     d['rr'] = []
     for r in o.rr:
         rname = str(r.get_rname())
@@ -105,14 +131,20 @@ def handleDatagram(udp):
         # Strip trailing dot.
         if rdata.endswith('.'):
             rdata = rdata[:-1]
-        dr = {
-              'rname': rname,
-              'rtype': QTYPE[r.rtype],
-              'rclass': CLASS[r.rclass],
-              'ttl': r.ttl,
-              'rdata': rdata
-            }
-        d['rr'].append(dr)
+
+        try:
+            dr = {'rname': rname,
+                  'rtype': QTYPE[r.rtype],
+                  'rclass': CLASS[r.rclass],
+                  'ttl': r.ttl,
+                  'rdata': rdata}
+            d['rr'].append(dr)
+        except DNSError as e:
+            chop.prnt("dnslib error: %s" % str(e))
+            return
+        except Exception as e:
+            chop.prnt("Unexpected exception: %s" % str(e))
+            return
 
     if sport == 53:
         chopp.serverData = d
@@ -122,6 +154,7 @@ def handleDatagram(udp):
         return chopp
 
     return None
-   
+
+
 def shutdown(module_data):
     return