From c8b2f0d5ccd4b393f55d80bab4541e2368048688 Mon Sep 17 00:00:00 2001 From: Jack Robison Date: Tue, 9 Oct 2018 11:43:12 -0400 Subject: [PATCH] try m-searching upnp and wifi alliance service strings by default --- aioupnp/__main__.py | 3 +-- aioupnp/protocols/ssdp.py | 35 ++++++++++++++++++++++++----------- aioupnp/upnp.py | 7 +++---- 3 files changed, 28 insertions(+), 17 deletions(-) diff --git a/aioupnp/__main__.py b/aioupnp/__main__.py index 9aba6cb..990d714 100644 --- a/aioupnp/__main__.py +++ b/aioupnp/__main__.py @@ -1,7 +1,6 @@ import logging import sys from aioupnp.upnp import UPnP -from aioupnp.constants import UPNP_ORG_IGD log = logging.getLogger("aioupnp") @@ -47,7 +46,7 @@ def main(): 'gateway_address': '', 'lan_address': '', 'timeout': 1, - 'service': UPNP_ORG_IGD, + 'service': '', # if not provided try all of them 'return_as_json': True } diff --git a/aioupnp/protocols/ssdp.py b/aioupnp/protocols/ssdp.py index a52296f..4e99e47 100644 --- a/aioupnp/protocols/ssdp.py +++ b/aioupnp/protocols/ssdp.py @@ -8,7 +8,8 @@ from asyncio.futures import Future from asyncio.transports import DatagramTransport from aioupnp.fault import UPnPError from aioupnp.serialization.ssdp import SSDPDatagram -from aioupnp.constants import UPNP_ORG_IGD, SSDP_IP_ADDRESS, SSDP_PORT, SSDP_DISCOVER, SSDP_ROOT_DEVICE +from aioupnp.constants import UPNP_ORG_IGD, WIFI_ALLIANCE_ORG_IGD +from aioupnp.constants import SSDP_IP_ADDRESS, SSDP_PORT, SSDP_DISCOVER, SSDP_ROOT_DEVICE from aioupnp.protocols.multicast import MulticastProtocol ADDRESS_REGEX = re.compile("^http:\/\/(\d+\.\d+\.\d+\.\d+)\:(\d*)(\/[\w|\/|\:|\-|\.]*)$") @@ -32,21 +33,33 @@ class SSDPProtocol(MulticastProtocol): log.debug("sending packet to %s:%i: %s", address, SSDP_PORT, packet) self.transport.sendto(packet.encode().encode(), (address, SSDP_PORT)) - async def m_search(self, address, timeout: int = 1, service=UPNP_ORG_IGD) -> SSDPDatagram: + async def m_search(self, address, timeout: int = 1, service='') -> SSDPDatagram: if (address, service) in self.discover_callbacks: return self.discover_callbacks[(address, service)] - # D-Link works with both + if not service: + services = [UPNP_ORG_IGD, WIFI_ALLIANCE_ORG_IGD] + else: + services = [service] - # Cisco only works with quotes - self.send_m_search_packet(service, address, '\"%s\"' % SSDP_DISCOVER) + search_futs: List[Future] = [] + outer_fut: Future = Future() - # DD-WRT only works without quotes - self.send_m_search_packet(service, address, SSDP_DISCOVER) + for service in services: + # D-Link works with both - f: Future = Future() - self.discover_callbacks[(address, service)] = f - return await asyncio.wait_for(f, timeout) + # Cisco only works with quotes + self.send_m_search_packet(service, address, '\"%s\"' % SSDP_DISCOVER) + + # DD-WRT only works without quotes + self.send_m_search_packet(service, address, SSDP_DISCOVER) + + f: Future = Future() + f.add_done_callback(lambda _f: outer_fut.set_result(_f.result())) + self.discover_callbacks[(address, service)] = f + search_futs.append(f) + + return await asyncio.wait_for(outer_fut, timeout) def datagram_received(self, data, addr) -> None: if addr[0] == self.lan_address: @@ -106,7 +119,7 @@ async def listen_ssdp(lan_address: str, gateway_address: str, async def m_search(lan_address: str, gateway_address: str, timeout: int = 1, - service: str = UPNP_ORG_IGD, ssdp_socket: socket.socket = None) -> SSDPDatagram: + service: str = '', ssdp_socket: socket.socket = None) -> SSDPDatagram: transport, protocol, gateway_address, lan_address = await listen_ssdp( lan_address, gateway_address, ssdp_socket ) diff --git a/aioupnp/upnp.py b/aioupnp/upnp.py index 3e0b85b..5458706 100644 --- a/aioupnp/upnp.py +++ b/aioupnp/upnp.py @@ -6,7 +6,6 @@ import asyncio from typing import Tuple, Dict, List, Union from aioupnp.fault import UPnPError from aioupnp.gateway import Gateway -from aioupnp.constants import UPNP_ORG_IGD from aioupnp.util import get_gateway_and_lan_addresses from aioupnp.protocols.ssdp import m_search @@ -43,7 +42,7 @@ class UPnP: @classmethod async def discover(cls, lan_address: str = '', gateway_address: str = '', timeout: int = 1, - service: str = UPNP_ORG_IGD, interface_name: str = 'default', + service: str = '', interface_name: str = 'default', ssdp_socket: socket.socket = None, soap_socket: socket.socket = None): try: lan_address, gateway_address = cls.get_lan_and_gateway(lan_address, gateway_address, interface_name) @@ -55,7 +54,7 @@ class UPnP: @classmethod @cli async def m_search(cls, lan_address: str = '', gateway_address: str = '', timeout: int = 1, - service: str = UPNP_ORG_IGD, interface_name: str = 'default') -> Dict: + service: str = '', interface_name: str = 'default') -> Dict: lan_address, gateway_address = cls.get_lan_and_gateway(lan_address, gateway_address, interface_name) datagram = await m_search(lan_address, gateway_address, timeout, service) return { @@ -215,7 +214,7 @@ class UPnP: @classmethod def run_cli(cls, method, lan_address: str = '', gateway_address: str = '', timeout: int = 60, - service: str = UPNP_ORG_IGD, interface_name: str = 'default', + service: str = '', interface_name: str = 'default', kwargs: dict = None) -> None: kwargs = kwargs or {}