async def run_me(loop):
        type_ = "_http._tcp.local."
        registration_name = "xxxyyy.%s" % type_

        def on_service_state_change(zeroconf, service_type, state_change,
                                    name):
            if name == registration_name:
                if state_change is ServiceStateChange.Added:
                    service_added.set()
                elif state_change is ServiceStateChange.Removed:
                    service_removed.set()

        zeroconf_browser = Zeroconf(loop, [netifaces.AF_INET], iface="lo")
        browser = ServiceBrowser(zeroconf_browser, type_,
                                 [on_service_state_change])

        zeroconf_registrar = Zeroconf(loop, [netifaces.AF_INET], iface="lo")
        desc = {'path': '/~paulsm/'}
        info = ServiceInfo(type_, registration_name,
                           socket.inet_aton("10.0.1.2"), 80, 0, 0, desc,
                           "ash-2.local.")
        zeroconf_registrar.register_service(info)

        try:
            await asyncio.sleep(1)
            assert service_added.is_set()
            # Don't remove service, allow close() to cleanup

        finally:
            await zeroconf_registrar.close()
            browser.cancel()
            await zeroconf_browser.close()
Example #2
0
    async def find(cls, zc, service, timeout=5, stop_after_first=STOP_AFTER_FIRST):
        """
        Return all of the advertised services on any local networks.
        :param zc: Zeroconf() instance.  Pass in an instance running
        :param timeout: seconds to wait for any responses
        :return: tuple of service type strings
        """
        listener = cls()
        browser = ServiceBrowser(zc, service, listener=listener)

        # wait for responses
        for i in range(100):
            await asyncio.sleep(timeout / 100)
            if len(listener.found_services) > 0 and STOP_AFTER_FIRST:
                break

        browser.cancel()
        # print(listener.found_services)
        return listener.found_services
class Client(object):
    '''TCP client object that searches for a server using zeroconf'''
    def __init__(self, service_type, port):
        '''Create a TCP client'''
        self.__logger = logging.getLogger(__name__)

        self.connection_changed = Event(sender='client')
        self.data_rx = Event(sender='client')

        self.__service_type = service_type
        self.__loop = None
        self.__browser = None
        self.__port = port
        self.__queue = asyncio.Queue()
        self.__server_connection = None
        self.__shutdown_in_progress = False

    def start(self, loop):
        '''Start the client'''
        # Exit here if we're already running, we don't want to randomly
        # restart, log as warning
        if self.is_running():
            return

        self.__loop = loop

        # Start zeroconf service broadcast
        self.__start_service_discovery()

    async def stop(self):
        '''Stop the client'''
        self.__shutdown_in_progress = True

        if self.__is_browsing():
            self.__logger.debug("I'm going to stop browsing now...")

            # Stop zeroconf service discovery first so that we don't collect more
            # clients as were trying to shutdown
            await self.__stop_service_discovery()

        # Only do this next bit if we're actually connected
        if self.__is_connected():
            self.__logger.debug("I'm going to start disconnecting now...")

            await self.__stop_server_connection()

        self.__shutdown_in_progress = False

    def is_running(self):
        '''Inidication that the client is running'''
        return self.__is_browsing() or self.__is_connected()

    async def __stop_server_connection(self):
        # Pushing an empty byte into the queue will cause the write_task to
        # end, and should take the read_task with it... fingers crossed...
        self.__queue.put_nowait(b'')

        self.__logger.debug('About to wait for server connection to terminate')

        await self.__server_connection

        self.__logger.debug('Server connection terminated')

    def send(self, data):
        '''Send data on the client interface'''
        try:
            self.__queue.put_nowait(data)
        except asyncio.QueueFull:
            self.__logger.warning('Queue full, data lost')

    def __is_browsing(self):
        '''Indication that the client is browsing for a server'''
        return (self.__browser is not None)

    def __is_connected(self):
        '''Indication that the client is connected to a server'''
        return (self.__server_connection is not None)

    def __on_service_state_change(self, zc, service_type, name, state_change):
        '''
        Handle service changes

        If we 'Added' a service then intiate a connection
        '''
        self.__logger.debug(
            'Service {0} of type {1} state changed: {2}'.format(
                name, service_type, state_change))

        if state_change is ServiceStateChange.Added:
            self.__loop.create_task(self.__found_service(name))

    async def __found_service(self, name):
        '''
        Start connection process

        Found the service we were looking for, start the connection process
        '''
        info = await self.__zc.get_service_info(self.__service_type, name)

        if info:
            self.__logger.debug("Address: %s:%d" %
                                (socket.inet_ntoa(info.address), info.port))
            self.__logger.debug("Server: %s" % (info.server, ))

            try:
                reader, writer = await asyncio.open_connection(
                    socket.inet_ntoa(info.address),
                    info.port,
                    loop=self.__loop)

                self.__server_connection = self.__loop.create_task(
                    self.__connected_process(reader, writer))
                self.__server_connection.add_done_callback(
                    self.__disconnected_process)
            except ConnectionRefusedError:
                self.__logger.debug('Connection refused')

    def __start_service_discovery(self):
        '''Start zeroconf service discovery'''
        self.__zc = Zeroconf(self.__loop, address_family=[netifaces.AF_INET])
        self.__browser = ServiceBrowser(
            self.__zc,
            self.__service_type,
            handlers=[self.__on_service_state_change])

    async def __stop_service_discovery(self):
        '''Stop zeroconf service discovery'''
        self.__browser.cancel()
        await self.__zc.close()
        self.__browser = None
        self.__zc = None

    async def __handle_server_read(self, reader):
        '''Server read process'''
        # Keep reading until we encounter a null byte
        while not reader.at_eof():
            # Read the size bytes first
            data = await reader.read(4)

            # If data is not null byte read data from stream
            if data:
                size = int.from_bytes(data, 'little')

                data = await reader.read(size)

                # Send data to receiving process
                self.data_rx(data)

        # If we aren't shutting dow (because of a client.stop) put a null byte
        # in the buffer to cause the write process to terminate
        if not self.__shutdown_in_progress:
            self.__queue.put_nowait(b'')

    async def __handle_server_write(self, writer):
        '''Server write process'''
        while True:
            # Wait for new data from the queue
            data = await self.__queue.get()

            if data:
                # send the size first so that the receiver knows how many bytes
                # to expect
                size = len(data)

                writer.write(size.to_bytes(4, 'little'))

                await writer.drain()

                # Now write the data out
                writer.write(data)

                # Pause the process to let the write out happen
                await writer.drain()

            self.__queue.task_done()

            if not data:
                break

        # Close the writer/transport, this may need to be paired with a
        # write_eof
        writer.close()

    async def __connected_process(self, reader, writer):
        '''
        Process that runs on connect
            1. Stop service discovery, we've already foudn someone that wants
               to talk to us
            2. Setup tasks to handle communication R/W with the server and then
               wait
        '''
        self.__logger.debug('connection changed')

        self.connection_changed(1)

        self.__logger.debug('stopping service discovery')

        # Stop service discovery because we've found something
        await self.__stop_service_discovery()

        self.__logger.debug('set up server r/w processes')

        await asyncio.gather(*[
            self.__handle_server_read(reader),
            self.__handle_server_write(writer)
        ])

        self.__logger.debug('connected process complete')

    def __disconnected_process(self, task):
        '''
        Process that runs on disconnect
            1. Invalidate server connection
            2. Notify external actor that we are disconnected
            3. If not shutting down restart service discovery

        Task is not used.
        '''
        self.__server_connection = None

        self.connection_changed(0)

        if not self.__shutdown_in_progress:
            # Start service discovery because we disconnected not because the
            # client is being shutdown
            self.__start_service_discovery()
        async def run_me(loop):
            # instantiate a zeroconf instance
            zc = Zeroconf(loop, [netifaces.AF_INET], iface="lo")

            # create a bunch of servers
            type_ = "_my-service._tcp.local."
            name = 'a wonderful service'
            server_count = 300
            self.generate_many_hosts(zc, type_, name, server_count)

            # verify that name changing works
            self.verify_name_change(zc, type_, name, server_count)

            # we are going to monkey patch the zeroconf send to check packet sizes
            old_send = zc.send

            # needs to be a list so that we can modify it in our phony send
            longest_packet = [0, None]

            def send(out, addr=r._MDNS_ADDR, port=r._MDNS_PORT):
                """Sends an outgoing packet."""
                packet = out.packet()
                if longest_packet[0] < len(packet):
                    longest_packet[0] = len(packet)
                    longest_packet[1] = out
                old_send(out, addr=addr, port=port)

            # monkey patch the zeroconf send
            zc.send = send

            # dummy service callback
            def on_service_state_change(zeroconf, service_type, state_change,
                                        name):
                pass

            # start a browser
            browser = ServiceBrowser(zc, type_, [on_service_state_change])

            # wait until the browse request packet has maxed out in size
            sleep_count = 0
            while sleep_count < 100 and \
                    longest_packet[0] < r._MAX_MSG_ABSOLUTE - 100:
                sleep_count += 1
                await asyncio.sleep(0.1)

            browser.cancel()
            await asyncio.sleep(0.5)

            import zeroconf
            zeroconf.log.debug('sleep_count %d, sized %d', sleep_count,
                               longest_packet[0])

            # now the browser has sent at least one request, verify the size
            assert longest_packet[0] <= r._MAX_MSG_ABSOLUTE
            assert longest_packet[0] >= r._MAX_MSG_ABSOLUTE - 100

            # mock zeroconf's logger warning() and debug()
            from mock import patch
            patch_warn = patch('zeroconf.log.warning')
            patch_debug = patch('zeroconf.log.debug')
            mocked_log_warn = patch_warn.start()
            mocked_log_debug = patch_debug.start()

            # now that we have a long packet in our possession, let's verify the
            # exception handling.
            out = longest_packet[1]
            out.data.append(b'\0' * 1000)

            # mock the zeroconf logger and check for the correct logging backoff
            call_counts = mocked_log_warn.call_count, mocked_log_debug.call_count
            # try to send an oversized packet
            zc.send(out)
            assert mocked_log_warn.call_count == call_counts[0] + 1
            assert mocked_log_debug.call_count == call_counts[0]
            zc.send(out)
            assert mocked_log_warn.call_count == call_counts[0] + 1
            assert mocked_log_debug.call_count == call_counts[0] + 1

            # force a receive of an oversized packet
            packet = out.packet()
            s = zc._respond_sockets[0]

            # mock the zeroconf logger and check for the correct logging backoff
            call_counts = mocked_log_warn.call_count, mocked_log_debug.call_count
            # force receive on oversized packet
            s.sendto(packet, 0, (r._MDNS_ADDR, r._MDNS_PORT))
            s.sendto(packet, 0, (r._MDNS_ADDR, r._MDNS_PORT))
            await asyncio.sleep(2.0)
            zeroconf.log.debug('warn %d debug %d was %s',
                               mocked_log_warn.call_count,
                               mocked_log_debug.call_count, call_counts)
            assert mocked_log_debug.call_count > call_counts[0]

            # close our zeroconf which will close the sockets
            zc.close()

            # pop the big chunk off the end of the data and send on a closed socket
            out.data.pop()
            zc._GLOBAL_DONE = False

            # mock the zeroconf logger and check for the correct logging backoff
            call_counts = mocked_log_warn.call_count, mocked_log_debug.call_count
            # send on a closed socket (force a socket error)
            zc.send(out)
            zeroconf.log.debug('warn %d debug %d was %s',
                               mocked_log_warn.call_count,
                               mocked_log_debug.call_count, call_counts)
            assert mocked_log_warn.call_count > call_counts[0]
            assert mocked_log_debug.call_count > call_counts[0]
            zc.send(out)
            zeroconf.log.debug('warn %d debug %d was %s',
                               mocked_log_warn.call_count,
                               mocked_log_debug.call_count, call_counts)
            assert mocked_log_debug.call_count > call_counts[0] + 2

            mocked_log_warn.stop()
            mocked_log_debug.stop()
            loop = asyncio.get_event_loop()
            loop.run_until_complete(run_me(loop))
Example #5
0
 async def find_hue(zc):
     browser = ServiceBrowser(zc,
                              "_hue._tcp.local.",
                              handlers=[on_service_state_change])
     await asyncio.sleep(timeout)
     browser.cancel()