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()
Exemple #2
0
async def appstart(loop):
    """Script starts here."""
    parser = argparse.ArgumentParser()
    parser.add_argument("--local-ip",
                        default="127.0.0.1",
                        help="local IP address")
    parser.add_argument("-d",
                        "--debug",
                        default=False,
                        action="store_true",
                        help="enable debug logs")
    args = parser.parse_args()

    level = logging.DEBUG if args.debug else logging.WARNING
    logging.basicConfig(level=level, stream=sys.stdout)

    state = FakeDeviceState()
    zconf = Zeroconf(loop)
    server = await loop.create_server(lambda: FakeAppleTV(loop, state=state),
                                      "0.0.0.0")
    port = server.sockets[0].getsockname()[1]
    _LOGGER.info("Started fake MRP device at port %d", port)

    service = await publish_zeroconf(zconf, args.local_ip, port)

    print("Press ENTER to quit")
    await loop.run_in_executor(None, sys.stdin.readline)

    await zconf.unregister_service(service)

    return 0
    def test_integration_with_subtype_and_listener(self):
        async def run_me(zeroconf_registrar):
            subtype_ = "_subtype._sub"
            type_ = "_type._tcp.local."
            name = "xxxyyy"
            # Note: discovery returns only DNS-SD type not subtype
            discovery_type = "%s.%s" % (subtype_, type_)
            registration_name = "%s.%s" % (name, type_)

            desc = {'path': '/~paulsm/'}
            info = ServiceInfo(discovery_type, registration_name,
                               socket.inet_aton("10.0.1.2"), 80, 0, 0, desc,
                               "ash-2.local.")
            await zeroconf_registrar.register_service(info)

            try:
                service_types = await ZeroconfServiceTypes.find(
                    zc=zeroconf_registrar, timeout=0.5)
                assert discovery_type in service_types

            finally:
                await zeroconf_registrar.close()

        event_loop = asyncio.new_event_loop()
        asyncio.set_event_loop(event_loop)
        zeroconf_registrar = Zeroconf(event_loop, [netifaces.AF_INET],
                                      iface="lo")
        event_loop.run_until_complete(run_me(zeroconf_registrar))
        event_loop.close()
Exemple #4
0
async def _start_mrp_proxy(loop, args):
    def proxy_factory():
        try:
            proxy = MrpAppleTVProxy(loop)
            asyncio.ensure_future(
                proxy.start(args.remote_ip, args.remote_port, args.credentials),
                loop=loop,
            )
        except Exception:
            _LOGGER.exception("failed to start proxy")
        return proxy

    zconf = Zeroconf(loop)

    # Setup server used to publish a fake MRP server
    server = await loop.create_server(proxy_factory, "0.0.0.0")
    port = server.sockets[0].getsockname()[1]
    _LOGGER.info("Started MRP server at port %d", port)

    service = await publish_zeroconf(zconf, args.local_ip, port)

    print("Press ENTER to quit")
    await loop.run_in_executor(None, sys.stdin.readline)

    await zconf.unregister_service(service)
Exemple #5
0
async def async_setup(hass, config):
    """Set up Zeroconf and make Home Assistant discoverable."""
    from aiozeroconf import Zeroconf, ServiceInfo

    zeroconf_name = '{}.{}'.format(hass.config.location_name, ZEROCONF_TYPE)

    params = {
        'version': __version__,
        'base_url': hass.config.api.base_url,
        # always needs authentication
        'requires_api_password': True,
    }

    info = ServiceInfo(ZEROCONF_TYPE,
                       zeroconf_name,
                       port=hass.http.server_port,
                       properties=params)

    zeroconf = Zeroconf(hass.loop)

    await zeroconf.register_service(info)

    async def stop_zeroconf(event):
        """Stop Zeroconf."""
        await zeroconf.unregister_service(info)
        await zeroconf.close()

    hass.bus.async_listen_once(EVENT_HOMEASSISTANT_STOP, stop_zeroconf)

    return True
Exemple #6
0
async def discover_mdns(websession, timeout=5):
    discovered_bridge = []

    async def do_close(zc):
        await zc.close()

    def on_service_state_change(zc, service_type, name, state_change):
        if state_change is ServiceStateChange.Added:
            asyncio.ensure_future(
                on_service_state_change_process(zc, service_type, name))

    async def on_service_state_change_process(zc, service_type, name):
        info = await zc.get_service_info(service_type, name)
        discovered_bridge.append(info)
        await do_close(zc)

    async def find_hue(zc):
        browser = ServiceBrowser(zc,
                                 "_hue._tcp.local.",
                                 handlers=[on_service_state_change])
        await asyncio.sleep(timeout)
        browser.cancel()

    loop = asyncio.get_event_loop()
    zeroconf = Zeroconf(loop)
    await find_hue(zeroconf)
    return [
        Bridge(ipaddress.ip_address(item.address), websession)
        for item in discovered_bridge
    ]
Exemple #7
0
def main():
    parser = argparse.ArgumentParser(
        description="Zeroconf service discovery tool")
    parser.add_argument('-i',
                        "--iface",
                        default="",
                        help="Name of the inteface to use.")
    parser.add_argument('-p',
                        "--protocol",
                        choices=['ipv4', 'ipv6', 'both'],
                        default="ipv4",
                        help="What IP protocol to use.")
    parser.add_argument("-s",
                        "--service",
                        default="_http._tcp.local.",
                        help="The service to browse.")
    parser.add_argument("-f",
                        "--find",
                        action='store_true',
                        default=False,
                        help="Find services")
    parser.add_argument("-d",
                        "--debug",
                        action='store_true',
                        default=False,
                        help="Set debug mode.")
    try:
        opts = parser.parse_args()
    except Exception as e:
        parser.error("Error: " + str(e))

    if opts.protocol == "ipv4":
        proto = [netifaces.AF_INET]
    elif opts.protocol == "ipv6":
        proto = [netifaces.AF_INET6]
    else:
        proto = [netifaces.AF_INET, netifaces.AF_INET6]

    loop = asyncio.get_event_loop()
    logging.basicConfig(level=logging.CRITICAL)
    if opts.debug:
        logging.getLogger('zeroconf').setLevel(logging.DEBUG)
        loop.set_debug(True)

    zc = Zeroconf(loop, proto, iface=opts.iface)
    print("\nBrowsing services, press Ctrl-C to exit...\n")

    try:
        if opts.find:
            loop.run_until_complete(list_service(zc))
        else:
            ServiceBrowser(zc,
                           guess(opts.service),
                           handlers=[on_service_state_change])
            loop.run_forever()
    except KeyboardInterrupt:
        print("Unregistering...")
        loop.run_until_complete(do_close(zc))
    finally:
        loop.close()
Exemple #8
0
def asyncio_zeroconf():
    """Builds an aiozeroconf service instance and starts browsing for WoT Servient services.
    Provides a deque that contains the service state change history."""

    from aiozeroconf import Zeroconf, ServiceBrowser
    from wotpy.wot.discovery.dnssd.service import DNSSDDiscoveryService

    loop = tornado.ioloop.IOLoop.current()

    service_history = collections.deque([])

    def on_change(zc, service_type, name, state_change):
        service_history.append((service_type, name, state_change))

    aio_zc = Zeroconf(loop.asyncio_loop, address_family=[socket.AF_INET])
    ServiceBrowser(aio_zc,
                   DNSSDDiscoveryService.WOT_SERVICE_TYPE,
                   handlers=[on_change])

    yield {"zeroconf": aio_zc, "service_history": service_history}

    @tornado.gen.coroutine
    def close():
        yield aio_zc.close()

    loop.run_sync(close)
 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])
Exemple #10
0
async def _discover(id: str, dns_type: str, single: bool):
    queue = asyncio.Queue()
    conf = Zeroconf(asyncio.get_event_loop(), address_family=[AF_INET])

    async def add_service(service_type, name):
        info = await conf.get_service_info(service_type, name)
        await queue.put(info)

    def sync_change_handler(_, service_type, name, state_change):
        if state_change is ServiceStateChange.Added:
            asyncio.create_task(add_service(service_type, name))

    try:
        ServiceBrowser(conf, dns_type, handlers=[sync_change_handler])
        match = f'{id}.local.'.lower() if id else None

        while True:
            info = await queue.get()
            addr = inet_ntoa(info.address)
            if addr == '0.0.0.0':
                continue  # discard simulators
            if match is None or info.server.lower() == match:
                serial = info.server[:-len('.local.')]
                yield addr, info.port, serial
                if single:
                    return
    finally:
        await conf.close()
Exemple #11
0
 def __init__(self, config, session, loop, **kwargs):
     """Initialize a new instance."""
     super().__init__(session, config.get_service(Protocol.DMAP))
     self._loop = loop
     self._zeroconf = kwargs.get(
         'zeroconf', Zeroconf(loop, address_family=[netifaces.AF_INET]))
     self._name = kwargs.get('name', 'pyatv')
     self._web_server = None
     self._server = None
     self._pin_code = None
     self._has_paired = False
     self._pairing_guid = (kwargs.get('pairing_guid', None)
                           or _generate_random_guid())[2:].upper()
Exemple #12
0
 async def discover(self, timeout):
     """Start discovery of devices and services."""
     zeroconf = Zeroconf(self.loop, address_family=[netifaces.AF_INET])
     try:
         ServiceBrowser(zeroconf, HOMESHARING_SERVICE, self)
         ServiceBrowser(zeroconf, DEVICE_SERVICE, self)
         ServiceBrowser(zeroconf, MEDIAREMOTE_SERVICE, self)
         ServiceBrowser(zeroconf, AIRPLAY_SERVICE, self)
         _LOGGER.debug('Discovering devices for %d seconds', timeout)
         await asyncio.sleep(timeout)
     finally:
         await zeroconf.close()
     return self._found_devices
Exemple #13
0
async def simple_find_service(
    loop,
    service=SERVICE,
    timeout=SERVICE_DISCOVERY_TIMEOUT,
    stop_after_first=STOP_AFTER_FIRST,
):
    # result = await z.get_service_info(type_=SERVICE, name=SERVICE, timeout=SERVICE_DISCOVERY_TIMEOUT)
    log.debug("Starting zeroconf discovery")
    zeroconf = Zeroconf(loop)
    result = await ZeroconfServiceTypes.find(
        zeroconf, service=service, timeout=timeout, stop_after_first=stop_after_first
    )
    # print(result)
    await do_close(zeroconf)
    return result
Exemple #14
0
 def __init__(self, config, session_manager: ClientSessionManager, loop,
              **kwargs):
     """Initialize a new instance."""
     super().__init__(session_manager, config.get_service(Protocol.DMAP))
     self._loop = loop
     self._zeroconf = kwargs.get(
         "zeroconf", Zeroconf(loop, address_family=[netifaces.AF_INET]))
     self._name = kwargs.get("name", "pyatv")
     self.app = web.Application()
     self.app.router.add_routes([web.get("/pairing", self.handle_request)])
     self.runner = web.AppRunner(self.app)
     self.site = None
     self._pin_code = None
     self._has_paired = False
     self._pairing_guid = (kwargs.get("pairing_guid", None)
                           or _generate_random_guid())[2:].upper()
Exemple #15
0
async def appstart(loop):
    """Start the asyncio event loop and runs the application."""
    parser = argparse.ArgumentParser()
    subparsers = parser.add_subparsers(title="sub-commands", dest="command")

    mrp = subparsers.add_parser("mrp", help="MRP proxy")
    mrp.add_argument("credentials", help="MRP credentials")
    mrp.add_argument("remote_ip", help="Apple TV IP address")
    mrp.add_argument("--name", default=None, help="proxy device name")
    mrp.add_argument("--local-ip", default=None, help="local IP address")
    mrp.add_argument("--remote_port", default=None, help="MRP port")

    relay = subparsers.add_parser("relay", help="Relay traffic to host")
    relay.add_argument("local_ip", help="local IP address")
    relay.add_argument("remote_ip", help="Remote host")
    relay.add_argument("remote_port", help="Remote port")
    relay.add_argument("name", help="Service name")
    relay.add_argument("service_type", help="Service type")
    relay.add_argument("-p",
                       "--properties",
                       nargs="+",
                       help="Service properties")

    args = parser.parse_args()
    if not args.command:
        parser.error("No command specified")
        return 1

    # To get logging from pyatv
    logging.basicConfig(
        level=logging.DEBUG,
        stream=sys.stdout,
        datefmt="%Y-%m-%d %H:%M:%S",
        format="%(asctime)s %(levelname)s: %(message)s",
    )

    zconf = Zeroconf(loop)
    if args.command == "mrp":
        service = await _start_mrp_proxy(loop, args, zconf)
    elif args.command == "relay":
        service = await _start_relay(loop, args, zconf)
    await zconf.unregister_service(service)

    return 0
Exemple #16
0
def main(loop):
    """Script starts here."""
    # To get logging from pyatv
    logging.basicConfig(level=logging.DEBUG)

    if len(sys.argv) != 6:
        print("Usage: {0} <credentials> <local ip> "
              "<atv ip> <atv port> <unique identifier>".format(sys.argv[0]))
        sys.exit(1)

    credentials = sys.argv[1]
    local_ip_addr = sys.argv[2]
    atv_ip_addr = sys.argv[3]
    atv_port = int(sys.argv[4])
    unique_identifier = sys.argv[5].encode()
    zconf = Zeroconf(loop)
    proxy = ProxyMrpAppleTV(loop, credentials, unique_identifier)

    proxy.start(atv_ip_addr, atv_port)
    loop.run_until_complete(publish_zeroconf(zconf, local_ip_addr, proxy.port))
    loop.run_forever()
Exemple #17
0
    async def discover(self, timeout):
        """Start discovery of devices and services."""
        zeroconf = Zeroconf(self.loop, address_family=[netifaces.AF_INET])
        browsers = []
        try:
            browsers += [
                ServiceBrowser(zeroconf, HOMESHARING_SERVICE, self),
                ServiceBrowser(zeroconf, DEVICE_SERVICE, self),
                ServiceBrowser(zeroconf, MEDIAREMOTE_SERVICE, self),
                ServiceBrowser(zeroconf, AIRPLAY_SERVICE, self),
            ]
            _LOGGER.debug("Discovering devices for %d seconds", timeout)
            await asyncio.sleep(timeout)

            if self.pending:
                await asyncio.wait(self.pending)
        finally:
            for browser in browsers:
                browser.cancel()
            await zeroconf.close()
        return self._found_devices
Exemple #18
0
async def main(loop):
    """Script starts here."""
    # To get logging from pyatv
    logging.basicConfig(level=logging.DEBUG, stream=sys.stdout)

    if len(sys.argv) != 5:
        print("Usage: {0} <credentials> <local ip> "
              "<atv ip> <atv port>".format(
                  sys.argv[0]))
        sys.exit(1)

    credentials = sys.argv[1]
    local_ip_addr = sys.argv[2]
    atv_ip_addr = sys.argv[3]
    atv_port = int(sys.argv[4])
    zconf = Zeroconf(loop)

    def proxy_factory():
        try:
            proxy = ProxyMrpAppleTV(loop)
            asyncio.ensure_future(
                proxy.start(atv_ip_addr, atv_port, credentials),
                loop=loop)
        except Exception:
            _LOGGER.exception("failed to start proxy")
        return proxy

    # Setup server used to publish a fake MRP server
    server = await loop.create_server(proxy_factory, '0.0.0.0')
    port = server.sockets[0].getsockname()[1]
    _LOGGER.error('Started MRP server at port %d', port)

    service = await publish_zeroconf(zconf, local_ip_addr, port)

    print("Press ENTER to quit")
    await loop.run_in_executor(None, sys.stdin.readline)

    await zconf.unregister_service(service)
async def _discover(
    desired_id: Optional[str],
    dns_type: str,
) -> Generator[ConnectInfo, None, None]:
    queue: asyncio.Queue[ServiceInfo] = asyncio.Queue()
    conf = Zeroconf(asyncio.get_event_loop(), address_family=[AF_INET])

    async def add_service(service_type, name):
        info = await conf.get_service_info(service_type, name)
        await queue.put(info)

    def sync_change_handler(_, service_type, name, state_change):
        if state_change is ServiceStateChange.Added:
            asyncio.create_task(add_service(service_type, name))

    try:
        ServiceBrowser(conf, dns_type, handlers=[sync_change_handler])

        while True:
            info = await queue.get()
            if info.address in [None, SIM_ADDR]:
                continue  # discard unknown addresses and simulators

            addr = inet_ntoa(info.address)
            id = info.properties.get(b'ID', bytes()).decode().lower()

            if not id:
                LOGGER.error(
                    f'Invalid device: {info.name} @ {addr}:{info.port} has no ID TXT property'
                )
                continue
            elif desired_id is None or desired_id.lower() == id:
                LOGGER.info(f'Discovered {id} @ {addr}:{info.port}')
                yield ConnectInfo(addr, info.port, id)
            else:
                LOGGER.info(f'Discarding {info.name} @ {addr}:{info.port}')
    finally:
        await conf.close()
Exemple #20
0
async def scan_devices(timeout=5):
    """
    Scan for devices on the network.

    :return: A list of devices found.
    """
    loop = asyncio.get_running_loop()
    zeroconf = Zeroconf(loop)
    devices = []

    def on_service_state_change(zc, service_type, name, change):
        host = name.split(".")[0] + ".local"
        host = socket.gethostbyname(host)
        device = Device(base_url=f"http://{host}")
        devices.append(device)

    ServiceBrowser(zeroconf,
                   "_ohm-led._tcp.local.",
                   handlers=[on_service_state_change])

    await asyncio.sleep(timeout)
    await zeroconf.close()

    return devices
Exemple #21
0
                    default=False,
                    help="Set debug mode.")
try:
    opts = parser.parse_args()
except Exception as e:
    parser.error("Error: " + str(e))

if opts.protocol == "ipv4":
    proto = [netifaces.AF_INET]
elif opts.protocol == "ipv6":
    proto = [netifaces.AF_INET6]
else:
    proto = [netifaces.AF_INET, netifaces.AF_INET6]

loop = asyncio.get_event_loop()
logging.basicConfig(level=logging.CRITICAL)
if opts.debug:
    logging.getLogger('zeroconf').setLevel(logging.DEBUG)
    loop.set_debug(True)

zc = Zeroconf(loop, proto, iface=opts.iface)
print("\nBrowsing services, press Ctrl-C to exit...\n")
browser = ServiceBrowser(zc, opts.service, handlers=[on_service_state_change])
try:
    loop.run_forever()
except KeyboardInterrupt:
    print("Unregistering...")
    loop.run_until_complete(do_close(zc))
finally:
    loop.close()
Exemple #22
0
            return ip4


async def do_close(zc):
    global info
    await zc.unregister_service(info)
    await zc.close()


async def on_shutdown(app):
    print("Unregistering...")
    await do_close(zc)


loop = asyncio.get_event_loop()
zc = Zeroconf(loop)

if __name__ == "__main__":
    logging.basicConfig(level=logging.DEBUG)
    if len(sys.argv) > 1:
        assert sys.argv[1:] == ["--debug"]
        logging.getLogger("aiozeroconf").setLevel(logging.DEBUG)

    try:
        xx = loop.create_task(run_test(zc))
        loop.run_forever()
    except KeyboardInterrupt:
        print("Unregistering...")
        loop.run_until_complete(do_close(zc))
    finally:
        loop.close()
Exemple #23
0
async def appstart(loop):
    """Script starts here."""
    parser = argparse.ArgumentParser()
    parser.add_argument("--local-ip",
                        default="127.0.0.1",
                        help="local IP address")
    parser.add_argument("--demo",
                        default=False,
                        action="store_true",
                        help="enable demo mode")
    parser.add_argument("-d",
                        "--debug",
                        default=False,
                        action="store_true",
                        help="enable debug logs")

    protocols = parser.add_argument_group("protocols")
    protocols.add_argument("--mrp",
                           default=False,
                           action="store_true",
                           help="enable MRP protocol")
    protocols.add_argument("--dmap",
                           default=False,
                           action="store_true",
                           help="enable DMAP protocol")
    protocols.add_argument("--airplay",
                           default=False,
                           action="store_true",
                           help="enable AirPlay protocol")
    args = parser.parse_args()

    if not (args.mrp or args.dmap or args.airplay):
        parser.error("no protocol enabled (see --help)")

    level = logging.DEBUG if args.debug else logging.WARNING
    logging.basicConfig(
        level=level,
        stream=sys.stdout,
        datefmt="%Y-%m-%d %H:%M:%S",
        format="%(asctime)s %(levelname)s: %(message)s",
    )

    tasks = []
    services = []
    zconf = Zeroconf(loop)
    fake_atv = FakeAppleTV(loop, test_mode=False)
    if args.mrp:
        _, usecase = fake_atv.add_service(Protocol.MRP)
        if args.demo:
            tasks.append(asyncio.ensure_future(_alter_playing(usecase)))

    if args.dmap:
        _, usecase = fake_atv.add_service(Protocol.DMAP,
                                          hsgid=HSGID,
                                          pairing_guid=PAIRING_GUID,
                                          session_id=SESSION_ID)
        if args.demo:
            tasks.append(asyncio.ensure_future(_alter_playing(usecase)))

    if args.airplay:
        _, usecase = fake_atv.add_service(Protocol.AirPlay)

    await fake_atv.start()

    if args.mrp:
        services.append(await
                        publish_mrp_zeroconf(zconf, args.local_ip,
                                             fake_atv.get_port(Protocol.MRP)))

    if args.dmap:
        services.append(await publish_dmap_zeroconf(
            zconf, args.local_ip, fake_atv.get_port(Protocol.DMAP)))

    if args.airplay:
        services.append(await publish_airplay_zeroconf(
            zconf, args.local_ip, fake_atv.get_port(Protocol.AirPlay)))

    print("Press ENTER to quit")
    await loop.run_in_executor(None, sys.stdin.readline)

    await fake_atv.stop()

    for task in tasks:
        task.cancel()

    for service in services:
        await zconf.unregister_service(service)

    print("Exiting")

    return 0
Exemple #24
0
 def start_yeelight_discovery(handler, gone_handler=None, iface=None):
     loop = aio.get_event_loop()
     zeroconf = Zeroconf(loop, iface=iface)
     listener = YeelightListener(handler, gone_handler)
     browser = ServiceBrowser(zeroconf, "_miio._udp.local.", listener)
     return zeroconf, browser
Exemple #25
0
class Server(object):
    '''TCP server object that broadcasts its availability using zeroconf'''

    def __init__(self, service_type, port):
        '''Create a TCP server'''
        self.__logger = logging.getLogger(__name__)

        self.connection_changed = Event(sender='server')

        # this keeps track of all the clients that connected to our
        # server.  It can be useful in some cases, for instance to
        # kill client connections or to broadcast some data to all
        # clients...
        self.__clients = {} # task -> (reader, writer)
        self.__port = port
        self.__server = None
        self.__loop = None
        self.__shutdown_in_progress = False
        self.__queue = asyncio.Queue()
        self.__info = ServiceInfo(
            service_type,
            'TTC-%s.%s' % (
                uuid.uuid3(uuid.NAMESPACE_DNS, socket.gethostname()),
                service_type),
            socket.inet_aton(
                socket.gethostbyname(socket.gethostname())),
            self.__port,
            0,
            0,
            {},
            socket.gethostname() + '.')

    def start(self, loop):
        '''Start the server'''
        self.__logger.debug('Server started')

        # 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 TCP server
        start_tcp_task = self.__loop.create_task(self.__start_tcp())
        start_tcp_task.add_done_callback(self.__start_broadcast)

    async def stop(self):
        '''
        Stop the server

        Stop zeroconf advertising and close listening socket(s). This method
        runs the loop until the server sockets are closed.
        '''
        if self.is_running():
            self.__shutdown_in_progress = True

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

            self.__queue.put_nowait(b'')

            await self.__server.wait_closed()

            self.__server = None

        self.__shutdown_in_progress = False


    def is_running(self):
        '''Inidication that the server is running'''
        return (self.__server is not None)

    async def __start_tcp(self):
        '''
        Start the TCP server process

        Starts a TCP streaming server that services all interfaces on the device
        Starts the write process to service the incoming message queue
        '''
        self.__server = await asyncio.start_server(
            self.__accept_client,
            '0.0.0.0',
            self.__port,
            loop=self.__loop)

        self.__loop.create_task(self.__write_process())

    def __accept_client(self, reader, writer):
        '''
        Handles incoming client connections
        '''
        # Start a new asyncio.Task to handle this specific client connection
        task = self.__loop.create_task(self.__handle_client_read(reader))

        # Store a tuple for the client connection indexed by the task for the
        # connection
        self.__clients[task] = (reader, writer)

        # Notify of connection change
        self.__connection_changed()

        # Add the client_done callback to be run when the future becomes done
        task.add_done_callback(self.__client_done)

    def __client_done(self, task):
        '''
        Client cleanup process
        '''
        # When the tasks that handles the specific client connection is done
        client = self.__clients[task]
        client[1].close()

        del self.__clients[task]

        # Notify of connection change
        self.__connection_changed()

        if not self.__clients and self.__shutdown_in_progress:
            self.__server.close()

    def __connection_changed(self):
        '''
        '''
        client_count = len(self.__clients)

        self.__logger.debug('Client(s) connected: {0}'.format(client_count))

        self.connection_changed(client_count)

    async def __handle_client_read(self, reader):
        '''
        Client read process
        '''
        while not reader.at_eof():
            data = await reader.read(4)

            if data:
                size = int.from_bytes(data, 'little')

                data = await reader.read(size)

                try:
                    self.__queue.put_nowait(data)
                except asyncio.QueueFull:
                    self.__logger.warning('Queue full, data lost')

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

            if data:
                # Valid data gets repeated to all clients including the one who produced it?
                for client in self.__clients:
                    # Pull the writer out of the client tuple (reader, writer)
                    writer = self.__clients[client][1]

                    # 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)

                    await writer.drain()

            self.__queue.task_done()

            if not data:
                for client in self.__clients:
                    # Pull the writer out of the client tuple (reader, writer)
                    writer = self.__clients[client][1]

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

                break

        if not self.__clients and self.__shutdown_in_progress:
            # There are no clients connected so shutdown the server
            self.__server.close()

    def __start_broadcast(self, task):
        '''Start zeroconf service broadcast'''
        self.__zc = Zeroconf(self.__loop, address_family = [netifaces.AF_INET])
        self.__loop.create_task(self.__zc.register_service(self.__info))

    async def __stop_broadcast(self):
        '''Stop zeroconf service broadcast'''
        await self.__zc.unregister_service(self.__info)
        await self.__zc.close()
    def test_integration_with_listener_class(self):

        service_added = Event()
        service_removed = Event()

        async def run_me(zcbrowser, zcregistrar):
            subtype_name = "My special Subtype"
            type_ = "_http._tcp.local."
            subtype = subtype_name + "._sub." + type_
            name = "xxxyyy"
            registration_name = "%s.%s" % (name, type_)

            class MyListener(object):
                def add_service(self, zeroconf, type, name):
                    asyncio.ensure_future(
                        self.async_add_service(zeroconf, type, name))

                async def async_add_service(self, zeroconf, type, name):
                    await zeroconf.get_service_info(type, name)
                    service_added.set()

                def remove_service(self, zeroconf, type, name):
                    service_removed.set()

            listener = MyListener()
            zcbrowser.add_service_listener(subtype, listener=listener)

            properties = dict(
                prop_none=None,
                prop_string=b'a_prop',
                prop_float=1.0,
                prop_blank=b'a blanked string',
                prop_true=1,
                prop_false=0,
            )

            desc = {'path': '/~paulsm/'}
            desc.update(properties)
            info_service = ServiceInfo(subtype, registration_name,
                                       socket.inet_aton("10.0.1.2"), 80, 0, 0,
                                       desc, "ash-2.local.")
            await zcregistrar.register_service(info_service)

            try:
                await asyncio.sleep(2)
                assert service_added.is_set()

                # short pause to allow multicast timers to expire
                await asyncio.sleep(2)

                # clear the answer cache to force query
                for record in zcbrowser.cache.entries():
                    zcbrowser.cache.remove(record)

                # get service info without answer cache
                info = await zcbrowser.get_service_info(
                    type_, registration_name)

                assert info.properties[b'prop_none'] is False
                assert info.properties[b'prop_string'] == properties[
                    'prop_string']
                assert info.properties[b'prop_float'] is False
                assert info.properties[b'prop_blank'] == properties[
                    'prop_blank']
                assert info.properties[b'prop_true'] is True
                assert info.properties[b'prop_false'] is False

                info = await zcbrowser.get_service_info(
                    subtype, registration_name)
                assert info.properties[b'prop_none'] is False

                await zcregistrar.unregister_service(info_service)
                await asyncio.sleep(1)
                assert service_removed.is_set()
            finally:
                await zcregistrar.close()
                zcbrowser.remove_service_listener(listener)
                await zcbrowser.close()

        event_loop = asyncio.new_event_loop()
        asyncio.set_event_loop(event_loop)
        zeroconf_browser = Zeroconf(event_loop, [netifaces.AF_INET],
                                    iface="lo")
        zeroconf_registrar = Zeroconf(event_loop, [netifaces.AF_INET],
                                      iface="lo")
        coro = run_me(zeroconf_browser, zeroconf_registrar)
        event_loop.run_until_complete(coro)
        event_loop.close()
Exemple #27
0
 def __start_broadcast(self, task):
     '''Start zeroconf service broadcast'''
     self.__zc = Zeroconf(self.__loop, address_family = [netifaces.AF_INET])
     self.__loop.create_task(self.__zc.register_service(self.__info))
 def setUpClass(cls):
     cls.loop = asyncio.get_event_loop()
     cls.browser = Zeroconf(cls.loop, [netifaces.AF_INET], iface="lo")
        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))