def create_request(
services: typing.List[str], qtype: QueryType = QueryType.PTR
) -> bytes:
"""Create a new DnsMessage requesting specified services."""
msg = DnsMessage(0x35FF)
msg.questions += [DnsQuestion(s, qtype, 0x8001) for s in services]
return msg.pack()
def create_service_queries(
services: typing.List[str], qtype: QueryType
) -> typing.List[bytes]:
"""Create service request messages."""
queries: typing.List[bytes] = []
for i in range(math.ceil(len(services) / SERVICES_PER_MSG)):
service_chunk = services[i * SERVICES_PER_MSG : i * SERVICES_PER_MSG + 4]
msg = DnsMessage(0x35FF)
msg.questions += [DnsQuestion(s, qtype, 0x8001) for s in service_chunk]
msg.questions += [DnsQuestion(SLEEP_PROXY_SERVICE, qtype, 0x8001)]
queries.append(msg.pack())
return queries
def __init__(self, services: typing.List[str], host: str, timeout: int):
"""Initialize a new UnicastDnsSdClientProtocol."""
self.message = create_request(services)
self.host = host
self.timeout = timeout
self.transport = None
self.semaphore: asyncio.Semaphore = asyncio.Semaphore(value=0)
self.result: DnsMessage = DnsMessage()
self._task: typing.Optional[asyncio.Future] = None
def datagram_received(self, data, addr) -> None:
"""DNS response packet received."""
log_binary(
_LOGGER,
f"Received DNS response from {addr}",
level=TRAFFIC_LEVEL,
Data=data,
)
query_resp = self.query_responses.setdefault(
addr[0], QueryResponse(count=0, deep_sleep=False, parser=ServiceParser())
)
# Suppress decode errors for now (but still log)
try:
decoded_msg = DnsMessage().unpack(data)
parser = ServiceParser()
services = parser.add_message(decoded_msg).parse()
except UnicodeDecodeError:
log_binary(_LOGGER, "Failed to decode message", Msg=data)
return
else:
if not services:
return
# Ignore responses from other services
for service in services:
if not (
service.type in self.services
or service.type in [DEVICE_INFO_SERVICE, SLEEP_PROXY_SERVICE]
):
return
is_sleep_proxy = all(service.port == 0 for service in services)
query_resp.count += 1
query_resp.deep_sleep |= is_sleep_proxy
query_resp.parser.add_message(decoded_msg)
if is_sleep_proxy:
self._unicasts[addr[0]] = create_service_queries(
[service.name + "." + service.type for service in services],
QueryType.ANY,
)
elif query_resp.count >= len(self.queries):
response = Response(
services=query_resp.parser.parse(),
deep_sleep=query_resp.deep_sleep,
model=_get_model(query_resp.parser.parse()),
)
if self.end_condition(response):
# Matches end condition: replace everything found so far and abort
self.query_responses = {addr[0]: self.query_responses[addr[0]]}
self.semaphore.release()
self.close()
def datagram_received(self, data: bytes, _) -> None:
"""DNS response packet received."""
log_binary(
_LOGGER,
"Received DNS response from " + self.host,
level=TRAFFIC_LEVEL,
Data=data,
)
self.result = DnsMessage().unpack(data)
self._finished()
if self.transport:
self.transport.close()
def datagram_received(self, data: bytes, _) -> None:
"""DNS response packet received."""
log_binary(
_LOGGER,
"Received DNS response from " + self.host,
level=TRAFFIC_LEVEL,
Data=data,
Index=self.received_responses + 1,
Total=len(self.queries),
)
self.parser.add_message(DnsMessage().unpack(data))
self.received_responses += 1
if self.received_responses == len(self.queries):
self._finished()
if self.transport:
self.transport.close()
def datagram_received(self, data, addr) -> None:
"""DNS response packet received."""
log_binary(
_LOGGER,
f"Received DNS response from {addr}",
level=TRAFFIC_LEVEL,
Data=data,
)
# Suppress decode errors for now (but still log)
try:
services = parse_services(DnsMessage().unpack(data))
except UnicodeDecodeError:
log_binary(_LOGGER, "Failed to decode message", Msg=data)
return
# Ignore responses from other services
for service in services:
if (
service.type not in self.services
and service.type != DEVICE_INFO_SERVICE
):
return
is_sleep_proxy = all(service.port == 0 for service in services)
if is_sleep_proxy:
self._unicasts[addr[0]] = create_request(
[service.name + "." + service.type for service in services],
qtype=QueryType.ANY,
)
else:
response = Response(
services=services,
deep_sleep=(addr[0] in self._unicasts),
model=_get_model(services),
)
if self.end_condition(response):
# Matches end condition: replace everything found so far and abort
self.responses = {IPv4Address(addr[0]): response}
self.semaphore.release()
self.close()
else:
self.responses[IPv4Address(addr[0])] = response