""" 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

"""

import asyncio

import queue

import random

import threading

import warnings

from collections import OrderedDict

from typing import Callable, Dict, Iterable, List, Optional, Set, TYPE_CHECKING, Tuple, Union, cast

from .._dns import DNSAddress, DNSPointer, DNSQuestion, DNSQuestionType, DNSRecord

from .._logger import log

from .._protocol.outgoing import DNSOutgoing

from .._services import (

ServiceListener,

ServiceStateChange,

Signal,

SignalRegistrationInterface,

)

from .._updates import RecordUpdate, RecordUpdateListener

from .._utils.name import possible_types, service_type_name

from .._utils.time import current_time_millis, millis_to_seconds

from ..const import (

_BROWSER_BACKOFF_LIMIT,

_BROWSER_TIME,

_CLASS_IN,

_DNS_PACKET_HEADER_LEN,

_EXPIRE_REFRESH_TIME_PERCENT,

_FLAGS_QR_QUERY,

_MAX_MSG_TYPICAL,

_MDNS_ADDR,

_MDNS_ADDR6,

_MDNS_PORT,

_TYPE_PTR,

)

# https://datatracker.ietf.org/doc/html/rfc6762#section-5.2

_FIRST_QUERY_DELAY_RANDOM_INTERVAL = (20, 120) # ms

_ON_CHANGE_DISPATCH = {

ServiceStateChange.Added: "add_service",

ServiceStateChange.Removed: "remove_service",

ServiceStateChange.Updated: "update_service",

}

if TYPE_CHECKING:

from .._core import Zeroconf

_QuestionWithKnownAnswers = Dict[DNSQuestion, Set[DNSPointer]]

class _DNSPointerOutgoingBucket:

"""A DNSOutgoing bucket."""

def __init__(self, now: float, multicast: bool) -> None:

"""Create a bucke to wrap a DNSOutgoing."""

self.now = now

self.out = DNSOutgoing(_FLAGS_QR_QUERY, multicast=multicast)

self.bytes = 0

def add(self, max_compressed_size: int, question: DNSQuestion, answers: Set[DNSPointer]) -> None:

"""Add a new set of questions and known answers to the outgoing."""

self.out.add_question(question)

for answer in answers:

self.out.add_answer_at_time(answer, self.now)

self.bytes += max_compressed_size

def _group_ptr_queries_with_known_answers(

now: float, multicast: bool, question_with_known_answers: _QuestionWithKnownAnswers

) -> List[DNSOutgoing]:

"""Aggregate queries so that as many known answers as possible fit in the same packet

without having known answers spill over into the next packet unless the

question and known answers are always going to exceed the packet size.

Some responders do not implement multi-packet known answer suppression

so we try to keep all the known answers in the same packet as the

questions.

"""

# This is the maximum size the query + known answers can be with name compression.

# The actual size of the query + known answers may be a bit smaller since other

# parts may be shared when the final DNSOutgoing packets are constructed. The

# goal of this algorithm is to quickly bucket the query + known answers without

# the overhead of actually constructing the packets.

query_by_size: Dict[DNSQuestion, int] = {

question: (question.max_size + sum([answer.max_size_compressed for answer in known_answers]))

for question, known_answers in question_with_known_answers.items()

}

max_bucket_size = _MAX_MSG_TYPICAL - _DNS_PACKET_HEADER_LEN

query_buckets: List[_DNSPointerOutgoingBucket] = []

for question in sorted(

query_by_size,

key=query_by_size.get, # type: ignore

reverse=True,

):

max_compressed_size = query_by_size[question]

answers = question_with_known_answers[question]

for query_bucket in query_buckets:

if query_bucket.bytes + max_compressed_size <= max_bucket_size:

query_bucket.add(max_compressed_size, question, answers)

break

else:

# If a single question and known answers won't fit in a packet

# we will end up generating multiple packets, but there will never

# be multiple questions

query_bucket = _DNSPointerOutgoingBucket(now, multicast)

query_bucket.add(max_compressed_size, question, answers)

query_buckets.append(query_bucket)

return [query_bucket.out for query_bucket in query_buckets]

def generate_service_query(

zc: 'Zeroconf',

now: float,

types_: List[str],

multicast: bool = True,

question_type: Optional[DNSQuestionType] = None,

) -> List[DNSOutgoing]:

"""Generate a service query for sending with zeroconf.send."""

questions_with_known_answers: _QuestionWithKnownAnswers = {}

qu_question = not multicast if question_type is None else question_type == DNSQuestionType.QU

for type_ in types_:

question = DNSQuestion(type_, _TYPE_PTR, _CLASS_IN)

question.unicast = qu_question

known_answers = {

cast(DNSPointer, record)

for record in zc.cache.get_all_by_details(type_, _TYPE_PTR, _CLASS_IN)

if not record.is_stale(now)

}

if not qu_question and zc.question_history.suppresses(

question, now, cast(Set[DNSRecord], known_answers)

):

log.debug("Asking %s was suppressed by the question history", question)

continue

questions_with_known_answers[question] = known_answers

if not qu_question:

zc.question_history.add_question_at_time(question, now, cast(Set[DNSRecord], known_answers))

return _group_ptr_queries_with_known_answers(now, multicast, questions_with_known_answers)

def _service_state_changed_from_listener(listener: ServiceListener) -> Callable[..., None]:

"""Generate a service_state_changed handlers from a listener."""

assert listener is not None

if not hasattr(listener, 'update_service'):

warnings.warn(

"%r has no update_service method. Provide one (it can be empty if you "

"don't care about the updates), it'll become mandatory." % (listener,),

FutureWarning,

)

def on_change(

zeroconf: 'Zeroconf', service_type: str, name: str, state_change: ServiceStateChange

) -> None:

getattr(listener, _ON_CHANGE_DISPATCH[state_change])(zeroconf, service_type, name)

return on_change

class QueryScheduler:

"""Schedule outgoing PTR queries for Continuous Multicast DNS Querying

https://datatracker.ietf.org/doc/html/rfc6762#section-5.2

"""

def __init__(

self,

types: Set[str],

delay: int,

first_random_delay_interval: Tuple[int, int],

):

self._schedule_changed_event: Optional[asyncio.Event] = None

self._types = types

self._next_time: Dict[str, float] = {}

self._first_random_delay_interval = first_random_delay_interval

self._delay: Dict[str, float] = {check_type_: delay for check_type_ in self._types}

def start(self, now: float) -> None:

"""Start the scheduler."""

self._schedule_changed_event = asyncio.Event()

self._generate_first_next_time(now)

def _generate_first_next_time(self, now: float) -> None:

"""Generate the initial next query times.

https://datatracker.ietf.org/doc/html/rfc6762#section-5.2

To avoid accidental synchronization when, for some reason, multiple

clients begin querying at exactly the same moment (e.g., because of

some common external trigger event), a Multicast DNS querier SHOULD

also delay the first query of the series by a randomly chosen amount

in the range 20-120 ms.

"""

delay = millis_to_seconds(random.randint(*self._first_random_delay_interval))

next_time = now + delay

self._next_time = {check_type_: next_time for check_type_ in self._types}

def millis_to_wait(self, now: float) -> float:

"""Returns the number of milliseconds to wait for the next event."""

# Wait for the type has the smallest next time

next_time = min(self._next_time.values())

return 0 if next_time <= now else next_time - now

def reschedule_type(self, type_: str, next_time: float) -> bool:

"""Reschedule the query for a type to happen sooner."""

if next_time >= self._next_time[type_]:

return False

self._next_time[type_] = next_time

return True

def process_ready_types(self, now: float) -> List[str]:

"""Generate a list of ready types that is due and schedule the next time."""

if self.millis_to_wait(now):

return []

ready_types: List[str] = []

for type_, due in self._next_time.items():

if due > now:

continue

ready_types.append(type_)

self._next_time[type_] = now + self._delay[type_]

self._delay[type_] = min(_BROWSER_BACKOFF_LIMIT * 1000, self._delay[type_] * 2)

return ready_types

class _ServiceBrowserBase(RecordUpdateListener):

"""Base class for ServiceBrowser."""

def __init__(

self,

zc: 'Zeroconf',

type_: Union[str, list],

handlers: Optional[Union[ServiceListener, List[Callable[..., None]]]] = None,

listener: Optional[ServiceListener] = None,

addr: Optional[str] = None,

port: int = _MDNS_PORT,

delay: int = _BROWSER_TIME,

question_type: Optional[DNSQuestionType] = None,

) -> None:

"""Used to browse for a service for specific type(s).

Constructor parameters are as follows:

* `zc`: A Zeroconf instance

* `type_`: fully qualified service type name

* `handler`: ServiceListener or Callable that knows how to process ServiceStateChange events

* `listener`: ServiceListener

* `addr`: address to send queries (will default to multicast)

* `port`: port to send queries (will default to mdns 5353)

* `delay`: The initial delay between answering questions

* `question_type`: The type of questions to ask (DNSQuestionType.QM or DNSQuestionType.QU)

The listener object will have its add_service() and

remove_service() methods called when this browser

discovers changes in the services availability.

"""

assert handlers or listener, 'You need to specify at least one handler'

self.types: Set[str] = set(type_ if isinstance(type_, list) else [type_])

for check_type_ in self.types:

# Will generate BadTypeInNameException on a bad name

service_type_name(check_type_, strict=False)

self.zc = zc

self.addr = addr

self.port = port

self.multicast = self.addr in (None, _MDNS_ADDR, _MDNS_ADDR6)

self.question_type = question_type

self._pending_handlers: OrderedDict[Tuple[str, str], ServiceStateChange] = OrderedDict()

self._service_state_changed = Signal()

self.query_scheduler = QueryScheduler(self.types, delay, _FIRST_QUERY_DELAY_RANDOM_INTERVAL)

self.queue: Optional[queue.SimpleQueue] = None

self.done = False

self._first_request: bool = True

self._next_send_timer: Optional[asyncio.TimerHandle] = None

if hasattr(handlers, 'add_service'):

listener = cast('ServiceListener', handlers)

handlers = None

handlers = cast(List[Callable[..., None]], handlers or [])

if listener:

handlers.append(_service_state_changed_from_listener(listener))

for h in handlers:

self.service_state_changed.register_handler(h)

def _async_start(self) -> None:

"""Generate the next time and setup listeners.

Must be called by uses of this base class after they

have finished setting their properties.

"""

self.query_scheduler.start(current_time_millis())

self.zc.async_add_listener(self, [DNSQuestion(type_, _TYPE_PTR, _CLASS_IN) for type_ in self.types])

# Only start queries after the listener is installed

asyncio.ensure_future(self._async_start_query_sender())

@property

def service_state_changed(self) -> SignalRegistrationInterface:

return self._service_state_changed.registration_interface

def _names_matching_types(self, names: Iterable[str]) -> List[Tuple[str, str]]:

"""Return the type and name for records matching the types we are browsing."""

return [(type_, name) for name in names for type_ in self.types.intersection(possible_types(name))]

def _enqueue_callback(

self,

state_change: ServiceStateChange,

type_: str,

name: str,

) -> None:

# Code to ensure we only do a single update message

# Precedence is; Added, Remove, Update

key = (name, type_)

if (

state_change is ServiceStateChange.Added

or (

state_change is ServiceStateChange.Removed

and self._pending_handlers.get(key) != ServiceStateChange.Added

)

or (state_change is ServiceStateChange.Updated and key not in self._pending_handlers)

):

self._pending_handlers[key] = state_change

def _async_process_record_update(

self, now: float, record: DNSRecord, old_record: Optional[DNSRecord]

) -> None:

"""Process a single record update from a batch of updates."""

if isinstance(record, DNSPointer):

for type_ in self.types.intersection(possible_types(record.name)):

if old_record is None:

self._enqueue_callback(ServiceStateChange.Added, type_, record.alias)

elif record.is_expired(now):

self._enqueue_callback(ServiceStateChange.Removed, type_, record.alias)

else:

self.reschedule_type(type_, now, record.get_expiration_time(_EXPIRE_REFRESH_TIME_PERCENT))

return

# If its expired or already exists in the cache it cannot be updated.

if old_record or record.is_expired(now):

return

if isinstance(record, DNSAddress):

# Iterate through the DNSCache and callback any services that use this address

for type_, name in self._names_matching_types(

{service.name for service in self.zc.cache.async_entries_with_server(record.name)}

):

self._enqueue_callback(ServiceStateChange.Updated, type_, name)

return

for type_, name in self._names_matching_types((record.name,)):

self._enqueue_callback(ServiceStateChange.Updated, type_, name)

def async_update_records(self, zc: 'Zeroconf', now: float, records: List[RecordUpdate]) -> None:

"""Callback invoked by Zeroconf when new information arrives.

Updates information required by browser in the Zeroconf cache.

Ensures that there is are no unecessary duplicates in the list.

This method will be run in the event loop.

"""

for record in records:

self._async_process_record_update(now, record[0], record[1])

def async_update_records_complete(self) -> None:

"""Called when a record update has completed for all handlers.

At this point the cache will have the new records.

This method will be run in the event loop.

"""

while self._pending_handlers:

event = self._pending_handlers.popitem(False)

# If there is a queue running (ServiceBrowser)

# get fired in dedicated thread

if self.queue:

self.queue.put(event)

else:

self._fire_service_state_changed_event(event)

def _fire_service_state_changed_event(self, event: Tuple[Tuple[str, str], ServiceStateChange]) -> None:

"""Fire a service state changed event.

When running with ServiceBrowser, this will happen in the dedicated

thread.

When running with AsyncServiceBrowser, this will happen in the event loop.

"""

name_type, state_change = event

self._service_state_changed.fire(

zeroconf=self.zc,

service_type=name_type[1],

name=name_type[0],

state_change=state_change,

)

def _async_cancel(self) -> None:

"""Cancel the browser."""

self.done = True

self._cancel_send_timer()

self.zc.async_remove_listener(self)

def _generate_ready_queries(self, first_request: bool, now: float) -> List[DNSOutgoing]:

"""Generate the service browser query for any type that is due."""

ready_types = self.query_scheduler.process_ready_types(now)

if not ready_types:

return []

# If they did not specify and this is the first request, ask QU questions

# https://datatracker.ietf.org/doc/html/rfc6762#section-5.4 since we are

# just starting up and we know our cache is likely empty. This ensures

# the next outgoing will be sent with the known answers list.

question_type = DNSQuestionType.QU if not self.question_type and first_request else self.question_type

return generate_service_query(self.zc, now, ready_types, self.multicast, question_type)

async def _async_start_query_sender(self) -> None:

"""Start scheduling queries."""

await self.zc.async_wait_for_start()

self._async_send_ready_queries_schedule_next()

def _cancel_send_timer(self) -> None:

"""Cancel the next send."""

if self._next_send_timer:

self._next_send_timer.cancel()

def reschedule_type(self, type_: str, now: float, next_time: float) -> None:

"""Reschedule a type to be refreshed in the future."""

if self.query_scheduler.reschedule_type(type_, next_time):

self._cancel_send_timer()

self._async_schedule_next(now)

self._async_send_ready_queries(now)

def _async_send_ready_queries(self, now: float) -> None:

"""Send any ready queries."""

outs = self._generate_ready_queries(self._first_request, now)

if outs:

self._first_request = False

for out in outs:

self.zc.async_send(out, addr=self.addr, port=self.port)

def _async_send_ready_queries_schedule_next(self) -> None:

"""Send ready queries and schedule next one checking for done first."""

if self.done or self.zc.done:

return

now = current_time_millis()

self._async_send_ready_queries(now)

self._async_schedule_next(now)

def _async_schedule_next(self, now: float) -> None:

"""Scheule the next time."""

assert self.zc.loop is not None

delay = millis_to_seconds(self.query_scheduler.millis_to_wait(now))

self._next_send_timer = self.zc.loop.call_later(delay, self._async_send_ready_queries_schedule_next)

class ServiceBrowser(_ServiceBrowserBase, threading.Thread):

"""Used to browse for a service of a specific type.

The listener object will have its add_service() and

remove_service() methods called when this browser

discovers changes in the services availability."""

def __init__(

self,

zc: 'Zeroconf',

type_: Union[str, list],

handlers: Optional[Union[ServiceListener, List[Callable[..., None]]]] = None,

listener: Optional[ServiceListener] = None,

addr: Optional[str] = None,

port: int = _MDNS_PORT,

delay: int = _BROWSER_TIME,

question_type: Optional[DNSQuestionType] = None,

) -> None:

assert zc.loop is not None

if not zc.loop.is_running():

raise RuntimeError("The event loop is not running")

threading.Thread.__init__(self)

super().__init__(zc, type_, handlers, listener, addr, port, delay, question_type)

# Add the queue before the listener is installed in _setup

# to ensure that events run in the dedicated thread and do

# not block the event loop

self.queue = queue.SimpleQueue()

self.daemon = True

self.start()

zc.loop.call_soon_threadsafe(self._async_start)

self.name = "zeroconf-ServiceBrowser-{}-{}".format(

'-'.join([type_[:-7] for type_ in self.types]),

getattr(self, 'native_id', self.ident),

)

def cancel(self) -> None:

"""Cancel the browser."""

assert self.zc.loop is not None

assert self.queue is not None

self.queue.put(None)

self.zc.loop.call_soon_threadsafe(self._async_cancel)

self.join()

def run(self) -> None:

"""Run the browser thread."""

assert self.queue is not None

while True:

event = self.queue.get()

if event is None:

return

self._fire_service_state_changed_event(event)