Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.11/site-packages/selkies/ice/mdns.py: 25%

1# This Source Code Form is subject to the terms of the Mozilla Public 

2# License, v. 2.0. If a copy of the MPL was not distributed with this 

3# file, You can obtain one at https://mozilla.org/MPL/2.0/. 

4# 

5# This file incorporates work covered by the following copyright and 

6# permission notice: 

7# 

8# Copyright (c) 2018-2019 Jeremy Lainé. 

9# All rights reserved. 

10# 

11# Redistribution and use in source and binary forms, with or without 

12# modification, are permitted provided that the following conditions are met: 

13# 

14# * Redistributions of source code must retain the above copyright notice, 

15# this list of conditions and the following disclaimer. 

16# * Redistributions in binary form must reproduce the above copyright notice, 

17# this list of conditions and the following disclaimer in the documentation 

18# and/or other materials provided with the distribution. 

19# * Neither the name of aioice nor the names of its contributors may 

20# be used to endorse or promote products derived from this software without 

21# specific prior written permission. 

22# 

23# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND 

24# ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED 

25# WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE 

26# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE 

27# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL 

28# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR 

29# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER 

30# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, 

31# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE 

32# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. 

33 

34import asyncio 

35import re 

36import socket 

37import uuid 

38from typing import Optional, Union, cast 

39 

40import dns.exception 

41import dns.flags 

42import dns.message 

43import dns.name 

44import dns.rdata 

45import dns.rdataclass 

46import dns.rdataset 

47import dns.rdatatype 

48import dns.zone 

49 

50MDNS_ADDRESS = "224.0.0.251" 

51MDNS_PORT = 5353 

52MDNS_HOSTNAME_RE = re.compile(r"^[a-zA-Z0-9-]{1,63}\.local$") 

53 

54MDNS_RDCLASS = dns.rdataclass.IN | 0x8000 

55 

56 

57def create_mdns_hostname() -> str: 

58 return str(uuid.uuid4()) + ".local" 

59 

60 

61def is_mdns_hostname(name: str) -> bool: 

62 return MDNS_HOSTNAME_RE.match(name) is not None 

63 

64 

65class MDnsProtocol(asyncio.DatagramProtocol): 

66 def __init__(self, tx_transport: asyncio.DatagramTransport) -> None: 

67 self.__closed: asyncio.Future[bool] = asyncio.Future() 

68 self.zone = dns.zone.Zone("", relativize=False, rdclass=MDNS_RDCLASS) 

69 self.queries: dict[dns.name.Name, set[asyncio.Future[str]]] = {} 

70 

71 self.rx_transport: Optional[asyncio.DatagramTransport] = None 

72 self.tx_transport = tx_transport 

73 

74 def connection_lost(self, exc: Exception) -> None: # type: ignore 

75 # abort any outstanding queries 

76 for name, futures in list(self.queries.items()): 

77 for future in futures: 

78 future.set_exception(asyncio.TimeoutError) 

79 self.__closed.set_result(True) 

80 

81 def connection_made(self, transport: asyncio.BaseTransport) -> None: 

82 self.rx_transport = cast(asyncio.DatagramTransport, transport) 

83 

84 def datagram_received(self, data: Union[bytes, str], addr: tuple) -> None: 

85 # parse message 

86 try: 

87 message = dns.message.from_wire(cast(bytes, data)) 

88 except dns.exception.FormError: 

89 return 

90 

91 if isinstance(message, dns.message.QueryMessage): 

92 # answer question 

93 for question in message.question: 

94 rdtypes: list[int] = [] 

95 if question.rdtype in ( 

96 dns.rdatatype.ANY, 

97 dns.rdatatype.A, 

98 dns.rdatatype.AAAA, 

99 ): 

100 rdtypes += [dns.rdatatype.A, dns.rdatatype.AAAA] 

101 

102 response = dns.message.QueryMessage(id=0) 

103 response.flags |= dns.flags.QR 

104 response.flags |= dns.flags.AA 

105 

106 for rdtype in rdtypes: 

107 try: 

108 response.answer.append( 

109 self.zone.find_rrset(name=question.name, rdtype=rdtype) 

110 ) 

111 except KeyError: 

112 continue 

113 if response.answer: 

114 self.tx_transport.sendto( 

115 response.to_wire(), (MDNS_ADDRESS, MDNS_PORT) 

116 ) 

117 

118 # handle answer 

119 for answer in message.answer: 

120 for item in answer: 

121 item = item.to_generic() 

122 if ( 

123 isinstance(item, dns.rdata.GenericRdata) 

124 and item.rdclass == MDNS_RDCLASS 

125 and item.rdtype in (dns.rdatatype.A, dns.rdatatype.AAAA) 

126 ): 

127 if item.rdtype == dns.rdatatype.A: 

128 result = socket.inet_ntop(socket.AF_INET, item.data) 

129 else: 

130 result = socket.inet_ntop(socket.AF_INET6, item.data) 

131 for future in self.queries.pop(answer.name, []): 

132 future.set_result(result) 

133 

134 # custom 

135 

136 async def close(self) -> None: 

137 self.rx_transport.close() 

138 self.tx_transport.close() 

139 await self.__closed 

140 

141 async def publish(self, hostname: str, addr: str) -> None: 

142 name = dns.name.from_text(hostname) 

143 try: 

144 data = socket.inet_pton(socket.AF_INET, addr) 

145 rdtype = dns.rdatatype.A 

146 except OSError: 

147 data = socket.inet_pton(socket.AF_INET6, addr) 

148 rdtype = dns.rdatatype.AAAA 

149 

150 rdata = dns.rdata.GenericRdata(rdclass=MDNS_RDCLASS, rdtype=rdtype, data=data) 

151 self.zone.replace_rdataset(name, dns.rdataset.from_rdata(120, rdata)) 

152 

153 async def resolve( 

154 self, hostname: str, timeout: Optional[float] = 1.0 

155 ) -> Optional[str]: 

156 name = dns.name.from_text(hostname) 

157 future: asyncio.Future[str] = asyncio.Future() 

158 

159 if name in self.queries: 

160 # a query for this name is already pending 

161 self.queries[name].add(future) 

162 else: 

163 # no query for this name is pending, send a request 

164 self.queries[name] = set([future]) 

165 message = dns.message.make_query(name, rdtype=dns.rdatatype.A) 

166 message.id = 0 

167 message.flags = 0 

168 self.tx_transport.sendto(message.to_wire(), (MDNS_ADDRESS, MDNS_PORT)) 

169 

170 try: 

171 return await asyncio.wait_for(future, timeout=timeout) 

172 except asyncio.TimeoutError: 

173 return None 

174 finally: 

175 if name in self.queries: 

176 self.queries[name].discard(future) 

177 if not self.queries[name]: 

178 del self.queries[name] 

179 

180 

181async def create_mdns_protocol() -> MDnsProtocol: 

182 """ 

183 Using a single socket works fine on Linux, but on OS X we need to use 

184 separate sockets for sending and receiving. 

185 """ 

186 loop = asyncio.get_running_loop() 

187 

188 # sender 

189 

190 tx_sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM) 

191 tx_sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) 

192 if hasattr(socket, "SO_REUSEPORT"): 

193 tx_sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEPORT, 1) 

194 tx_sock.bind(("", MDNS_PORT)) 

195 

196 tx_transport, _ = await loop.create_datagram_endpoint( 

197 lambda: asyncio.DatagramProtocol(), 

198 sock=tx_sock, 

199 ) 

200 

201 # receiver 

202 

203 rx_sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM) 

204 rx_sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) 

205 if hasattr(socket, "SO_REUSEPORT"): 

206 rx_sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEPORT, 1) 

207 rx_sock.setsockopt( 

208 socket.IPPROTO_IP, 

209 socket.IP_ADD_MEMBERSHIP, 

210 socket.inet_aton(MDNS_ADDRESS) + b"\x00\x00\x00\x00", 

211 ) 

212 rx_sock.bind((MDNS_ADDRESS, MDNS_PORT)) 

213 

214 _, protocol = await loop.create_datagram_endpoint( 

215 lambda: MDnsProtocol(tx_transport=tx_transport), 

216 sock=rx_sock, 

217 ) 

218 

219 return protocol