Python dnslib.A Examples

def test_on_upstream_response_BD(greendns):
    qname = "google.com"
    s = init_greendns_session(greendns, qname, dnslib.QTYPE.A)
    res = dnslib.DNSRecord(dnslib.DNSHeader(qr=1, aa=1, ra=1),
                           q=dnslib.DNSQuestion(qname),
                           a=dnslib.RR(qname,
                                       rdata=dnslib.A("1.2.3.4"),
                                       ttl=3))
    s.server_resps[local_dns1] = bytes(res.pack())
    resp = greendns.on_upstream_response(s, local_dns1)
    assert not resp

    res.rr[0].rdata = dnslib.A("172.217.24.14")
    s.server_resps[foreign_dns] = bytes(res.pack())
    resp = greendns.on_upstream_response(s, foreign_dns)
    assert resp
    d = dnslib.DNSRecord.parse(resp)
    assert str(d.rr[0].rdata) == "172.217.24.14" 

def test_on_upstream_response_AD(greendns):
    qname = "www.microsoft.com"
    s = init_greendns_session(greendns, qname, dnslib.QTYPE.A)
    res = dnslib.DNSRecord(dnslib.DNSHeader(qr=1, aa=1, ra=1),
                           q=dnslib.DNSQuestion(qname),
                           a=dnslib.RR(qname,
                                       rdata=dnslib.A("184.85.123.14"),
                                       ttl=3))
    s.server_resps[foreign_dns] = bytes(res.pack())
    resp = greendns.on_upstream_response(s, foreign_dns)
    assert not resp

    res.rr[0].rdata = dnslib.A("183.136.212.50")
    s.server_resps[local_dns1] = bytes(res.pack())
    resp = greendns.on_upstream_response(s, local_dns1)
    assert resp
    d = dnslib.DNSRecord.parse(resp)
    assert str(d.rr[0].rdata) == "183.136.212.50" 

def test_on_upstream_response_AC(greendns):
    qname = "www.coding.net"
    s = init_greendns_session(greendns, qname, dnslib.QTYPE.A)
    res = dnslib.DNSRecord(dnslib.DNSHeader(qr=1, aa=1, ra=1),
                           q=dnslib.DNSQuestion(qname),
                           a=dnslib.RR(qname,
                                       rdata=dnslib.A("120.132.59.101"),
                                       ttl=3))
    s.server_resps[foreign_dns] = bytes(res.pack())
    resp = greendns.on_upstream_response(s, foreign_dns)
    assert not resp

    res.rr[0].rdata = dnslib.A("219.146.244.91")
    s.server_resps[local_dns1] = bytes(res.pack())
    resp = greendns.on_upstream_response(s, local_dns1)
    assert resp
    d = dnslib.DNSRecord.parse(resp)
    assert str(d.rr[0].rdata) == "219.146.244.91" 

def test_on_upstream_response_BC(greendns):
    qname = "www.x.net"
    s = init_greendns_session(greendns, qname, dnslib.QTYPE.A)
    res = dnslib.DNSRecord(dnslib.DNSHeader(qr=1, aa=1, ra=1),
                           q=dnslib.DNSQuestion(qname),
                           a=dnslib.RR(qname,
                                       rdata=dnslib.A("8.8.8.8"),
                                       ttl=3))
    s.server_resps[local_dns1] = bytes(res.pack())
    resp = greendns.on_upstream_response(s, local_dns1)
    assert not resp

    res.rr[0].rdata = dnslib.A("1.2.4.8")
    s.server_resps[foreign_dns] = bytes(res.pack())
    resp = greendns.on_upstream_response(s, foreign_dns)
    assert resp
    d = dnslib.DNSRecord.parse(resp)
    assert str(d.rr[0].rdata) == "1.2.4.8" 

def test_on_upstream_response_invalid_A(greendns):
    qname = "www.x.net"
    s = init_greendns_session(greendns, qname, dnslib.QTYPE.A)
    res = dnslib.DNSRecord(dnslib.DNSHeader(qr=1, aa=1, ra=1),
                           q=dnslib.DNSQuestion(qname),
                           a=dnslib.RR(qname,
                                       rdata=dnslib.A("1.2.4.8"),
                                       ttl=3))
    s.server_resps[local_dns1] = bytes(res.pack())
    resp = greendns.on_upstream_response(s, local_dns1)
    d = dnslib.DNSRecord.parse(resp)
    assert str(d.rr[0].rdata) == "1.2.4.8"

    s.server_resps[local_dns1] = b'123456'
    resp = greendns.on_upstream_response(s, local_dns1)
    assert not resp 

def do_add_trigger(self, line):
        """Add trigger\nUsage: add_trigger <REGEXP> <ACTION> [<ACTION ARG>]\n\nAction Values: 0 = Pass; 1 = Offline; 2 = Alert Message; 3 = DDoS; 4 = Hijack\n\nExamples:
        add_trigger ".*com" 1                                      -- This will disable access to any *.COM website
        add_trigger ".*google*" 2 "Sorry, Google is closed today!" -- Display alert if trying to browse any URL that matches: *GOOGLE* Regexp
        add_trigger ".*" 3 "www.ikotler.org:80"                    -- Will start TCP SYN DDoS on http://www.ikotler.org on 80/tcp
        add_trigger ".*token*" 4 "grin.host:123"                   -- Redirect every URL contains *token* to PROXY: grin.host on 123/tcp
        """
        global triggers

        if not line:
            self.do_help("add_trigger")
        else:
            trigger_args = shlex.split(line)
            trigger_data = [len(triggers), trigger_args[0], int(trigger_args[1])]
            if trigger_args[1] == "1":
                trigger_data.append("N/A")
            else:
                trigger_data.append(trigger_args[2])
            triggers.append(trigger_data)
            print "Done!" 

def send_request(self, id, domain, server):
        try:
            d = DNSRecord(DNSHeader(id))
            d.add_question(DNSQuestion(domain, QTYPE.A))
            req4_pack = d.pack()

            d = DNSRecord(DNSHeader(id))
            d.add_question(DNSQuestion(domain, QTYPE.AAAA))
            req6_pack = d.pack()

            self.sock.sendto(req4_pack, (server, 53))
            # xlog.debug("send req:%s to:%s", domain, server)

            self.sock.sendto(req6_pack, (server, 53))
            # xlog.debug("send req:%s to:%s", domain, server)
        except Exception as e:
            xlog.warn("send_request except:%r", e) 

def on_udp_query(self, rsock, req_data, addr):
        start_time = time.time()
        try:
            request = DNSRecord.parse(req_data)
            if len(request.questions) != 1:
                xlog.warn("query num:%d %s", len(request.questions), request)
                return

            domain = utils.to_bytes(str(request.questions[0].qname))

            if domain.endswith(b"."):
                domain = domain[:-1]

            type = request.questions[0].qtype
            if type not in [1, 28]:
                xlog.info("direct_query:%s type:%d", domain, type)
                return self.direct_query(rsock, request, addr)

            xlog.debug("DNS query:%s type:%d from %s", domain, type, addr)

            ips = self.query(domain, type)
            if not ips:
                xlog.debug("query:%s type:%d from:%s, get fail, cost:%d", domain, type, addr,
                           (time.time() - start_time) * 1000)

            reply = DNSRecord(DNSHeader(id=request.header.id, qr=1, aa=1, ra=1, auth=1), q=request.q)
            ips = utils.to_bytes(ips)
            for ip_cn in ips:
                ipcn_p = ip_cn.split(b"|")
                ip = ipcn_p[0]
                if b"." in ip and type == 1:
                    reply.add_answer(RR(domain, ttl=60, rdata=A(ip)))
                elif b":" in ip and type == 28:
                    reply.add_answer(RR(domain, rtype=type, ttl=60, rdata=AAAA(ip)))
            res_data = reply.pack()

            rsock.sendto(res_data, addr)
            xlog.debug("query:%s type:%d from:%s, return ip num:%d cost:%d", domain, type, addr,
                       len(reply.rr), (time.time()-start_time)*1000)
        except Exception as e:
            xlog.exception("on_query except:%r", e) 

def handle(self):
        data = self.request[0]
        d = dnslib.DNSRecord.parse(data)
        id = d.header.id
        sock = self.request[1]
        res = dnslib.DNSRecord(dnslib.DNSHeader(qr=1, aa=1, ra=1, id=id),
                               q=dnslib.DNSQuestion(qname),
                               a=dnslib.RR(qname,
                                           rdata=dnslib.A("101.226.103.106"),
                                           ttl=3))
        sock.sendto(bytes(res.pack()), self.client_address) 

def handle(self):
        data = self.request[0]
        d = dnslib.DNSRecord.parse(data)
        id = d.header.id
        sock = self.request[1]
        res = dnslib.DNSRecord(dnslib.DNSHeader(qr=1, aa=1, ra=1, id=id),
                               q=dnslib.DNSQuestion(qname),
                               a=dnslib.RR(qname,
                                           rdata=dnslib.A("182.254.8.146"),
                                           ttl=3))
        time.sleep(0.3)
        sock.sendto(bytes(res.pack()), self.client_address) 

def test_on_client_request_invalid(greendns):
    s = init_greendns_session(greendns, "google.com", dnslib.QTYPE.A)
    s.req_data = b'123456'
    is_continue, raw_resp = greendns.on_client_request(s)
    assert not is_continue
    assert not raw_resp 

def test_on_client_request_without_cached(greendns):
    s = init_greendns_session(greendns, "google.com", dnslib.QTYPE.A)
    is_continue, raw_resp = greendns.on_client_request(s)
    assert is_continue
    assert not raw_resp 

def test_on_client_request_with_cached(greendns):
    qname = "qq.com"
    id = 1024
    s = init_greendns_session(greendns, qname, dnslib.QTYPE.A, id)
    res = dnslib.DNSRecord(dnslib.DNSHeader(qr=1, aa=1, ra=1),
                           q=dnslib.DNSQuestion(qname),
                           a=dnslib.RR(qname,
                                       rdata=dnslib.A("101.226.103.106"),
                                       ttl=3))
    greendns.cache.add(("qq.com.", 1), res, 3)
    is_continue, raw_resp = greendns.on_client_request(s)
    assert not is_continue
    assert raw_resp
    d = dnslib.DNSRecord.parse(raw_resp)
    assert d.header.id == id 

def get_reply_record(self, data):
        request = dnslib.DNSRecord.parse(data)
        qname = str(request.q.qname).lower()
        qtype = request.q.qtype
        dnsservers = self.dns_servers
        if qname.endswith('.in-addr.arpa'):
            ipaddr = '.'.join(reversed(qname[:-13].split('.')))
            record = dnslib.DNSRecord(header=dnslib.DNSHeader(id=request.header.id, qr=1,aa=1,ra=1), a=dnslib.RR(qname, rdata=dnslib.A(ipaddr)))
            return record
        if 'USERDNSDOMAIN' in os.environ:
            user_dnsdomain = '.' + os.environ['USERDNSDOMAIN'].lower()
            if qname.endswith(user_dnsdomain):
                qname = qname[:-len(user_dnsdomain)]
                if '.' not in qname:
                    if not self.dns_intranet_servers:
                        logging.warning('qname=%r is a plain hostname, need intranet dns server!!!', qname)
                        return dnslib.DNSRecord(header=dnslib.DNSHeader(id=request.header.id, rcode=3))
                    qname += user_dnsdomain
                    dnsservers = self.dns_intranet_servers
        try:
            return self.dns_cache.get((qname, qtype))
        except KeyError:
            pass
        try:
            dns_resolve = dnslib_resolve_over_tcp if qname.endswith(self.dns_tcpover) else dnslib_resolve_over_udp
            kwargs = {'blacklist': self.dns_blacklist, 'turstservers': self.dns_trust_servers}
            record = dns_resolve(request, dnsservers, self.dns_timeout, **kwargs)
            ttl = max(x.ttl for x in record.rr) if record.rr else 600
            self.dns_cache.set((qname, qtype), record, ttl * 2)
            return record
        except socket.gaierror as e:
            logging.warning('resolve %r failed: %r', qname, e)
            return dnslib.DNSRecord(header=dnslib.DNSHeader(id=request.header.id, rcode=3)) 

def dnslib_resolve_over_tcp(query, dnsservers, timeout, **kwargs):
    """dns query over tcp"""
    if not isinstance(query, (basestring, dnslib.DNSRecord)):
        raise TypeError('query argument requires string/DNSRecord')
    blacklist = kwargs.get('blacklist', ())
    def do_resolve(query, dnsserver, timeout, queobj):
        if isinstance(query, basestring):
            qtype = dnslib.QTYPE.AAAA if ':' in dnsserver else dnslib.QTYPE.A
            query = dnslib.DNSRecord(q=dnslib.DNSQuestion(query, qtype=qtype))
        query_data = query.pack()
        sock_family = socket.AF_INET6 if ':' in dnsserver else socket.AF_INET
        sock = socket.socket(sock_family)
        rfile = None
        try:
            sock.settimeout(timeout or None)
            sock.connect(parse_hostport(dnsserver, 53))
            sock.send(struct.pack('>h', len(query_data)) + query_data)
            rfile = sock.makefile('r', 1024)
            reply_data_length = rfile.read(2)
            if len(reply_data_length) < 2:
                raise socket.gaierror(11004, 'getaddrinfo %r from %r failed' % (query, dnsserver))
            reply_data = rfile.read(struct.unpack('>h', reply_data_length)[0])
            record = dnslib.DNSRecord.parse(reply_data)
            iplist = [str(x.rdata) for x in record.rr if x.rtype in (1, 28, 255)]
            if any(x in blacklist for x in iplist):
                logging.debug('query=%r dnsserver=%r record bad iplist=%r', query, dnsserver, iplist)
                raise socket.gaierror(11004, 'getaddrinfo %r from %r failed' % (query, dnsserver))
            else:
                logging.debug('query=%r dnsserver=%r record iplist=%s', query, dnsserver, iplist)
                queobj.put(record)
        except socket.error as e:
            logging.debug('query=%r dnsserver=%r failed %r', query, dnsserver, e)
            queobj.put(e)
        finally:
            if rfile:
                rfile.close()
            sock.close()
    queobj = Queue.Queue()
    for dnsserver in dnsservers:
        thread.start_new_thread(do_resolve, (query, dnsserver, timeout, queobj))
    for i in range(len(dnsservers)):
        try:
            result = queobj.get(timeout)
        except Queue.Empty:
            raise socket.gaierror(11004, 'getaddrinfo %r from %r failed' % (query, dnsservers))
        if result and not isinstance(result, Exception):
            return result
        elif i == len(dnsservers) - 1:
            logging.warning('dnslib_resolve_over_tcp %r with %s return %r', query, dnsservers, result)
    raise socket.gaierror(11004, 'getaddrinfo %r from %r failed' % (query, dnsservers))