controller/smartdns/sdns.py at cb37b7ef1d9d6d2d632e602f640abd4e51a7a7b4 · drycc/controller · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
# -*- coding: utf-8 -*-

# version
__version__ = '1.0.1.1'

import socket
import time
import signal
import sys
import os
import yaml
import logging
from os.path import isfile
from zope.interface import implements
from twisted.internet import defer, interfaces
from twisted.python import failure
from twisted.internet.protocol import DatagramProtocol
from twisted.application import service, internet
from twisted.names import dns, server, client, cache, common, resolve
from twisted.internet import reactor
from . import dnsserver, ippool, monitor


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 get_local_ip():
    import sys
    import socket
    import fcntl
    import array
    import struct
    is_64bits = sys.maxsize > 2**32
    struct_size = 40 if is_64bits else 32
    s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    max_possible = 8  # initial value
    while True:
        bytes = max_possible * struct_size
        names = array.array('B', b'\0' * bytes)
        outbytes = struct.unpack('iL', fcntl.ioctl(
            s.fileno(),
            0x8912,  # SIOCGIFCONF
            struct.pack('iL', bytes, names.buffer_info()[0])
        ))[0]
        if outbytes == bytes:
            max_possible *= 2
        else:
            break
    namestr = names.tostring()
    return [(namestr[i:i+16].split(b'\0', 1)[0],
             socket.inet_ntoa(namestr[i+20:i+24]))
            for i in range(0, outbytes, struct_size)]


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 ......')
    Amapping = loadconfig(os.path.join(run_env['conf'], 'a.yaml'))
    NSmapping = loadconfig(os.path.join(run_env['conf'], 'ns.yaml'))
    SOAmapping = 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, Amapping)
    # load dns record config file
    logger.info('start to init IP pool ......')
    Finder = ippool.IPPool(
        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
    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)
        f = dnsserver.SmartDNSFactory(
            caches=[cache.CacheResolver()], clients=[SmartResolver])
        p = dns.DNSDatagramProtocol(f)
        f.noisy = p.noisy = False
        run_env['tcp'].append([f, ip])
        run_env['udp'].append([p, ip])


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'

    prepare_run(run_env)
    for e in run_env['tcp']:
        reactor.listenTCP(53, e[0], interface=e[1])
    for e in run_env['udp']:
        reactor.listenUDP(53, e[0], interface=e[1])
    reactor.run()