diff --git a/proxy.py b/proxy.py index 2b966a5..0db8336 100644 --- a/proxy.py +++ b/proxy.py @@ -1,13 +1,14 @@ #!/usr/bin/env python3 -from zeroconf import Zeroconf, ServiceBrowser, ServiceStateChange, ServiceInfo, DNSQuestion, DNSOutgoing, RecordUpdateListener -from zeroconf.const import _TYPE_A, _CLASS_IN, _FLAGS_QR_QUERY +from zeroconf import Zeroconf, ServiceBrowser, ServiceStateChange, ServiceInfo, DNSQuestion, DNSOutgoing, RecordUpdateListener, IPVersion +from zeroconf.const import _TYPE_A, _TYPE_AAAA, _CLASS_IN, _FLAGS_QR_QUERY import sys import time from twisted.internet import reactor, defer, threads from twisted.names import client, dns, error, server from twisted.python import log import socket +import ipaddress domain = sys.argv[1] port = int(sys.argv[2]) @@ -16,7 +17,7 @@ timeout = 2 class DynamicResolver(object): def __init__(self): - self.zeroconf = Zeroconf(ip_version=4) + self.zeroconf = Zeroconf(ip_version=IPVersion.All) def _dynamicResponseRequired(self, query): if str(query.name).endswith(domain): @@ -84,22 +85,33 @@ class DynamicResolver(object): answers = [dns.RRHeader(name=localname[:-6] + domain, ttl=ttl, type=dns.SRV, payload=dns.Record_SRV( info.priority, info.weight, info.port, info.server[:-6] + domain ))] - additional = [dns.RRHeader(name=info.server[:-6] + domain, ttl=ttl, type=dns.A, payload=dns.Record_A( - socket.inet_ntop(socket.AF_INET, addr) - )) for addr in info.addresses] + target = info.server[:-6] + domain + additional = [dns.RRHeader(name=target, ttl=ttl, type=dns.A, payload=dns.Record_A( + str(addr) + )) for addr in info.ip_addresses_by_version(IPVersion.V4Only)] + additional += [dns.RRHeader(name=target, ttl=ttl, type=dns.AAAA, payload=dns.Record_AAAA( + str(addr) + )) for addr in info.ip_addresses_by_version(IPVersion.V6Only) if not addr.is_link_local] return answers, [], additional - def host(localname): + def host(localname, qtype): + if qtype == dns.AAAA: + mdns_type, addr_len, parse_addr = _TYPE_AAAA, 16, ipaddress.IPv6Address + else: + mdns_type, addr_len, parse_addr = _TYPE_A, 4, ipaddress.IPv4Address + class listener(RecordUpdateListener): def __init__(self): self.addrs = [] self.time = time.time() def update_record(self, zc, now, record): - if record.type == _TYPE_A and len(record.address) == 4: - self.addrs.append(socket.inet_ntop(socket.AF_INET, record.address)) - + if record.type == mdns_type and len(record.address) == addr_len: + addr = parse_addr(record.address) + if not addr.is_link_local: + self.addrs.append(str(addr)) + l = listener() - q = DNSQuestion(localname, _TYPE_A, _CLASS_IN) + q = DNSQuestion(localname, mdns_type, _CLASS_IN) self.zeroconf.add_listener(l, q) out = DNSOutgoing(_FLAGS_QR_QUERY) out.add_question(q) @@ -107,11 +119,12 @@ class DynamicResolver(object): while len(l.addrs) == 0 and time.time() - l.time < timeout: time.sleep(0.1) self.zeroconf.remove_listener(l) - - answers = [dns.RRHeader(name=query.name.name, ttl=ttl, type=dns.A, payload=dns.Record_A( + + record_cls = dns.Record_AAAA if qtype == dns.AAAA else dns.Record_A + answers = [dns.RRHeader(name=query.name.name, ttl=ttl, type=qtype, payload=record_cls( addr )) for addr in l.addrs] - + return answers, [], [] d = defer.Deferred() @@ -125,10 +138,10 @@ class DynamicResolver(object): elif query.type == dns.SRV: d = threads.deferToThread(srv, localname) return d - elif query.type == dns.A: - d = threads.deferToThread(host, localname) + elif query.type in (dns.A, dns.AAAA): + d = threads.deferToThread(host, localname, query.type) return d - elif query.type != dns.AAAA: + else: print("Unsupported request", query) d.callback(([], [], [])) return d