def _init_udp(self):
"""
Create a new Datagram instance and listen on a socket.
"""
self._udp = Datagram(engine=self.engine)
self._udp.on_read = self.receive_message
start = port = random.randrange(10005, 65535)
while True:
try:
self._udp.listen(('', port))
break
except Exception:
port += 1
if port > 65535:
port = 10000
if port == start:
raise Exception("Can't listen on any port.")
def _init_udp(self):
"""
Create a new Datagram instance and listen on a socket.
"""
self._udp = Datagram()
self._udp.on_read = self.receive_message
start = port = random.randrange(10005, 65535)
while True:
try:
self._udp.listen(("", port))
break
except Exception:
port += 1
if port > 65535:
port = 10000
if port == start:
raise Exception("Can't listen on any port.")
class Resolver(object):
"""
The Resolver class generates DNS messages, sends them to remote servers,
and processes any responses. The bulk of the heavy lifting is done in
DNSMessage and the RDATA handling functions, however.
========= ============
Argument Description
========= ============
servers *Optional.* A list of DNS servers to query. If a list isn't provided, Pants will attempt to retrieve a list of servers from the OS, falling back to a list of default servers if none are available.
========= ============
"""
def __init__(self, servers=None):
self.servers = servers or list_dns_servers()
# Internal State
self._messages = {}
self._cache = {}
self._queries = {}
self._tcp = {}
self._udp = None
self._last_id = -1
def _safely_call(self, callback, *args, **kwargs):
try:
callback(*args, **kwargs)
except Exception:
log.exception("Error calling callback for DNS result.")
def _error(self, message, err=DNS_TIMEOUT):
if not message in self._messages:
return
if message in self._tcp:
try:
self._tcp[message].close()
except Exception:
pass
del self._tcp[message]
callback, message, df_timeout, media, data = self._messages[message]
del self._messages[message.id]
try:
df_timeout.cancel()
except Exception:
pass
if err == DNS_TIMEOUT and data:
self._safely_call(callback, DNS_OK, data)
else:
self._safely_call(callback, err, None)
def _init_udp(self):
"""
Create a new Datagram instance and listen on a socket.
"""
self._udp = Datagram()
self._udp.on_read = self.receive_message
start = port = random.randrange(10005, 65535)
while True:
try:
self._udp.listen(("", port))
break
except Exception:
port += 1
if port > 65535:
port = 10000
if port == start:
raise Exception("Can't listen on any port.")
def send_message(self, message, callback=None, timeout=10, media=None):
"""
Send an instance of DNSMessage to a DNS server, and call the provided
callback when a response is received, or if the action times out.
========= ======== ============
Argument Default Description
========= ======== ============
message The :class:`DNSMessage` instance to send to the server.
callback None *Optional.* The function to call once the response has been received or the attempt has timed out.
timeout 10 *Optional.* How long, in seconds, to wait before timing out.
media None *Optional.* Whether to use UDP or TCP. UDP is used by default.
========= ======== ============
"""
while message.id is None or message.id in self._messages:
self._last_id += 1
if self._last_id > 65535:
self._last_id = 0
message.id = self._last_id
# Timeout in timeout seconds.
df_timeout = pants.engine.defer(timeout, self._error, message.id)
# Send the Message
msg = str(message)
if media is None:
media = "udp"
# if len(msg) > 512:
# media = 'tcp'
# else:
# media = 'udp'
# Store Info
self._messages[message.id] = callback, message, df_timeout, media, None
if media == "udp":
if self._udp is None:
self._init_udp()
try:
self._udp.write(msg, (self.servers[0], DNS_PORT))
except Exception:
# Pants gummed up. Try again.
self._next_server(message.id)
pants.engine.defer(0.5, self._next_server, message.id)
else:
tcp = self._tcp[message.id] = _DNSStream(self, message.id)
tcp.connect((self.servers[0], DNS_PORT))
def _next_server(self, id):
if not id in self._messages or id in self._tcp:
return
# Cycle the list.
self.servers.append(self.servers.pop(0))
msg = str(self._messages[id][1])
try:
self._udp.write(msg, (self.servers[0], DNS_PORT))
except Exception:
try:
self._udp.close()
except Exception:
pass
del self._udp
self._init_udp()
self._udp.write(msg, (self.servers[0], DNS_PORT))
def receive_message(self, data):
if not isinstance(data, DNSMessage):
try:
data = DNSMessage.from_string(data)
except TooShortError:
if len(data) < 2:
return
id = struct.unpack("!H", data[:2])
if not id in self._messages:
return
self._error(id, err=DNS_BADRESPONSE)
return
if not data.id in self._messages:
return
callback, message, df_timeout, media, _ = self._messages[data.id]
# if data.tc and media == 'udp':
# self._messages[data.id] = callback, message, df_timeout, 'tcp', data
# tcp = self._tcp[data.id] = _DNSStream(self, message.id)
# tcp.connect((self.servers[0], DNS_PORT))
# return
if not data.server:
if self._udp and isinstance(self._udp.remote_addr, tuple):
data.server = "%s:%d" % self._udp.remote_addr
else:
data.server = "%s:%d" % (self.servers[0], DNS_PORT)
try:
df_timeout.cancel()
except Exception:
pass
del self._messages[data.id]
self._safely_call(callback, DNS_OK, data)
def query(self, name, qtype=A, qclass=IN, callback=None, timeout=10, allow_cache=True, allow_hosts=True):
"""
Make a DNS request of the given QTYPE for the given name.
============ ======== ============
Argument Default Description
============ ======== ============
name The name to query.
qtype A *Optional.* The QTYPE to query.
qclass IN *Optional.* The QCLASS to query.
callback None *Optional.* The function to call when a response for the query has been received, or when the request has timed out.
timeout 10 *Optional.* The time, in seconds, to wait before timing out.
allow_cache True *Optional.* Whether or not to use the cache. If you expect to be performing thousands of requests, you may want to disable the cache to avoid excess memory usage.
allow_hosts True *Optional.* Whether or not to use any records gathered from the OS hosts file.
============ ======== ============
"""
if allow_hosts:
if host_time + 30 < time.time():
load_hosts()
cname = None
if name in self._cache and CNAME in self._cache[name]:
cname = self._cache[name][CNAME]
if qtype == A and name in hosts[A]:
self._safely_call(callback, DNS_OK, cname, None, (hosts[A][name],))
elif qtype == AAAA and name in hosts[AAAA]:
self._safely_call(callback, DNS_OK, cname, None, (hosts[AAAA][name],))
if allow_cache and name in self._cache and (qtype, qclass) in self._cache[name]:
cname = None
if CNAME in self._cache[name]:
cname = self._cache[name][CNAME]
death, ttl, rdata = self._cache[name][(qtype, qclass)]
if death < time.time():
# Clear out the old record.
del self._cache[name][(qtype, qclass)]
else:
if callback:
self._safely_call(callback, DNS_OK, cname, ttl, rdata)
return
# Build a message and add our question.
m = DNSMessage()
m.questions.append((name, qtype, qclass))
# Make the function for handling our response.
def handle_response(status, data):
cname = None
# TTL is 30 by default, so answers with no records we want will be
# repeated, but not too often.
ttl = 30
if not data:
self._safely_call(callback, status, None, None, None)
return
rdata = []
for (aname, atype, aclass, attl, ardata) in data.answers:
if atype == CNAME:
cname = ardata[0]
if atype == qtype and aclass == qclass:
ttl = attl
if len(ardata) == 1:
rdata.append(ardata[0])
else:
rdata.append(ardata)
rdata = tuple(rdata)
if allow_cache:
if not name in self._cache:
self._cache[name] = {}
if cname:
self._cache[name][CNAME] = cname
self._cache[name][(qtype, qclass)] = time.time() + ttl, ttl, rdata
if data.rcode != DNS_OK:
status = data.rcode
self._safely_call(callback, status, cname, ttl, rdata)
# Send it, so we get an ID.
self.send_message(m, handle_response)
class Resolver(object):
"""
The Resolver class generates DNS messages, sends them to remote servers,
and processes any responses. The bulk of the heavy lifting is done in
DNSMessage and the RDATA handling functions, however.
========= ============
Argument Description
========= ============
servers *Optional.* A list of DNS servers to query. If a list isn't provided, Pants will attempt to retrieve a list of servers from the OS, falling back to a list of default servers if none are available.
engine *Optional.* The :class:`pants.engine.Engine` instance to use.
========= ============
"""
def __init__(self, servers=None, engine=None):
self.servers = servers or list_dns_servers()
self.engine = engine or Engine.instance()
# Internal State
self._messages = {}
self._cache = {}
self._queries = {}
self._tcp = {}
self._udp = None
self._last_id = -1
def _safely_call(self, callback, *args, **kwargs):
try:
callback(*args, **kwargs)
except Exception:
log.exception('Error calling callback for DNS result.')
def _error(self, message, err=DNS_TIMEOUT):
if not message in self._messages:
return
if message in self._tcp:
try:
self._tcp[message].close()
except Exception:
pass
del self._tcp[message]
callback, message, df_timeout, media, data = self._messages[message]
del self._messages[message.id]
try:
df_timeout.cancel()
except Exception:
pass
if err == DNS_TIMEOUT and data:
self._safely_call(callback, DNS_OK, data)
else:
self._safely_call(callback, err, None)
def _init_udp(self):
"""
Create a new Datagram instance and listen on a socket.
"""
self._udp = Datagram(engine=self.engine)
self._udp.on_read = self.receive_message
start = port = random.randrange(10005, 65535)
while True:
try:
self._udp.listen(('', port))
break
except Exception:
port += 1
if port > 65535:
port = 10000
if port == start:
raise Exception("Can't listen on any port.")
def send_message(self, message, callback=None, timeout=10, media=None):
"""
Send an instance of DNSMessage to a DNS server, and call the provided
callback when a response is received, or if the action times out.
========= ======== ============
Argument Default Description
========= ======== ============
message The :class:`DNSMessage` instance to send to the server.
callback None *Optional.* The function to call once the response has been received or the attempt has timed out.
timeout 10 *Optional.* How long, in seconds, to wait before timing out.
media None *Optional.* Whether to use UDP or TCP. UDP is used by default.
========= ======== ============
"""
while message.id is None or message.id in self._messages:
self._last_id += 1
if self._last_id > 65535:
self._last_id = 0
message.id = self._last_id
# Timeout in timeout seconds.
df_timeout = self.engine.defer(timeout, self._error, message.id)
# Send the Message
msg = str(message)
if media is None:
media = 'udp'
#if len(msg) > 512:
# media = 'tcp'
#else:
# media = 'udp'
# Store Info
self._messages[message.id] = callback, message, df_timeout, media, None
if media == 'udp':
if self._udp is None:
self._init_udp()
try:
self._udp.write(msg, (self.servers[0], DNS_PORT))
except Exception:
# Pants gummed up. Try again.
self._next_server(message.id)
self.engine.defer(0.5, self._next_server, message.id)
else:
tcp = self._tcp[message.id] = _DNSStream(self, message.id)
tcp.connect((self.servers[0], DNS_PORT))
def _next_server(self, id):
if not id in self._messages or id in self._tcp:
return
# Cycle the list.
self.servers.append(self.servers.pop(0))
msg = str(self._messages[id][1])
try:
self._udp.write(msg, (self.servers[0], DNS_PORT))
except Exception:
try:
self._udp.close()
except Exception:
pass
del self._udp
self._init_udp()
self._udp.write(msg, (self.servers[0], DNS_PORT))
def receive_message(self, data):
if not isinstance(data, DNSMessage):
try:
data = DNSMessage.from_string(data)
except TooShortError:
if len(data) < 2:
return
id = struct.unpack("!H", data[:2])
if not id in self._messages:
return
self._error(id, err=DNS_FORMATERROR)
return
if not data.id in self._messages:
return
callback, message, df_timeout, media, _ = self._messages[data.id]
#if data.tc and media == 'udp':
# self._messages[data.id] = callback, message, df_timeout, 'tcp', data
# tcp = self._tcp[data.id] = _DNSStream(self, message.id)
# tcp.connect((self.servers[0], DNS_PORT))
# return
if not data.server:
if self._udp and isinstance(self._udp.remote_address, tuple):
data.server = '%s:%d' % self._udp.remote_address
else:
data.server = '%s:%d' % (self.servers[0], DNS_PORT)
try:
df_timeout.cancel()
except Exception:
pass
del self._messages[data.id]
self._safely_call(callback, DNS_OK, data)
def query(self,
name,
qtype=A,
qclass=IN,
callback=None,
timeout=10,
allow_cache=True,
allow_hosts=True):
"""
Make a DNS request of the given QTYPE for the given name.
============ ======== ============
Argument Default Description
============ ======== ============
name The name to query.
qtype A *Optional.* The QTYPE to query.
qclass IN *Optional.* The QCLASS to query.
callback None *Optional.* The function to call when a response for the query has been received, or when the request has timed out.
timeout 10 *Optional.* The time, in seconds, to wait before timing out.
allow_cache True *Optional.* Whether or not to use the cache. If you expect to be performing thousands of requests, you may want to disable the cache to avoid excess memory usage.
allow_hosts True *Optional.* Whether or not to use any records gathered from the OS hosts file.
============ ======== ============
"""
if not isinstance(qtype, (list, tuple)):
qtype = (qtype, )
if allow_hosts:
if host_time + 30 < time.time():
load_hosts()
cname = None
if name in self._cache and CNAME in self._cache[name]:
cname = self._cache[name][CNAME]
result = []
if AAAA in qtype and name in hosts[AAAA]:
result.append(hosts[AAAA][name])
if A in qtype and name in hosts[A]:
result.append(hosts[A][name])
if result:
if callback:
self._safely_call(callback, DNS_OK, cname, None,
tuple(result))
return
if allow_cache and name in self._cache:
cname = self._cache[name].get(CNAME, None)
tm = time.time()
result = []
min_ttl = sys.maxint
for t in qtype:
death, ttl, rdata = self._cache[name][(t, qclass)]
if death < tm:
del self._cache[name][(t, qclass)]
continue
min_ttl = min(ttl, min_ttl)
if rdata:
result.extend(rdata)
if callback:
self._safely_call(callback, DNS_OK, cname, min_ttl,
tuple(result))
return
# Build a message and add our question.
m = DNSMessage()
m.questions.append((name, qtype[0], qclass))
# Make the function for handling our response.
def handle_response(status, data):
cname = None
# TTL is 30 by default, so answers with no records we want will be
# repeated, but not too often.
ttl = sys.maxint
if not data:
self._safely_call(callback, status, None, None, None)
return
rdata = {}
final_rdata = []
for (aname, atype, aclass, attl, ardata) in data.answers:
if atype == CNAME:
cname = ardata[0]
if atype in qtype and aclass == qclass:
ttl = min(attl, ttl)
if len(ardata) == 1:
rdata.setdefault(atype, []).append(ardata[0])
final_rdata.append(ardata[0])
else:
rdata.setdefault(atype, []).append(ardata)
final_rdata.append(ardata)
final_rdata = tuple(final_rdata)
ttl = min(30, ttl)
if allow_cache:
if not name in self._cache:
self._cache[name] = {}
if cname:
self._cache[name][CNAME] = cname
for t in qtype:
self._cache[name][(
t, qclass)] = time.time() + ttl, ttl, rdata.get(
t, [])
if data.rcode != DNS_OK:
status = data.rcode
self._safely_call(callback, status, cname, ttl, final_rdata)
# Send it, so we get an ID.
self.send_message(m, handle_response)
