add ping support

Author: duanhongyi

bin/dnsserver.py

@@ -54,9 +54,10 @@ def __call__(self, failure):
 
 
 class MapResolver(client.Resolver):
-    def __init__(self, Finder, Amapping, NSmapping, SOAmapping, servers):
+    def __init__(self, Finder, Amapping, AmappingBlacklist, NSmapping, SOAmapping, servers):
         self.Finder = Finder
         self.Amapping = Amapping
+        self.AmappingBlacklist = AmappingBlacklist
         self.NSmapping = NSmapping
         self.SOAmapping = SOAmapping
         client.Resolver.__init__(self, servers=servers)
@@ -77,6 +78,7 @@ def packResult( value ):
                 ret = []
                 add = []
                 for x in value:
+                    if x in self.AmappingBlacklist: continue
                     ret.append(dns.RRHeader(name, dns.A, dns.IN, ttl, dns.Record_A(x, ttl), True))
 
                 if edns is not None:

bin/icmp.py

@@ -0,0 +1,114 @@
+import os, sys, socket, struct, select, time
+ 
+# From /usr/include/linux/icmp.h; your milage may vary.
+ICMP_ECHO_REQUEST = 8 # Seems to be the same on Solaris.
+ 
+ 
+def checksum(source_string):
+    """
+    I'm not too confident that this is right but testing seems
+    to suggest that it gives the same answers as in_cksum in ping.c
+    """
+    sum = 0
+    countTo = (len(source_string)/2)*2
+    count = 0
+    while count<countTo:
+        thisVal = ord(source_string[count + 1])*256 + ord(source_string[count])
+        sum = sum + thisVal
+        sum = sum & 0xffffffff # Necessary?
+        count = count + 2
+ 
+    if countTo<len(source_string):
+        sum = sum + ord(source_string[len(source_string) - 1])
+        sum = sum & 0xffffffff # Necessary?
+ 
+    sum = (sum >> 16)  +  (sum & 0xffff)
+    sum = sum + (sum >> 16)
+    answer = ~sum
+    answer = answer & 0xffff
+ 
+    # Swap bytes. Bugger me if I know why.
+    answer = answer >> 8 | (answer << 8 & 0xff00)
+ 
+    return answer
+ 
+ 
+def receive_one_ping(my_socket, ID, timeout):
+    """
+    receive the ping from the socket.
+    """
+    timeLeft = timeout
+    while True:
+        startedSelect = time.time()
+        whatReady = select.select([my_socket], [], [], timeLeft)
+        howLongInSelect = (time.time() - startedSelect)
+        if whatReady[0] == []: # Timeout
+            return
+ 
+        timeReceived = time.time()
+        recPacket, addr = my_socket.recvfrom(1024)
+        icmpHeader = recPacket[20:28]
+        type, code, checksum, packetID, sequence = struct.unpack(
+            "bbHHh", icmpHeader
+        )
+        if packetID == ID:
+            bytesInDouble = struct.calcsize("d")
+            timeSent = struct.unpack("d", recPacket[28:28 + bytesInDouble])[0]
+            return timeReceived - timeSent
+ 
+        timeLeft = timeLeft - howLongInSelect
+        if timeLeft <= 0:
+            return
+ 
+ 
+def send_one_ping(my_socket, dest_addr, ID):
+    """
+    Send one ping to the given >dest_addr<.
+    """
+    dest_addr  =  socket.gethostbyname(dest_addr)
+ 
+    # Header is type (8), code (8), checksum (16), id (16), sequence (16)
+    my_checksum = 0
+ 
+    # Make a dummy heder with a 0 checksum.
+    header = struct.pack("bbHHh", ICMP_ECHO_REQUEST, 0, my_checksum, ID, 1)
+    bytesInDouble = struct.calcsize("d")
+    data = (192 - bytesInDouble) * "Q"
+    data = struct.pack("d", time.time()) + data
+ 
+    # Calculate the checksum on the data and the dummy header.
+    my_checksum = checksum(header + data)
+ 
+    # Now that we have the right checksum, we put that in. It's just easier
+    # to make up a new header than to stuff it into the dummy.
+    header = struct.pack(
+        "bbHHh", ICMP_ECHO_REQUEST, 0, socket.htons(my_checksum), ID, 1
+    )
+    packet = header + data
+    my_socket.sendto(packet, (dest_addr, 1)) # Don't know about the 1
+ 
+ 
+def ping(dest_addr, timeout):
+    """
+    Returns either the delay (in seconds) or none on timeout.
+    """
+    icmp = socket.getprotobyname("icmp")
+    try:
+        my_socket = socket.socket(socket.AF_INET, socket.SOCK_RAW, icmp)
+    except socket.error, (errno, msg):
+        if errno == 1:
+            # Operation not permitted
+            msg = msg + (
+                " - Note that ICMP messages can only be sent from processes"
+                " running as root."
+            )
+            raise socket.error(msg)
+        raise # raise the original error
+ 
+    my_ID = os.getpid() & 0xFFFF
+ 
+    send_one_ping(my_socket, dest_addr, my_ID)
+    delay = receive_one_ping(my_socket, my_ID, timeout)
+ 
+    my_socket.close()
+    return delay

bin/sdns.py

@@ -6,6 +6,7 @@
 import sys, os
 import signal
 import time
+import threading
 from os.path import isfile
 sys.path.append('../lib')
 import yaml
@@ -28,6 +29,7 @@
 sys.path.append('ippool.py')
 sys.path.append('dnsserver.py')
 import ippool, dnsserver
+from icmp import ping
 
 def loadconfig(path):
 	if not isfile(path):
@@ -61,6 +63,23 @@ def get_local_ip():
 		socket.inet_ntoa(namestr[i+20:i+24]))
 		for i in range(0, outbytes, struct_size)]
 
+def prepare_ip_blacklist(amapping_blacklist, ip_list):
+    def task(amapping_blacklist, ip_list):
+        while True:
+            for ip in ip_list:
+                try:
+                    delay  =  ping(ip, 9)
+                except socket.gaierror as e:
+                    delay = None
+                if delay == None and ip not in amapping_blacklist:
+                    amapping_blacklist.add(ip)
+                if delay != None and ip in amapping_blacklist:
+                    amapping_blacklist.remove(ip)
+    thread = threading.Thread(target=task, args=(amapping_blacklist, ip_list)) 
+    thread.setDaemon(True)
+    thread.start()
+
+
 def prepare_run(run_env):
 	#load main config
 	logger.info('start to load conf/sdns.yaml ......')
@@ -75,15 +94,22 @@ def prepare_run(run_env):
 	Amapping = loadconfig(conf['AFILE'])
 	NSmapping = loadconfig(conf['NSFILE'])
 	SOAmapping = loadconfig(conf['SOAFILE'])
-
+        AmappingBlacklist = set()
+        ip_list = set()
+        for item in Amapping.values():
+            for key, value in item.items():
+                if key == 'ttl': continue
+                ip_list.update(value.split(' '))
+        prepare_ip_blacklist(AmappingBlacklist, ip_list) 
 	# set up a resolver that uses the mapping or a secondary nameserver
 	dnsforward = []
 	for i in conf['dnsforward_ip']:
 		dnsforward.append((i, conf['dnsforward_port']))
 
 	for ifc,ip in get_local_ip():
 		# create the protocols
-		SmartResolver = dnsserver.MapResolver(Finder, Amapping, NSmapping, SOAmapping, servers=dnsforward)
+		SmartResolver = dnsserver.MapResolver(
+                    Finder, Amapping, AmappingBlacklist, NSmapping, SOAmapping, servers=dnsforward)
 		f = dnsserver.SmartDNSFactory(caches=[cache.CacheResolver()], clients=[SmartResolver])
 		p = dns.DNSDatagramProtocol(f)
 		f.noisy = p.noisy = False