2dns.py

#!/usr/bin/env python3

from time import time, sleep
try:
    from time import time_ns
except ImportError:
    def time_ns(): return int(str(time()).replace('.', ''))

from socket import socket, timeout, AF_INET, AF_INET6, SOCK_DGRAM, SOCK_STREAM
from threading import Thread


def Send(data, server, f, count=3, timeout=10, drop=-1):
    assert type(data) is bytes and type(server) is tuple
    assert f in (AF_INET, AF_INET6)
    assert type(count) is int and type(timeout) in (int, float)

    with socket(f, SOCK_DGRAM) as c:
        c.settimeout(int(timeout))
        for i in range(0, int(count)):
            for ii in server:
                c.sendto(data, ii)
            del ii
        del i
        t = time()
        while True:
            rdata, raddr = c.recvfrom(4096)
            if raddr in server and (time() - t) > drop:
                break
            del rdata, raddr
    del c

    return (rdata, raddr)


def DNSPackDomainGet(data):
    assert type(data) is bytes
    raw_domain = data[12:]

    for i in range(0, len(raw_domain)):
        if chr(raw_domain[i]) == '\x00':
            break
    raw_domain = raw_domain[0:i]
    end = i + 12
    del i

    domain = raw_domain.decode()

    domain = domain[1:]
    domain = str(domain.encode())[1:-1]

    dots = ['\\t', '\\n']
    for i in range(0, 256):
        ii = str(hex(i)).replace('0x', '')
        if len(ii) < 2:
            ii = '0'+ii
        dots.append('\\x'+ii)
        del ii
    del i

    for i in dots:
        domain = domain.replace(i, '.')
    del i

    domain = domain[1:]
    domain += '.'

    return ((domain, raw_domain), end)


class Master(Thread):
    Count = {
        'Accepted': 0,
        'Processed': {
            'All': 0,
            'Checked': [0, 0],
            'Queryed': 0,
            'Replyed': 0
        }
    }
    Time = {
        'Waiting': 0.0,
        'Processing': {
            'All': 0.0,
            'Check': 0.0,
            'Query': 0.0,
            'Other': 0.0
        }
    }
    Log = []

    TRUE_REPLY_DOMAINS = set()
    FALSE_REPLY_DOMAINS = set()

    Lt = time()
    L = hash(str(time_ns()))
    for i in range(0, int(str(time_ns())[-3:])+1):
        L += (hash(str(time_ns())) * id(time_ns()+i))
        L -= (hash(format(L, 'x')) ** int(str(time_ns())[-3:]))
    del i
    L = format(L, 'x')
    Lt = time() - Lt
    Log.append('{t}; L is generated. data: #{d}#, hash: {dh}, len: {dl}, use time: [{ut}].'.format(
        t=time(), d=L, dh=hash(L), dl=len(L), ut=Lt))

    RootServer4 = (
        ('198.41.0.4', 53),      # a.root-servers.net
        ('199.9.14.201', 53),    # b.root-servers.net
        ('192.33.4.12', 53),     # c.root-servers.net
        ('199.7.91.13', 53),     # d.root-servers.net
        ('192.203.230.10', 53),  # e.root-servers.net
        ('192.112.36.4', 53),    # g.root-servers.net
        ('198.97.190.53', 53),   # h.root-servers.net
        ('202.12.27.33', 53)     # m.root-servers.net
    )
    RootServer6 = (
        ('2001:503:ba3e::2:30', 53),  # a.root-servers.net
        ('2001:500:200::b', 53),      # b.root-servers.net
        ('2001:500:2::c', 53),        # c.root-servers.net
        ('2001:500:2d::d', 53),       # d.root-servers.net
        ('2001:500:a8::e', 53),       # e.root-servers.net
        ('2001:500:12::d0d', 53),     # g.root-servers.net
        ('2001:500:1::53', 53),       # h.root-servers.net
        ('2001:dc3::35', 53)          # m.root-servers.net
    )

    CheckServer4 = [
        ('8.8.8.8', 53),  # dns.google
        ('8.8.4.4', 53)   # dns.google
    ]
    CheckServer6 = [
        ('2001:4860:4860::8888', 53),  # dns.google
        ('2001:4860:4860::8844', 53)   # dns.google
    ]

    def __init__(self, CMODE, QMODE, SERVER, LOCAL_UPSTREAM, REMOTE_UPSTREAM):
        assert type(CMODE) is str and CMODE in (
            'root', 'reset', 'dual', 'void')
        self.CMode = CMODE

        assert type(QMODE) is str and QMODE in ('non-53port', 'non-first')
        self.QMode = QMODE

        assert type(SERVER) is socket
        self.Server = SERVER

        assert type(LOCAL_UPSTREAM) is tuple and len(LOCAL_UPSTREAM) > 0
        for i in LOCAL_UPSTREAM:
            assert type(i[0]) is str and ('.' in i[0] or ':' in i[0])
            assert not ('.' in i[0] and ':' in i[0])
            assert type(i[1]) is int and 0 < i[1] <= 65535
        del i
        if '.' in LOCAL_UPSTREAM[0][0]:
            self.localupf = AF_INET
        elif ':' in LOCAL_UPSTREAM[0][0]:
            self.localupf = AF_INET6
        else:
            raise ValueError
        self.LocalUpstream = LOCAL_UPSTREAM

        assert type(REMOTE_UPSTREAM) is tuple and len(REMOTE_UPSTREAM) > 0
        for i in REMOTE_UPSTREAM:
            assert type(i[0]) is str and ('.' in i[0] or ':' in i[0])
            assert not ('.' in i[0] and ':' in i[0])
            assert type(i[1]) is int and 0 < i[1] <= 65535

            if self.QMode == 'non-53port':
                assert i[1] != 53
        del i
        if '.' in REMOTE_UPSTREAM[0][0]:
            self.remoteupf = AF_INET
        elif ':' in REMOTE_UPSTREAM[0][0]:
            self.remoteupf = AF_INET6
        else:
            raise ValueError
        self.RemoteUpstream = REMOTE_UPSTREAM

        Thread.__init__(self)

    def run(self):
        while True:
            Waitt = time()
            data, addr = self.Server.recvfrom(4096)
            Waitt = time() - Waitt

            self.__class__.Time['Waiting'] += Waitt
            del Waitt

            self.__class__.Count['Accepted'] += 1
            Thread(target=self.Task, args=(self, data, addr)).start()
            del data, addr

    @staticmethod
    def Task(self, request, addr):
        self.__class__.Count['Processed']['All'] += 1
        Allt = time()
        et = tuple(self.__class__.TRUE_REPLY_DOMAINS &
                   self.__class__.FALSE_REPLY_DOMAINS)
        if len(et) != 0:
            for i in et:
                try:
                    self.__class__.TRUE_REPLY_DOMAINS.remove(i)
                except:
                    pass

                try:
                    self.__class__.FALSE_REPLY_DOMAINS.remove(i)
                except:
                    pass
            del i
        del et

        domain, domainhash = None, None

        domain = DNSPackDomainGet(request)[0][0]
        domainhash = hash((domain, self.__class__.L))

        assert not (
            domainhash in self.__class__.TRUE_REPLY_DOMAINS and domainhash in self.__class__.FALSE_REPLY_DOMAINS)
        TrueReply = None

        Cmode = None
        CheckUpstream = None
        if domainhash in self.__class__.TRUE_REPLY_DOMAINS:
            TrueReply = True
            Cmode = 'cache'
            Checkt = 0.0
        elif domainhash in self.__class__.FALSE_REPLY_DOMAINS:
            TrueReply = False
            Cmode = 'cache'
            Checkt = 0.0
        else:
            Cmode = self.CMode
            Checkt = time()
            if Cmode == 'root':
                if self.Server.family is AF_INET:
                    checkup0 = self.__class__.RootServer4
                elif self.Server.family is AF_INET6:
                    checkup0 = self.__class__.RootServer6
                else:
                    raise ValueError

                CheckUpstream = checkup0
                raw_checkdata, CheckUpstream = Send(
                    request, checkup0, self.Server.family)

                authcount = raw_checkdata[9:10]
                answercount = raw_checkdata[7:8]
                checktype = (
                    raw_checkdata[(DNSPackDomainGet(raw_checkdata)[1]):])[8:9]

                if authcount != b'\x00' and answercount == b'\x00' and checktype == b'\x02':
                    TrueReply = True
                else:
                    TrueReply = False
            elif Cmode == 'reset':
                with socket(self.Server.family, SOCK_STREAM) as c:
                    c.settimeout(10)

                    if c.family is AF_INET:
                        checkup1 = self.__class__.CheckServer4[0]
                    elif c.family is AF_INET6:
                        checkup1 = self.__class__.CheckServer6[0]
                    else:
                        raise ValueError
                    self.__class__.CheckServer4.append(
                        self.__class__.CheckServer4.pop(0))
                    self.__class__.CheckServer6.append(
                        self.__class__.CheckServer6.pop(0))

                    CheckUpstream = checkup1
                    c.connect(checkup1)
                    c.send(b'\x00\x00'+request)
                    sleep(0.2)
                    try:
                        c.send(b'')
                    except ConnectionResetError:
                        TrueReply = False
                    else:
                        TrueReply = True
                del c
            elif Cmode == 'dual':
                with socket(self.Server.family, SOCK_DGRAM) as c:
                    c.settimeout(10)

                    if c.family is AF_INET:
                        checkup2 = self.__class__.CheckServer4[0]
                    elif c.family is AF_INET6:
                        checkup2 = self.__class__.CheckServer6[0]
                    else:
                        raise ValueError

                    self.__class__.CheckServer4.append(
                        self.__class__.CheckServer4.pop(0))
                    self.__class__.CheckServer6.append(
                        self.__class__.CheckServer6.pop(0))

                    CheckUpstream = checkup2
                    c.sendto(request, checkup2)
                    while True:
                        rdata, raddr = c.recvfrom(4096)
                        if raddr == checkup2:
                            c.settimeout(0.3)
                            try:
                                c.recvfrom(4096)
                            except timeout:
                                TrueReply = True
                            else:
                                TrueReply = False
                            break
                        del rdata, raddr
            elif Cmode == 'void':
                def getrandip(f, v4h='2', v6h='2620'):
                    assert f in (AF_INET, AF_INET6)
                    assert type(v4h) is str and type(v6h) is str

                    from random import randint

                    ip = []
                    if f is AF_INET:
                        ip.append(v4h)

                        for i in range(0, 3):
                            ip.append(str(randint(0, 255)))
                        del i

                        ip = '.'.join(ip)
                    elif f is AF_INET6:
                        ip.append(v6h)

                        for i in range(0, 7):
                            ip.append(format(randint(0x0, 0xffff), 'x'))
                        del i

                        ip = ':'.join(ip)
                    else:
                        raise ValueError
                    assert type(ip) is str

                    return (ip, 53)

                checkup3 = []
                for i in range(0, 16):
                    checkup3.append(getrandip(self.Server.family))
                del i
                checkup3 = tuple(checkup3)

                try:
                    CheckUpstream = Send(
                        request, checkup3, self.Server.family, timeout=1)[1]
                except timeout:
                    TrueReply = True
                else:
                    TrueReply = False
            else:
                raise ValueError
            Checkt = time() - Checkt
            self.__class__.Count['Processed']['Checked'][0] += 1
            self.__class__.Time['Processing']['Check'] += Checkt

        if Cmode == 'cache':
            self.__class__.Count['Processed']['Checked'][1] += 1

        assert not (TrueReply is None)

        if TrueReply is True:
            (domainhash not in self.__class__.FALSE_REPLY_DOMAINS) and self.__class__.TRUE_REPLY_DOMAINS.add(
                domainhash)
            upstream, upstreamfamily, upstreamtype = self.LocalUpstream, self.localupf, 'local'
        elif TrueReply is False:
            (domainhash not in self.__class__.TRUE_REPLY_DOMAINS) and self.__class__.FALSE_REPLY_DOMAINS.add(
                domainhash)
            upstream, upstreamfamily, upstreamtype = self.RemoteUpstream, self.remoteupf, 'remote'
        else:
            raise ValueError

        up4, up6 = [], []
        for i in upstream:
            if '.' in i[0]:
                up4.append(i)
            elif ':' in i[0]:
                up6.append(i)
            else:
                raise ValueError
        del i
        up4, up6 = tuple(up4), tuple(up6)

        if upstreamfamily is AF_INET:
            upstream = up4
        elif upstreamfamily is AF_INET6:
            upstream = up6
        else:
            raise ValueError
        del up4, up6

        Queryt = time()
        if self.QMode == 'non-53port':
            reply, upaddr = Send(request, upstream, upstreamfamily)
        elif self.QMode == 'non-first':
            di = -1
            if TrueReply is False:
                di = 0.2
            reply, upaddr = Send(request, upstream, upstreamfamily, drop=di)
        else:
            raise ValueError
        Queryt = time() - Queryt
        self.__class__.Count['Processed']['Queryed'] += 1
        self.__class__.Time['Processing']['Query'] += Queryt
        for i in range(0, 3):
            self.Server.sendto(reply, addr)
        del i
        self.__class__.Count['Processed']['Replyed'] += 1

        disclient = (':' in addr[0] and '['+addr[0] +
                     ']' or addr[0])+':'+str(addr[1])
        disupaddr = (
            ':' in upaddr[0] and '['+upaddr[0]+']' or upaddr[0])+':'+str(upaddr[1])
        if CheckUpstream != None:
            discaddr = (':' in CheckUpstream[0] and '['+CheckUpstream[0] +
                        ']' or CheckUpstream[0])+':'+str(CheckUpstream[1])
        else:
            discaddr = None

        self.__class__.Log.append('{time}; from <{client}>, domain: \'{domain}\'({domainhash}): check upstream: <{checkup}>, check mode: ({cmode}), check use time: [{checktime}]; use {uptype} upstream <{upaddr}>, query mode: ({qmode}), query use time: [{socktime}].'.format(
            time=time(), client=disclient, domain=domain, domainhash=domainhash, checkup=discaddr, cmode=Cmode, qmode=((TrueReply is False) and self.QMode or 'default'), checktime=Checkt, uptype=upstreamtype, upaddr=disupaddr, socktime=Queryt))
        del domain

        Allt = time() - Allt
        Othert = Allt - (Checkt + Queryt)
        self.__class__.Time['Processing']['All'] += Allt
        self.__class__.Time['Processing']['Other'] += Othert


def selfkill(mode=()):
    assert type(mode) is tuple
    mode = tuple(set(mode))
    from os import kill, getpid, _exit, abort

    if len(mode) == 0:
        isall = True
    else:
        isall = False
    while True:
        if isall is True or '_exit' in mode:
            try:
                Thread(target=_exit, args=(0,)).start()
            except:
                pass
        if isall is True or 'kill-15' in mode:
            try:
                Thread(target=kill, args=(0, 15)).start()
            except:
                pass
            try:
                Thread(target=kill, args=(getpid(), 15)).start()
            except:
                pass
        if isall is True or 'kill-9' in mode:
            try:
                Thread(target=kill, args=(0, 9)).start()
            except:
                pass
            try:
                Thread(target=kill, args=(getpid(), 9)).start()
            except:
                pass
        if isall is True or 'abort' in mode:
            try:
                Thread(target=abort).start()
            except:
                pass
        isall = True


def main(argv=__import__('sys').argv):
    default_port = 52

    default_cmode = 'root'
    default_qmode = 'non-53port'

    default_local_upstream = (
        ('118.118.118.118', 53),
        #('114.114.114.114', 53),
        #('180.76.76.76', 53),
        ('223.5.5.5', 53),
        ('223.6.6.6', 53),
        #('101.226.4.6', 53),
        #('218.30.118.6', 53)
    )

    default_remote_upstream = (
        (
            #Quad9: dns10.quad9.net
            ('9.9.9.10', 9953),
            ('149.112.112.10', 9953),
            ('2620:fe::10', 9953),
            ('2620:fe::fe:10', 9953),
            #OpenDNS: dns.opendns.com
            ('208.67.222.222', 443),
            ('208.67.220.220', 5353)
        ),
        (
            # one.one.one.one
            ('1.0.0.1', 53),
            # dns.google
            ('8.8.8.8', 53),
            ('8.8.4.4', 53),
            # dns10.quad9.net
            ('9.9.9.10', 53),
            ('149.112.112.10', 53),
            ('2620:fe::10', 53),
            ('2620:fe::fe:10', 53),
            # dns.opendns.com
            ('208.67.222.222', 53),
            ('208.67.220.220', 53)
        )
    )

    point = 0

    showlog = False
    try:
        argv[1]
    except:
        pass
    else:
        if argv[1] == '--showlog':
            showlog = True
            point += 1

    try:
        argv[point+1]
    except:
        port = default_port
    else:
        port = argv[point+1]

    port = int(port)
    assert 0 < port <= 65535

    try:
        argv[point+2]
    except:
        cmode = default_cmode
        qmode = default_qmode
    else:
        tmp = tuple(argv[point+2].split(','))
        assert len(tmp) == 2
        cmode = tmp[0]
        qmode = tmp[1]
        del tmp

    cmode = str(cmode).lower()
    qmode = str(qmode).lower()

    try:
        argv[point+3]
    except:
        local_upstream = default_local_upstream
    else:
        tmp = tuple(str(argv[point+3]).split(';'))

        tmp1 = []
        for i in tmp:
            tmp1.append(i)
        del i, tmp
        local_upstream = []
        for i in tmp1:
            i = i.split(',')
            assert len(i) == 2
            local_upstream.append((str(i[0]), int(i[1])))
        del i, tmp1

    local_upstream = tuple(local_upstream)

    try:
        argv[point+4]
    except:
        if qmode == 'non-53port':
            remote_upstream = default_remote_upstream[0]
        elif qmode == 'non-first':
            remote_upstream = default_remote_upstream[1]
        else:
            remote_upstream = default_remote_upstream
    else:
        tmp = tuple(str(argv[point+4]).split(';'))

        tmp1 = []
        for i in tmp:
            tmp1.append(i)
        del i, tmp
        remote_upstream = []
        for i in tmp1:
            i = i.split(',')
            assert len(i) == 2
            remote_upstream.append((str(i[0]), int(i[1])))
        del i, tmp1

    remote_upstream = tuple(remote_upstream)

    with socket(AF_INET, SOCK_DGRAM) as s:
        s.bind(('0.0.0.0', port))

        for i in range(0, 16):
            Master(cmode, qmode, s, local_upstream, remote_upstream).start()
        del i

        def clear():
            print('\033c', end='')

        def seelog(t=1, other=False):
            num = 0
            while True:
                if other is True:
                    seeall()
                log = Master.Log.copy()
                for i in log[num:]:
                    print(i)
                del i
                num = len(log)
                while len(Master.Log) == num:
                    sleep(t)

        def seecount():
            print('{0}; Count: {1}'.format(time(), Master.Count))

        def seetime():
            print('{0}; Time: {1}'.format(time(), Master.Time))

        def seelen():
            print('{0}; Cache Length: TRUE_REPLY_DOMAINS: {t}, FALSE_REPLY_DOMAINS: {f}'.format(
                time(), t=len(Master.TRUE_REPLY_DOMAINS), f=len(Master.FALSE_REPLY_DOMAINS)))
            print('{0}; Log Length:'.format(time()), len(Master.Log))

        def seeall():
            seelen()
            seecount()
            seetime()

        def auto(t=1):
            while True:
                clear()
                seeall()
                sleep(t)

        if showlog is True:
            seelog(other=True)
            return

        exit = quit = selfkill
        while True:
            try:
                exec(input(';; '))
            except BaseException as e:
                print('{', type(e).__name__, *(len(str(e))
                                               != 0 and (': ', e) or ''), '}', sep='')


if __name__ == '__main__':
    try:
        main()
    except(SystemExit, KeyboardInterrupt, EOFError):
        selfkill()