111 lines
3.7 KiB
Python
111 lines
3.7 KiB
Python
# -*- coding: utf-8 -*-
|
|
from __future__ import print_function
|
|
import time
|
|
|
|
import dnslib
|
|
from dnslib import server
|
|
from dnslib import dns
|
|
|
|
from IPy import IP
|
|
|
|
from onedns import utils
|
|
|
|
|
|
class DynamicResolver(server.BaseResolver):
|
|
"""
|
|
Dynamic In-Memory DNS Resolver
|
|
"""
|
|
|
|
def __init__(self, domain, one_kwargs={}):
|
|
"""
|
|
Initialise resolver from zone list
|
|
Stores RRs as a list of (label, type, rr) tuples
|
|
"""
|
|
self.domain = domain
|
|
self.zone = []
|
|
self._tcp_server = None
|
|
self._udp_server = None
|
|
|
|
def resolve(self, request, handler):
|
|
"""
|
|
Respond to DNS request - parameters are request packet & handler.
|
|
Method is expected to return DNS response
|
|
"""
|
|
reply = request.reply()
|
|
qname = request.q.qname
|
|
qtype = dnslib.QTYPE[request.q.qtype]
|
|
A_RECORDS = ['A', 'AAAA']
|
|
for name, rtype, rr in self.zone:
|
|
# Check if label & type match
|
|
if qname == name and (qtype in [rtype, 'ANY'] or rtype == 'CNAME'):
|
|
reply.add_answer(rr)
|
|
# Check for A/AAAA records associated with reply and
|
|
# add in additional section
|
|
if rtype in ['CNAME', 'NS', 'MX', 'PTR']:
|
|
for a_name, a_rtype, a_rr in self.zone:
|
|
if a_name == rr.rdata.label and a_rtype in A_RECORDS:
|
|
reply.add_ar(a_rr)
|
|
if not reply.rr:
|
|
reply.header.rcode = dnslib.RCODE.NXDOMAIN
|
|
return reply
|
|
|
|
def clear(self):
|
|
self.zone = []
|
|
|
|
def add_host(self, name, ip):
|
|
self._add_forward(name, ip)
|
|
self._add_reverse(ip, name)
|
|
|
|
def _get_fqdn(self, name):
|
|
return utils.get_fqdn(name, self.domain)
|
|
|
|
def _add_forward(self, name, ip):
|
|
f = dnslib.RR(rname=dnslib.DNSLabel(self._get_fqdn(name)),
|
|
rtype=dnslib.QTYPE.reverse['A'],
|
|
rclass=dnslib.CLASS.reverse['IN'],
|
|
rdata=dns.A(ip))
|
|
self.zone.append((f.rname, 'A', f))
|
|
|
|
def _add_reverse(self, ip, name):
|
|
ip = IP(ip)
|
|
r = dnslib.RR(rname=dnslib.DNSLabel(ip.reverseName()),
|
|
rtype=dnslib.QTYPE.reverse['PTR'],
|
|
rclass=dnslib.CLASS.reverse['IN'],
|
|
rdata=dns.PTR(self._get_fqdn(name)))
|
|
self.zone.append((r.rname, 'PTR', r))
|
|
|
|
def start(self, dns_address='0.0.0.0', dns_port=53,
|
|
api_address='127.0.0.1', api_port=8000, tcp=False, udplen=0,
|
|
log="request,reply,truncated,error", log_prefix=False):
|
|
logger = server.DNSLogger(log, log_prefix)
|
|
|
|
print("Starting OneDNS (%s:%d) [%s]" %
|
|
(dns_address or "*", dns_port, "UDP/TCP" if tcp else "UDP"))
|
|
|
|
server.DNSHandler.udplen = udplen
|
|
|
|
self._udp_server = server.DNSServer(self, port=dns_port,
|
|
address=dns_address, logger=logger)
|
|
self._udp_server.start_thread()
|
|
|
|
if tcp:
|
|
self._tcp_server = server.DNSServer(self, port=dns_port,
|
|
address=dns_address, tcp=True,
|
|
logger=logger)
|
|
self._tcp_server.start_thread()
|
|
|
|
def close(self):
|
|
for srv in [self._tcp_server, self._udp_server]:
|
|
if srv:
|
|
srv.stop()
|
|
srv.server.socket.close()
|
|
|
|
def daemon(self, *args, **kwargs):
|
|
testing = kwargs.pop('testing', False)
|
|
if self._udp_server is None or not self._udp_server.isAlive():
|
|
self.start(*args, **kwargs)
|
|
while self._udp_server.isAlive():
|
|
time.sleep(1)
|
|
if testing:
|
|
break
|