You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

155 lines
4.7 KiB

4 years ago
  1. # -*- coding: utf-8 -*-
  2. """UPnP discovery via Simple Service Discovery Protocol (SSDP)."""
  3. import asyncio
  4. import logging
  5. import socket
  6. from datetime import datetime
  7. from ipaddress import IPv4Address
  8. _LOGGER = logging.getLogger(__name__)
  9. _LOGGER_TRAFFIC = logging.getLogger("async_upnp_client.traffic")
  10. SSDP_TARGET = ('239.255.255.250', 1900)
  11. SSDP_ST_ALL = 'ssdp:all'
  12. SSDP_ST_ROOTDEVICE = 'upnp:rootdevice'
  13. SSDP_MX = 4
  14. RECV_SIZE = 32678
  15. def _discovery_message(ssdp_target: str, ssdp_mx: int, ssdp_st: str):
  16. """Construct a SSDP packet."""
  17. return 'M-SEARCH * HTTP/1.1\r\n' \
  18. 'HOST:{target}:{port}\r\n' \
  19. 'MAN:"ssdp:discover"\r\n' \
  20. 'MX:{mx}\r\n' \
  21. 'ST:{st}\r\n' \
  22. '\r\n'.format(target=ssdp_target[0], port=ssdp_target[1],
  23. mx=ssdp_mx, st=ssdp_st).encode()
  24. def _parse_response(data):
  25. headers = {}
  26. for line in data.splitlines():
  27. decoded = line.decode()
  28. if ':' in decoded:
  29. key, value = decoded.split(':', 1)
  30. key = key.strip().lower()
  31. headers[key] = value.strip()
  32. # custom data
  33. headers['_timestamp'] = datetime.now()
  34. return headers
  35. def discover(timeout: int = SSDP_MX,
  36. service_type: str = SSDP_ST_ALL,
  37. source_ip: IPv4Address = None):
  38. """Discover devices via SSDP."""
  39. # create socket
  40. sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM, socket.IPPROTO_UDP)
  41. sock.setsockopt(socket.IPPROTO_IP, socket.IP_MULTICAST_TTL, timeout)
  42. sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
  43. sock.settimeout(timeout)
  44. if source_ip:
  45. sock.bind((source_ip, 0))
  46. # send discovery packet
  47. message = _discovery_message(SSDP_TARGET, timeout, service_type)
  48. _LOGGER.debug('Sending discovery message')
  49. _LOGGER_TRAFFIC.debug('Sending message:\n%s', message)
  50. sock.sendto(message, SSDP_TARGET)
  51. # handle replies
  52. responses = []
  53. try:
  54. while True:
  55. # parse response
  56. data, _ = sock.recvfrom(RECV_SIZE)
  57. _LOGGER_TRAFFIC.debug('Received packet:\n%s', data)
  58. response = _parse_response(data)
  59. _LOGGER.debug('Received response: %s', response)
  60. if response not in responses:
  61. responses.append(response)
  62. except socket.timeout:
  63. pass
  64. sock.close()
  65. return responses
  66. async def async_discover(timeout: int = SSDP_MX,
  67. service_type: str = SSDP_ST_ALL,
  68. source_ip: IPv4Address = None,
  69. async_callback=None,
  70. loop=None):
  71. """Discover devices via SSDP."""
  72. loop = loop or asyncio.get_event_loop()
  73. # create socket
  74. sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM, socket.IPPROTO_UDP)
  75. sock.setsockopt(socket.IPPROTO_IP, socket.IP_MULTICAST_TTL, timeout)
  76. # sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
  77. if source_ip:
  78. sock.bind((source_ip, 0))
  79. # create protocol and send discovery packet
  80. connect = loop.create_datagram_endpoint(
  81. lambda: SsdpDiscoveryProtocol(loop, timeout, service_type, async_callback),
  82. sock=sock,
  83. )
  84. transport, protocol = await connect
  85. # wait for devices to respond
  86. await asyncio.sleep(timeout)
  87. transport.close()
  88. return protocol.responses
  89. class SsdpDiscoveryProtocol:
  90. """SSDP Discovery Protocol."""
  91. def __init__(self, loop, timeout, service_type, async_callback):
  92. """Initializer."""
  93. self.loop = loop
  94. self.timeout = timeout
  95. self.service_type = service_type
  96. self.async_callback = async_callback
  97. self.on_con_lost = loop.create_future()
  98. self.transport = None
  99. self.responses = []
  100. def connection_made(self, transport):
  101. """Handle connection made."""
  102. self.transport = transport
  103. message = _discovery_message(SSDP_TARGET, self.timeout, self.service_type)
  104. self.transport.sendto(message, SSDP_TARGET)
  105. def datagram_received(self, data, addr):
  106. """Handle a discovery-response."""
  107. _LOGGER_TRAFFIC.debug('Received packet from %s:\n%s', addr, data)
  108. response = _parse_response(data)
  109. _LOGGER.debug('Received response: %s', response)
  110. if response not in self.responses:
  111. self.responses.append(response)
  112. if self.async_callback:
  113. callback = self.async_callback(response)
  114. self.loop.create_task(callback)
  115. def error_received(self, exc):
  116. """Handle an error."""
  117. # pylint: disable=no-self-use
  118. _LOGGER.error('Received error: %s', exc)
  119. def connection_lost(self, exc):
  120. """Handle connection lost."""