133 lines
4.4 KiB
Python
133 lines
4.4 KiB
Python
#!/usr/bin/env python3
|
|
|
|
from zeroconf import Zeroconf, ServiceBrowser, ServiceStateChange
|
|
import sys
|
|
import time
|
|
from twisted.internet import reactor, defer, threads
|
|
from twisted.names import client, dns, error, server
|
|
from twisted.python import log
|
|
|
|
domain = sys.argv[1]
|
|
port = int(sys.argv[2])
|
|
ttl = 120
|
|
timeout = 2
|
|
|
|
class DynamicResolver(object):
|
|
def _dynamicResponseRequired(self, query):
|
|
if str(query.name).endswith(domain):
|
|
return True
|
|
|
|
return False
|
|
|
|
def _doDynamicResponse(self, query):
|
|
if query.type == dns.SOA:
|
|
return defer.succeed(([], [], []))
|
|
|
|
localname = str(query.name)[:-len(domain)] + "local."
|
|
|
|
def browse(localname):
|
|
services = []
|
|
def handler(zeroconf, service_type, name, state_change):
|
|
if state_change is ServiceStateChange.Added:
|
|
services.append(name)
|
|
|
|
zeroconf = Zeroconf(ip_version=4) # TODO: move out?
|
|
sb = ServiceBrowser(zeroconf, localname, [handler])
|
|
print("waiting for", localname) # TODO: remove
|
|
time.sleep(timeout)
|
|
sb.cancel()
|
|
print("done waiting for", localname) # TODO: remove
|
|
|
|
answers = []
|
|
for service in services:
|
|
answers.append(dns.RRHeader(name=query.name.name, ttl=ttl, type=dns.PTR, payload=dns.Record_PTR(
|
|
name=service[:-6] + domain
|
|
)))
|
|
return answers
|
|
|
|
def txt(localname):
|
|
zeroconf = Zeroconf(ip_version=4) # TODO: move out?
|
|
print("waiting for txt", localname) # TODO: remove
|
|
info = zeroconf.get_service_info(localname, localname)
|
|
print("done waiting for txt", localname) # TODO: remove
|
|
|
|
if info.text == b"\0":
|
|
answers = []
|
|
else:
|
|
answers = [dns.RRHeader(name=query.name.name, ttl=ttl, type=dns.TXT, payload=dns.Record_TXT(
|
|
info.text # TODO: something wrong here?
|
|
))]
|
|
print(answers, type(info.text), info.text)
|
|
return answers
|
|
|
|
def srv(localname):
|
|
zeroconf = Zeroconf(ip_version=4) # TODO: move out?
|
|
print("waiting for srv", localname)
|
|
info = zeroconf.get_service_info(localname, localname)
|
|
print("done waiting for srv", localname)
|
|
|
|
answers = [dns.RRHeader(name=query.name.name, ttl=ttl, type=dns.SRV, payload=dns.Record_SRV(
|
|
info.priority, info.weight, info.port, info.server[:-6] + domain
|
|
))]
|
|
|
|
print(answers)
|
|
|
|
return answers
|
|
|
|
def host(localname):
|
|
answers = [dns.RRHeader(name=query.name.name, ttl=ttl, type=dns.A, payload=dns.Record_A(
|
|
"192.0.2.1" # TODO
|
|
))]
|
|
|
|
print(answers)
|
|
|
|
return answers
|
|
|
|
d = defer.Deferred()
|
|
|
|
if query.type == dns.PTR:
|
|
d = threads.deferToThread(browse, localname)
|
|
d.addCallback(lambda a: (a, [], []))
|
|
return d
|
|
elif query.type == dns.TXT:
|
|
d = threads.deferToThread(txt, localname)
|
|
d.addCallback(lambda a: (a, [], []))
|
|
return d
|
|
elif query.type == dns.SRV:
|
|
d = threads.deferToThread(srv, localname)
|
|
d.addCallback(lambda a: (a, [], []))
|
|
return d
|
|
elif query.type == dns.A:
|
|
d = threads.deferToThread(host, localname)
|
|
d.addCallback(lambda a: (a, [], []))
|
|
return d
|
|
else:
|
|
answers = []
|
|
d.callback((answers, [], []))
|
|
if query.type != dns.AAAA:
|
|
print("Unsupported request", query)
|
|
return d
|
|
|
|
def query(self, query, timeout=None):
|
|
if self._dynamicResponseRequired(query):
|
|
return self._doDynamicResponse(query)
|
|
else:
|
|
return defer.fail(error.DomainError())
|
|
|
|
def main():
|
|
factory = server.DNSServerFactory(
|
|
clients=[DynamicResolver()],
|
|
verbose=0
|
|
)
|
|
|
|
protocol = dns.DNSDatagramProtocol(controller=factory)
|
|
|
|
reactor.listenUDP(port, protocol)
|
|
reactor.listenTCP(port, factory)
|
|
|
|
log.startLogging(sys.stdout)
|
|
reactor.run()
|
|
|
|
if __name__ == '__main__':
|
|
raise SystemExit(main())
|