fix: fix no edns bug

Author: duanhongyi

.gitignore

@@ -21,3 +21,5 @@ develop-eggs
 *.pyc
 *.pid
 .venv
+.idea
+.vscode

setup.py

@@ -25,7 +25,7 @@
 
     name='smartdns',
     description='An Intelligent DNS Similar to DNSPod',
-    version='3.3.3',
+    version='3.3.4',
     author='duanhongyi',
     author_email='duanhyi@gmail.com',
     packages=find_packages(),
@@ -38,5 +38,5 @@
     license='BSD',
     long_description_content_type='text/markdown',
     # install
-    entry_points={'console_scripts': ['sdns=smartdns.sdns:main']}
+    entry_points={'console_scripts': ['sdns=smartdns.runner:run']}
 )

smartdns/runner.py

@@ -0,0 +1,93 @@
+# -*- coding: utf-8 -*-
+
+# version
+__version__ = '1.0.1.1'
+
+import logging
+import os
+import sys
+import yaml
+from multiprocessing import cpu_count, Process
+from os.path import isfile
+from twisted.internet import reactor
+from twisted.names import dns
+
+from . import server, monitor
+from .finder import Finder
+
+logger = logging.getLogger(__name__)
+
+
+def loadconfig(path):
+    if not isfile(path):
+        print("[FATAL] can't find config file %s !" % path)
+        exit(1)
+    with open(path, 'r') as f:
+        return yaml.load(f, Loader=yaml.FullLoader)
+
+
+def prepare_run(run_env):
+    # load main config
+    logger.info('start to load %s ......' % run_env['conf'])
+    conf = loadconfig(os.path.join(run_env['conf'], 'sdns.yaml'))
+    logger.info('start to load A,SOA,NS record ......')
+    a_mapping = loadconfig(os.path.join(run_env['conf'], 'a.yaml'))
+    ns_mapping = loadconfig(os.path.join(run_env['conf'], 'ns.yaml'))
+    soa_mapping = loadconfig(os.path.join(run_env['conf'], 'soa.yaml'))
+
+    # start monitor
+    monitor_config = loadconfig(os.path.join(run_env['conf'], 'monitor.yaml'))
+    monitor_mapping = monitor.MonitorMapping(monitor_config, a_mapping)
+    # load dns record config file
+    logger.info('start to init IP pool ......')
+    finder = Finder(
+        os.path.join(run_env['conf'], 'ip.csv'),
+        os.path.join(run_env['conf'], 'a.yaml'),
+        monitor_mapping)
+
+    run_env['finder'] = finder
+
+    # set up a resolver that uses the mapping or a secondary nameserver
+    dns_forwards = []
+    for i in conf['dns_forwards']:
+        dns_forward_ip, dns_forward_port = i.split(":")
+        dns_forwards.append((dns_forward_ip, int(dns_forward_port)))
+
+    # create the protocols
+    for listen_tcp in conf['listen']['tcp']:
+        listen_tcp_ip, listen_tcp_port = listen_tcp.split(":")
+        f = server.SmartDNSFactory(clients=[
+            server.MapResolver(
+                finder, a_mapping, ns_mapping, soa_mapping, servers=dns_forwards)])
+        f.noisy = False
+        run_env['tcp'].append([int(listen_tcp_port), f, listen_tcp_ip])
+    for listen_udp in conf['listen']['tcp']:
+        listen_udp_ip, listen_udp_port = listen_udp.split(":")
+        p = dns.DNSDatagramProtocol(server.SmartDNSFactory(clients=[
+            server.MapResolver(
+                finder, a_mapping, ns_mapping, soa_mapping, servers=dns_forwards)]))
+        p.noisy = False
+        run_env['udp'].append([int(listen_udp_port), p, listen_udp_ip])
+    return conf
+
+
+def run():
+    run_env = {'udp': [], 'tcp': [], 'closed': 0,
+               'updated': False, 'finder': None}
+    if len(sys.argv) > 1:
+        run_env['conf'] = sys.argv[1]
+    else:
+        run_env['conf'] = '/etc/smartdns'
+
+    conf = prepare_run(run_env)
+    for e in run_env['tcp']:
+        reactor.listenTCP(e[0], e[1], interface=e[2])
+    for e in run_env['udp']:
+        reactor.listenUDP(e[0], e[1], interface=e[2])
+    workers = []
+    for _ in range(conf.get("workers", cpu_count() * 2)):
+        process = Process(target=reactor.run)
+        process.daemon = True
+        process.start()
+        workers.append(process)
+    [worker.join() for worker in workers]

smartdns/sdns.py

@@ -1,93 +1,55 @@
-# -*- coding: utf-8 -*-
+import socket
+import struct
+from zope.interface import implementer
+from twisted.names.dns import str2time, tputil, readPrecisely, IRecord, IEncodable
 
-# version
-__version__ = '1.0.1.1'
+EDNS = 41
 
-import logging
-import os
-import sys
-import yaml
-from multiprocessing import cpu_count, Process
-from os.path import isfile
-from twisted.internet import reactor
-from twisted.names import dns
+@implementer(IEncodable, IRecord)
+class Record_EDNS(tputil.FancyEqMixin):
+    """
+    An EDNS
+    """
 
-from . import server, monitor
-from .finder import Finder
+    compareAttributes = ('address', 'family', 'opt_code', 'opt_length', 'source_nm', 'scope_nm')
 
-logger = logging.getLogger(__name__)
+    TYPE = EDNS
+    address = None
+    opt_code = None
+    opt_length = None
+    family =  None
+    source_nm = None
+    scope_nm = None
 
+    def __init__(self, address='0.0.0.0', ttl=None):
+        if address is not None:
+            address = socket.inet_aton(address)
+            self.address = address
+            self.ttl = str2time(ttl)
 
-def loadconfig(path):
-    if not isfile(path):
-        print("[FATAL] can't find config file %s !" % path)
-        exit(1)
-    with open(path, 'r') as f:
-        return yaml.load(f, Loader=yaml.FullLoader)
 
+    def encode(self, strio, compDict = None):
+        if self.address is not None:
+            strio.write(struct.pack('!HHHcc', self.opt_code, self.opt_length, self.family, self.source_nm, self.scope_nm))
+            strio.write(self.address)
 
-def prepare_run(run_env):
-    # load main config
-    logger.info('start to load %s ......' % run_env['conf'])
-    conf = loadconfig(os.path.join(run_env['conf'], 'sdns.yaml'))
-    logger.info('start to load A,SOA,NS record ......')
-    a_mapping = loadconfig(os.path.join(run_env['conf'], 'a.yaml'))
-    ns_mapping = loadconfig(os.path.join(run_env['conf'], 'ns.yaml'))
-    soa_mapping = loadconfig(os.path.join(run_env['conf'], 'soa.yaml'))
+    def decode(self, strio, length = None):
+        if length > 8:
+            rdata = readPrecisely(strio, 8)
+            self.opt_code, self.opt_length, self.family, self.source_nm, self.scope_nm = struct.unpack("!HHHcc", rdata)
+            #hack let scope netmask = source netmask
+            self.scope_nm = self.source_nm
+            self.address = readPrecisely(strio, length - 8)
 
-    # start monitor
-    monitor_config = loadconfig(os.path.join(run_env['conf'], 'monitor.yaml'))
-    monitor_mapping = monitor.MonitorMapping(monitor_config, a_mapping)
-    # load dns record config file
-    logger.info('start to init IP pool ......')
-    finder = Finder(
-        os.path.join(run_env['conf'], 'ip.csv'),
-        os.path.join(run_env['conf'], 'a.yaml'),
-        monitor_mapping)
+    def __hash__(self):
+        return hash(self.address)
 
-    run_env['finder'] = finder
 
-    # set up a resolver that uses the mapping or a secondary nameserver
-    dns_forwards = []
-    for i in conf['dns_forwards']:
-        dns_forward_ip, dns_forward_port = i.split(":")
-        dns_forwards.append((dns_forward_ip, int(dns_forward_port)))
+    def __str__(self):
+        return '<A address=%s ttl=%s>' % (self.address, self.ttl)
+    __repr__ = __str__
 
-    # create the protocols
-    for listen_tcp in conf['listen']['tcp']:
-        listen_tcp_ip, listen_tcp_port = listen_tcp.split(":")
-        f = server.SmartDNSFactory(clients=[
-            server.MapResolver(
-                finder, a_mapping, ns_mapping, soa_mapping, servers=dns_forwards)])
-        f.noisy = False
-        run_env['tcp'].append([int(listen_tcp_port), f, listen_tcp_ip])
-    for listen_udp in conf['listen']['tcp']:
-        listen_udp_ip, listen_udp_port = listen_udp.split(":")
-        p = dns.DNSDatagramProtocol(server.SmartDNSFactory(clients=[
-            server.MapResolver(
-                finder, a_mapping, ns_mapping, soa_mapping, servers=dns_forwards)]))
-        p.noisy = False
-        run_env['udp'].append([int(listen_udp_port), p, listen_udp_ip])
-    return conf
 
-
-def main():
-    run_env = {'udp': [], 'tcp': [], 'closed': 0,
-               'updated': False, 'finder': None}
-    if len(sys.argv) > 1:
-        run_env['conf'] = sys.argv[1]
-    else:
-        run_env['conf'] = '/etc/smartdns'
-
-    conf = prepare_run(run_env)
-    for e in run_env['tcp']:
-        reactor.listenTCP(e[0], e[1], interface=e[2])
-    for e in run_env['udp']:
-        reactor.listenUDP(e[0], e[1], interface=e[2])
-    workers = []
-    for _ in range(conf.get("workers", cpu_count() * 2)):
-        process = Process(target=reactor.run)
-        process.daemon = True
-        process.start()
-        workers.append(process)
-    [worker.join() for worker in workers]
+    def dottedQuad(self):
+        n = 4 - len(self.address)
+        return socket.inet_ntoa(self.address + '\x00' * n)

smartdns/server.py

@@ -6,6 +6,8 @@
 from twisted.names import dns, server, client, common, resolve
 from twisted.python import failure
 
+from . import sdns
+
 logger = logging.getLogger(__name__)
 
 typeToMethod = {
@@ -81,9 +83,9 @@ def packResult(value):
 
                 if edns is not None:
                     if edns.rdlength > 8:
-                        add.append(dns.RRHeader('', dns.EDNS, 4096, edns.ttl, edns.payload, True))
+                        add.append(dns.RRHeader('', sdns.EDNS, 4096, edns.ttl, edns.payload, True))
                     else:
-                        add.append(dns.RRHeader('', dns.EDNS, 4096, 0, dns.Record_EDNS(None, 0), True))
+                        add.append(dns.RRHeader('', sdns.EDNS, 4096, 0, sdns.Record_EDNS(None, 0), True))
                 return [ret, (), add]
 
             result = self.finder.findIP(str(addr[0]), name)
@@ -119,7 +121,7 @@ def packResultSOA(value):
                     if edns.rdlength > 8:
                         add.append(dns.RRHeader('', dns.EDNS, 4096, edns.ttl, edns.payload, True))
                     else:
-                        add.append(dns.RRHeader('', dns.EDNS, 4096, 0, dns.Record_EDNS(None, 0), True))
+                        add.append(dns.RRHeader('', dns.EDNS, 4096, 0, sdns.Record_EDNS(None, 0), True))
 
                 return [(dns.RRHeader(name, dns.SOA, dns.IN, value['ttl'],
                                       dns.Record_SOA(value['record'], value['email'], value['serial'], value['refresh'],