438 lines
17 KiB
Python
438 lines
17 KiB
Python
|
""" Multicast DNS Service Discovery for Python, v0.14-wmcbrine
|
||
|
|
Copyright 2003 Paul Scott-Murphy, 2014 William McBrine
|
||
|
|
|
||
|
|
This module provides a framework for the use of DNS Service Discovery
|
||
|
|
using IP multicast.
|
||
|
|
|
||
|
|
This library is free software; you can redistribute it and/or
|
||
|
|
modify it under the terms of the GNU Lesser General Public
|
||
|
|
License as published by the Free Software Foundation; either
|
||
|
|
version 2.1 of the License, or (at your option) any later version.
|
||
|
|
|
||
|
|
This library is distributed in the hope that it will be useful,
|
||
|
|
but WITHOUT ANY WARRANTY; without even the implied warranty of
|
||
|
|
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
|
||
|
|
Lesser General Public License for more details.
|
||
|
|
|
||
|
|
You should have received a copy of the GNU Lesser General Public
|
||
|
|
License along with this library; if not, write to the Free Software
|
||
|
|
Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA 02110-1301
|
||
|
|
USA
|
||
|
|
"""
|
||
|
|
|
||
|
|
from typing import TYPE_CHECKING, List, Optional, Set, Tuple, Union, cast
|
||
|
|
|
||
|
|
from .._cache import DNSCache, _UniqueRecordsType
|
||
|
|
from .._dns import DNSAddress, DNSPointer, DNSQuestion, DNSRecord, DNSRRSet
|
||
|
|
from .._protocol.incoming import DNSIncoming
|
||
|
|
from .._services.info import ServiceInfo
|
||
|
|
from .._transport import _WrappedTransport
|
||
|
|
from .._utils.net import IPVersion
|
||
|
|
from ..const import (
|
||
|
|
_ADDRESS_RECORD_TYPES,
|
||
|
|
_CLASS_IN,
|
||
|
|
_DNS_OTHER_TTL,
|
||
|
|
_MDNS_PORT,
|
||
|
|
_ONE_SECOND,
|
||
|
|
_SERVICE_TYPE_ENUMERATION_NAME,
|
||
|
|
_TYPE_A,
|
||
|
|
_TYPE_AAAA,
|
||
|
|
_TYPE_ANY,
|
||
|
|
_TYPE_NSEC,
|
||
|
|
_TYPE_PTR,
|
||
|
|
_TYPE_SRV,
|
||
|
|
_TYPE_TXT,
|
||
|
|
)
|
||
|
|
from .answers import (
|
||
|
|
QuestionAnswers,
|
||
|
|
_AnswerWithAdditionalsType,
|
||
|
|
construct_outgoing_multicast_answers,
|
||
|
|
construct_outgoing_unicast_answers,
|
||
|
|
)
|
||
|
|
|
||
|
|
_RESPOND_IMMEDIATE_TYPES = {_TYPE_NSEC, _TYPE_SRV, *_ADDRESS_RECORD_TYPES}
|
||
|
|
|
||
|
|
_EMPTY_SERVICES_LIST: List[ServiceInfo] = []
|
||
|
|
_EMPTY_TYPES_LIST: List[str] = []
|
||
|
|
|
||
|
|
_IPVersion_ALL = IPVersion.All
|
||
|
|
|
||
|
|
_int = int
|
||
|
|
_str = str
|
||
|
|
|
||
|
|
_ANSWER_STRATEGY_SERVICE_TYPE_ENUMERATION = 0
|
||
|
|
_ANSWER_STRATEGY_POINTER = 1
|
||
|
|
_ANSWER_STRATEGY_ADDRESS = 2
|
||
|
|
_ANSWER_STRATEGY_SERVICE = 3
|
||
|
|
_ANSWER_STRATEGY_TEXT = 4
|
||
|
|
|
||
|
|
if TYPE_CHECKING:
|
||
|
|
from .._core import Zeroconf
|
||
|
|
|
||
|
|
|
||
|
|
class _AnswerStrategy:
|
||
|
|
|
||
|
|
__slots__ = ("question", "strategy_type", "types", "services")
|
||
|
|
|
||
|
|
def __init__(
|
||
|
|
self,
|
||
|
|
question: DNSQuestion,
|
||
|
|
strategy_type: _int,
|
||
|
|
types: List[str],
|
||
|
|
services: List[ServiceInfo],
|
||
|
|
) -> None:
|
||
|
|
"""Create an answer strategy."""
|
||
|
|
self.question = question
|
||
|
|
self.strategy_type = strategy_type
|
||
|
|
self.types = types
|
||
|
|
self.services = services
|
||
|
|
|
||
|
|
|
||
|
|
class _QueryResponse:
|
||
|
|
"""A pair for unicast and multicast DNSOutgoing responses."""
|
||
|
|
|
||
|
|
__slots__ = (
|
||
|
|
"_is_probe",
|
||
|
|
"_questions",
|
||
|
|
"_now",
|
||
|
|
"_cache",
|
||
|
|
"_additionals",
|
||
|
|
"_ucast",
|
||
|
|
"_mcast_now",
|
||
|
|
"_mcast_aggregate",
|
||
|
|
"_mcast_aggregate_last_second",
|
||
|
|
)
|
||
|
|
|
||
|
|
def __init__(self, cache: DNSCache, questions: List[DNSQuestion], is_probe: bool, now: float) -> None:
|
||
|
|
"""Build a query response."""
|
||
|
|
self._is_probe = is_probe
|
||
|
|
self._questions = questions
|
||
|
|
self._now = now
|
||
|
|
self._cache = cache
|
||
|
|
self._additionals: _AnswerWithAdditionalsType = {}
|
||
|
|
self._ucast: Set[DNSRecord] = set()
|
||
|
|
self._mcast_now: Set[DNSRecord] = set()
|
||
|
|
self._mcast_aggregate: Set[DNSRecord] = set()
|
||
|
|
self._mcast_aggregate_last_second: Set[DNSRecord] = set()
|
||
|
|
|
||
|
|
def add_qu_question_response(self, answers: _AnswerWithAdditionalsType) -> None:
|
||
|
|
"""Generate a response to a multicast QU query."""
|
||
|
|
for record, additionals in answers.items():
|
||
|
|
self._additionals[record] = additionals
|
||
|
|
if self._is_probe:
|
||
|
|
self._ucast.add(record)
|
||
|
|
if not self._has_mcast_within_one_quarter_ttl(record):
|
||
|
|
self._mcast_now.add(record)
|
||
|
|
elif not self._is_probe:
|
||
|
|
self._ucast.add(record)
|
||
|
|
|
||
|
|
def add_ucast_question_response(self, answers: _AnswerWithAdditionalsType) -> None:
|
||
|
|
"""Generate a response to a unicast query."""
|
||
|
|
self._additionals.update(answers)
|
||
|
|
self._ucast.update(answers)
|
||
|
|
|
||
|
|
def add_mcast_question_response(self, answers: _AnswerWithAdditionalsType) -> None:
|
||
|
|
"""Generate a response to a multicast query."""
|
||
|
|
self._additionals.update(answers)
|
||
|
|
for answer in answers:
|
||
|
|
if self._is_probe:
|
||
|
|
self._mcast_now.add(answer)
|
||
|
|
continue
|
||
|
|
|
||
|
|
if self._has_mcast_record_in_last_second(answer):
|
||
|
|
self._mcast_aggregate_last_second.add(answer)
|
||
|
|
continue
|
||
|
|
|
||
|
|
if len(self._questions) == 1:
|
||
|
|
question = self._questions[0]
|
||
|
|
if question.type in _RESPOND_IMMEDIATE_TYPES:
|
||
|
|
self._mcast_now.add(answer)
|
||
|
|
continue
|
||
|
|
|
||
|
|
self._mcast_aggregate.add(answer)
|
||
|
|
|
||
|
|
def answers(
|
||
|
|
self,
|
||
|
|
) -> QuestionAnswers:
|
||
|
|
"""Return answer sets that will be queued."""
|
||
|
|
ucast = {r: self._additionals[r] for r in self._ucast}
|
||
|
|
mcast_now = {r: self._additionals[r] for r in self._mcast_now}
|
||
|
|
mcast_aggregate = {r: self._additionals[r] for r in self._mcast_aggregate}
|
||
|
|
mcast_aggregate_last_second = {r: self._additionals[r] for r in self._mcast_aggregate_last_second}
|
||
|
|
return QuestionAnswers(ucast, mcast_now, mcast_aggregate, mcast_aggregate_last_second)
|
||
|
|
|
||
|
|
def _has_mcast_within_one_quarter_ttl(self, record: DNSRecord) -> bool:
|
||
|
|
"""Check to see if a record has been mcasted recently.
|
||
|
|
|
||
|
|
https://datatracker.ietf.org/doc/html/rfc6762#section-5.4
|
||
|
|
When receiving a question with the unicast-response bit set, a
|
||
|
|
responder SHOULD usually respond with a unicast packet directed back
|
||
|
|
to the querier. However, if the responder has not multicast that
|
||
|
|
record recently (within one quarter of its TTL), then the responder
|
||
|
|
SHOULD instead multicast the response so as to keep all the peer
|
||
|
|
caches up to date
|
||
|
|
"""
|
||
|
|
if TYPE_CHECKING:
|
||
|
|
record = cast(_UniqueRecordsType, record)
|
||
|
|
maybe_entry = self._cache.async_get_unique(record)
|
||
|
|
return bool(maybe_entry is not None and maybe_entry.is_recent(self._now))
|
||
|
|
|
||
|
|
def _has_mcast_record_in_last_second(self, record: DNSRecord) -> bool:
|
||
|
|
"""Check if an answer was seen in the last second.
|
||
|
|
Protect the network against excessive packet flooding
|
||
|
|
https://datatracker.ietf.org/doc/html/rfc6762#section-14
|
||
|
|
"""
|
||
|
|
if TYPE_CHECKING:
|
||
|
|
record = cast(_UniqueRecordsType, record)
|
||
|
|
maybe_entry = self._cache.async_get_unique(record)
|
||
|
|
return bool(maybe_entry is not None and self._now - maybe_entry.created < _ONE_SECOND)
|
||
|
|
|
||
|
|
|
||
|
|
class QueryHandler:
|
||
|
|
"""Query the ServiceRegistry."""
|
||
|
|
|
||
|
|
__slots__ = ("zc", "registry", "cache", "question_history", "out_queue", "out_delay_queue")
|
||
|
|
|
||
|
|
def __init__(self, zc: 'Zeroconf') -> None:
|
||
|
|
"""Init the query handler."""
|
||
|
|
self.zc = zc
|
||
|
|
self.registry = zc.registry
|
||
|
|
self.cache = zc.cache
|
||
|
|
self.question_history = zc.question_history
|
||
|
|
self.out_queue = zc.out_queue
|
||
|
|
self.out_delay_queue = zc.out_delay_queue
|
||
|
|
|
||
|
|
def _add_service_type_enumeration_query_answers(
|
||
|
|
self, types: List[str], answer_set: _AnswerWithAdditionalsType, known_answers: DNSRRSet
|
||
|
|
) -> None:
|
||
|
|
"""Provide an answer to a service type enumeration query.
|
||
|
|
|
||
|
|
https://datatracker.ietf.org/doc/html/rfc6763#section-9
|
||
|
|
"""
|
||
|
|
for stype in types:
|
||
|
|
dns_pointer = DNSPointer(
|
||
|
|
_SERVICE_TYPE_ENUMERATION_NAME, _TYPE_PTR, _CLASS_IN, _DNS_OTHER_TTL, stype, 0.0
|
||
|
|
)
|
||
|
|
if not known_answers.suppresses(dns_pointer):
|
||
|
|
answer_set[dns_pointer] = set()
|
||
|
|
|
||
|
|
def _add_pointer_answers(
|
||
|
|
self, services: List[ServiceInfo], answer_set: _AnswerWithAdditionalsType, known_answers: DNSRRSet
|
||
|
|
) -> None:
|
||
|
|
"""Answer PTR/ANY question."""
|
||
|
|
for service in services:
|
||
|
|
# Add recommended additional answers according to
|
||
|
|
# https://tools.ietf.org/html/rfc6763#section-12.1.
|
||
|
|
dns_pointer = service._dns_pointer(None)
|
||
|
|
if known_answers.suppresses(dns_pointer):
|
||
|
|
continue
|
||
|
|
answer_set[dns_pointer] = {
|
||
|
|
service._dns_service(None),
|
||
|
|
service._dns_text(None),
|
||
|
|
*service._get_address_and_nsec_records(None),
|
||
|
|
}
|
||
|
|
|
||
|
|
def _add_address_answers(
|
||
|
|
self,
|
||
|
|
services: List[ServiceInfo],
|
||
|
|
answer_set: _AnswerWithAdditionalsType,
|
||
|
|
known_answers: DNSRRSet,
|
||
|
|
type_: _int,
|
||
|
|
) -> None:
|
||
|
|
"""Answer A/AAAA/ANY question."""
|
||
|
|
for service in services:
|
||
|
|
answers: List[DNSAddress] = []
|
||
|
|
additionals: Set[DNSRecord] = set()
|
||
|
|
seen_types: Set[int] = set()
|
||
|
|
for dns_address in service._dns_addresses(None, _IPVersion_ALL):
|
||
|
|
seen_types.add(dns_address.type)
|
||
|
|
if dns_address.type != type_:
|
||
|
|
additionals.add(dns_address)
|
||
|
|
elif not known_answers.suppresses(dns_address):
|
||
|
|
answers.append(dns_address)
|
||
|
|
missing_types: Set[int] = _ADDRESS_RECORD_TYPES - seen_types
|
||
|
|
if answers:
|
||
|
|
if missing_types:
|
||
|
|
assert service.server is not None, "Service server must be set for NSEC record."
|
||
|
|
additionals.add(service._dns_nsec(list(missing_types), None))
|
||
|
|
for answer in answers:
|
||
|
|
answer_set[answer] = additionals
|
||
|
|
elif type_ in missing_types:
|
||
|
|
assert service.server is not None, "Service server must be set for NSEC record."
|
||
|
|
answer_set[service._dns_nsec(list(missing_types), None)] = set()
|
||
|
|
|
||
|
|
def _answer_question(
|
||
|
|
self,
|
||
|
|
question: DNSQuestion,
|
||
|
|
strategy_type: _int,
|
||
|
|
types: List[str],
|
||
|
|
services: List[ServiceInfo],
|
||
|
|
known_answers: DNSRRSet,
|
||
|
|
) -> _AnswerWithAdditionalsType:
|
||
|
|
"""Answer a question."""
|
||
|
|
answer_set: _AnswerWithAdditionalsType = {}
|
||
|
|
|
||
|
|
if strategy_type == _ANSWER_STRATEGY_SERVICE_TYPE_ENUMERATION:
|
||
|
|
self._add_service_type_enumeration_query_answers(types, answer_set, known_answers)
|
||
|
|
elif strategy_type == _ANSWER_STRATEGY_POINTER:
|
||
|
|
self._add_pointer_answers(services, answer_set, known_answers)
|
||
|
|
elif strategy_type == _ANSWER_STRATEGY_ADDRESS:
|
||
|
|
self._add_address_answers(services, answer_set, known_answers, question.type)
|
||
|
|
elif strategy_type == _ANSWER_STRATEGY_SERVICE:
|
||
|
|
# Add recommended additional answers according to
|
||
|
|
# https://tools.ietf.org/html/rfc6763#section-12.2.
|
||
|
|
service = services[0]
|
||
|
|
dns_service = service._dns_service(None)
|
||
|
|
if not known_answers.suppresses(dns_service):
|
||
|
|
answer_set[dns_service] = service._get_address_and_nsec_records(None)
|
||
|
|
elif strategy_type == _ANSWER_STRATEGY_TEXT: # pragma: no branch
|
||
|
|
service = services[0]
|
||
|
|
dns_text = service._dns_text(None)
|
||
|
|
if not known_answers.suppresses(dns_text):
|
||
|
|
answer_set[dns_text] = set()
|
||
|
|
|
||
|
|
return answer_set
|
||
|
|
|
||
|
|
def async_response( # pylint: disable=unused-argument
|
||
|
|
self, msgs: List[DNSIncoming], ucast_source: bool
|
||
|
|
) -> Optional[QuestionAnswers]:
|
||
|
|
"""Deal with incoming query packets. Provides a response if possible.
|
||
|
|
|
||
|
|
This function must be run in the event loop as it is not
|
||
|
|
threadsafe.
|
||
|
|
"""
|
||
|
|
strategies: List[_AnswerStrategy] = []
|
||
|
|
for msg in msgs:
|
||
|
|
for question in msg._questions:
|
||
|
|
strategies.extend(self._get_answer_strategies(question))
|
||
|
|
|
||
|
|
if not strategies:
|
||
|
|
# We have no way to answer the question because we have
|
||
|
|
# nothing in the ServiceRegistry that matches or we do not
|
||
|
|
# understand the question.
|
||
|
|
return None
|
||
|
|
|
||
|
|
is_probe = False
|
||
|
|
msg = msgs[0]
|
||
|
|
questions = msg._questions
|
||
|
|
# Only decode known answers if we are not a probe and we have
|
||
|
|
# at least one answer strategy
|
||
|
|
answers: List[DNSRecord] = []
|
||
|
|
for msg in msgs:
|
||
|
|
if msg.is_probe():
|
||
|
|
is_probe = True
|
||
|
|
else:
|
||
|
|
answers.extend(msg.answers())
|
||
|
|
|
||
|
|
query_res = _QueryResponse(self.cache, questions, is_probe, msg.now)
|
||
|
|
known_answers = DNSRRSet(answers)
|
||
|
|
known_answers_set: Optional[Set[DNSRecord]] = None
|
||
|
|
now = msg.now
|
||
|
|
for strategy in strategies:
|
||
|
|
question = strategy.question
|
||
|
|
is_unicast = question.unique # unique and unicast are the same flag
|
||
|
|
if not is_unicast:
|
||
|
|
if known_answers_set is None: # pragma: no branch
|
||
|
|
known_answers_set = known_answers.lookup_set()
|
||
|
|
self.question_history.add_question_at_time(question, now, known_answers_set)
|
||
|
|
answer_set = self._answer_question(
|
||
|
|
question, strategy.strategy_type, strategy.types, strategy.services, known_answers
|
||
|
|
)
|
||
|
|
if not ucast_source and is_unicast:
|
||
|
|
query_res.add_qu_question_response(answer_set)
|
||
|
|
continue
|
||
|
|
if ucast_source:
|
||
|
|
query_res.add_ucast_question_response(answer_set)
|
||
|
|
# We always multicast as well even if its a unicast
|
||
|
|
# source as long as we haven't done it recently (75% of ttl)
|
||
|
|
query_res.add_mcast_question_response(answer_set)
|
||
|
|
|
||
|
|
return query_res.answers()
|
||
|
|
|
||
|
|
def _get_answer_strategies(
|
||
|
|
self,
|
||
|
|
question: DNSQuestion,
|
||
|
|
) -> List[_AnswerStrategy]:
|
||
|
|
"""Collect strategies to answer a question."""
|
||
|
|
name = question.name
|
||
|
|
question_lower_name = name.lower()
|
||
|
|
type_ = question.type
|
||
|
|
strategies: List[_AnswerStrategy] = []
|
||
|
|
|
||
|
|
if type_ == _TYPE_PTR and question_lower_name == _SERVICE_TYPE_ENUMERATION_NAME:
|
||
|
|
types = self.registry.async_get_types()
|
||
|
|
if types:
|
||
|
|
strategies.append(
|
||
|
|
_AnswerStrategy(
|
||
|
|
question, _ANSWER_STRATEGY_SERVICE_TYPE_ENUMERATION, types, _EMPTY_SERVICES_LIST
|
||
|
|
)
|
||
|
|
)
|
||
|
|
return strategies
|
||
|
|
|
||
|
|
if type_ in (_TYPE_PTR, _TYPE_ANY):
|
||
|
|
services = self.registry.async_get_infos_type(question_lower_name)
|
||
|
|
if services:
|
||
|
|
strategies.append(
|
||
|
|
_AnswerStrategy(question, _ANSWER_STRATEGY_POINTER, _EMPTY_TYPES_LIST, services)
|
||
|
|
)
|
||
|
|
|
||
|
|
if type_ in (_TYPE_A, _TYPE_AAAA, _TYPE_ANY):
|
||
|
|
services = self.registry.async_get_infos_server(question_lower_name)
|
||
|
|
if services:
|
||
|
|
strategies.append(
|
||
|
|
_AnswerStrategy(question, _ANSWER_STRATEGY_ADDRESS, _EMPTY_TYPES_LIST, services)
|
||
|
|
)
|
||
|
|
|
||
|
|
if type_ in (_TYPE_SRV, _TYPE_TXT, _TYPE_ANY):
|
||
|
|
service = self.registry.async_get_info_name(question_lower_name)
|
||
|
|
if service is not None:
|
||
|
|
if type_ in (_TYPE_SRV, _TYPE_ANY):
|
||
|
|
strategies.append(
|
||
|
|
_AnswerStrategy(question, _ANSWER_STRATEGY_SERVICE, _EMPTY_TYPES_LIST, [service])
|
||
|
|
)
|
||
|
|
if type_ in (_TYPE_TXT, _TYPE_ANY):
|
||
|
|
strategies.append(
|
||
|
|
_AnswerStrategy(question, _ANSWER_STRATEGY_TEXT, _EMPTY_TYPES_LIST, [service])
|
||
|
|
)
|
||
|
|
|
||
|
|
return strategies
|
||
|
|
|
||
|
|
def handle_assembled_query(
|
||
|
|
self,
|
||
|
|
packets: List[DNSIncoming],
|
||
|
|
addr: _str,
|
||
|
|
port: _int,
|
||
|
|
transport: _WrappedTransport,
|
||
|
|
v6_flow_scope: Union[Tuple[()], Tuple[int, int]],
|
||
|
|
) -> None:
|
||
|
|
"""Respond to a (re)assembled query.
|
||
|
|
|
||
|
|
If the protocol recieved packets with the TC bit set, it will
|
||
|
|
wait a bit for the rest of the packets and only call
|
||
|
|
handle_assembled_query once it has a complete set of packets
|
||
|
|
or the timer expires. If the TC bit is not set, a single
|
||
|
|
packet will be in packets.
|
||
|
|
"""
|
||
|
|
first_packet = packets[0]
|
||
|
|
ucast_source = port != _MDNS_PORT
|
||
|
|
question_answers = self.async_response(packets, ucast_source)
|
||
|
|
if question_answers is None:
|
||
|
|
return
|
||
|
|
if question_answers.ucast:
|
||
|
|
questions = first_packet._questions
|
||
|
|
id_ = first_packet.id
|
||
|
|
out = construct_outgoing_unicast_answers(question_answers.ucast, ucast_source, questions, id_)
|
||
|
|
# When sending unicast, only send back the reply
|
||
|
|
# via the same socket that it was recieved from
|
||
|
|
# as we know its reachable from that socket
|
||
|
|
self.zc.async_send(out, addr, port, v6_flow_scope, transport)
|
||
|
|
if question_answers.mcast_now:
|
||
|
|
self.zc.async_send(construct_outgoing_multicast_answers(question_answers.mcast_now))
|
||
|
|
if question_answers.mcast_aggregate:
|
||
|
|
self.out_queue.async_add(first_packet.now, question_answers.mcast_aggregate)
|
||
|
|
if question_answers.mcast_aggregate_last_second:
|
||
|
|
# https://datatracker.ietf.org/doc/html/rfc6762#section-14
|
||
|
|
# If we broadcast it in the last second, we have to delay
|
||
|
|
# at least a second before we send it again
|
||
|
|
self.out_delay_queue.async_add(first_packet.now, question_answers.mcast_aggregate_last_second)
|