From 29b148cbcc85da6d73c459844d144c4eb8676fc9 Mon Sep 17 00:00:00 2001
From: henices <zhouzhenster@gmail.com>
Date: Fri, 18 Jul 2014 15:50:44 +0800
Subject: [PATCH] bug fixed

---
 tcpdns.py | 57 ++++++++++++++++++++++++++++---------------------------
 1 file changed, 29 insertions(+), 28 deletions(-)

diff --git a/tcpdns.py b/tcpdns.py
index 8402ae3..61b5c60 100755
--- a/tcpdns.py
+++ b/tcpdns.py
@@ -138,32 +138,33 @@ def private_dns_response(data):
     print 'domain:%s, qtype:%x' % (q_domain, qtype)
     sys.stdout.flush()
 
-    if qtype != 0x0001:
-        return None
-
-    if Questions != '\x00\x01' or AnswerRRs != '\x00\x00' or \
-        AuthorityRRs != '\x00\x00' or AdditionalRRs != '\x00\x00':
-            return None
-
-    items = cfg['private_host'].items()
-
-    for domain, ip in items:
-        if fnmatch(q_domain, domain):
-            ret = TID
-            ret += '\x81\x80'
-            ret += '\x00\x01'
-            ret += '\x00\x01'
-            ret += '\x00\x00'
-            ret += '\x00\x00'
-            ret += data[12:]
-            ret += '\xc0\x0c'
-            ret += '\x00\x01'
-            ret += '\x00\x01'
-            ret += '\x00\x00\xff\xff'
-            ret += '\x00\x04'
-            ret +=  socket.inet_aton(ip)
-
-    return ret
+    try:
+        if qtype != 0x0001:
+            return
+
+        if Questions != '\x00\x01' or AnswerRRs != '\x00\x00' or \
+            AuthorityRRs != '\x00\x00' or AdditionalRRs != '\x00\x00':
+                return
+
+        items = cfg['private_host'].items()
+
+        for domain, ip in items:
+            if fnmatch(q_domain, domain):
+                ret = TID
+                ret += '\x81\x80'
+                ret += '\x00\x01'
+                ret += '\x00\x01'
+                ret += '\x00\x00'
+                ret += '\x00\x00'
+                ret += data[12:]
+                ret += '\xc0\x0c'
+                ret += '\x00\x01'
+                ret += '\x00\x01'
+                ret += '\x00\x00\xff\xff'
+                ret += '\x00\x04'
+                ret +=  socket.inet_aton(ip)
+    finally:
+        return (q_domain, ret)
 
 
 def check_dns_packet(data):
@@ -199,7 +200,7 @@ def transfer(querydata, addr, server):
     t_id = querydata[:2]
     key = querydata[2:].encode('hex')
 
-    response = private_dns_response(querydata)
+    q_domain, response = private_dns_response(querydata)
     if response:
         server.sendto(response, addr)
         return
@@ -232,7 +233,7 @@ def transfer(querydata, addr, server):
         break
 
     if response is None:
-        print "[ERROR] Tried many times and failed to resolve %s" % domain
+        print "[ERROR] Tried many times and failed to resolve %s" % q_domain
 
 
 class ThreadedUDPServer(SocketServer.ThreadingMixIn, SocketServer.UDPServer):