Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.11/site-packages/aiodns/compat.py: 60%

1""" 

2Compatibility layer for pycares 5.x API. 

3 

4This module provides result types compatible with pycares 4.x API 

5to maintain backward compatibility with existing code. 

6""" 

7 

8from __future__ import annotations 

9 

10from dataclasses import dataclass 

11from typing import Union, cast 

12 

13import pycares 

14 

15 

16def _maybe_str(data: bytes) -> str | bytes: 

17 """Decode bytes as ASCII, return bytes if decode fails (pycares 4.x).""" 

18 try: 

19 return data.decode('ascii') 

20 except UnicodeDecodeError: 

21 return data 

22 

23 

24@dataclass(frozen=True, slots=True) 

25class AresQueryAResult: 

26 """A record result (compatible with pycares 4.x ares_query_a_result).""" 

27 

28 host: str 

29 ttl: int 

30 

31 

32@dataclass(frozen=True, slots=True) 

33class AresQueryAAAAResult: 

34 """AAAA record result (pycares 4.x compat).""" 

35 

36 host: str 

37 ttl: int 

38 

39 

40@dataclass(frozen=True, slots=True) 

41class AresQueryCNAMEResult: 

42 """CNAME record result (pycares 4.x compat).""" 

43 

44 cname: str 

45 ttl: int 

46 

47 

48@dataclass(frozen=True, slots=True) 

49class AresQueryMXResult: 

50 """MX record result (pycares 4.x compat).""" 

51 

52 host: str 

53 priority: int 

54 ttl: int 

55 

56 

57@dataclass(frozen=True, slots=True) 

58class AresQueryNSResult: 

59 """NS record result (pycares 4.x compat).""" 

60 

61 host: str 

62 ttl: int 

63 

64 

65@dataclass(frozen=True, slots=True) 

66class AresQueryTXTResult: 

67 """TXT record result (pycares 4.x compat).""" 

68 

69 text: str | bytes # str if ASCII, bytes otherwise (pycares 4.x behavior) 

70 ttl: int 

71 

72 

73@dataclass(frozen=True, slots=True) 

74class AresQuerySOAResult: 

75 """SOA record result (pycares 4.x compat).""" 

76 

77 nsname: str 

78 hostmaster: str 

79 serial: int 

80 refresh: int 

81 retry: int 

82 expires: int 

83 minttl: int 

84 ttl: int 

85 

86 

87@dataclass(frozen=True, slots=True) 

88class AresQuerySRVResult: 

89 """SRV record result (pycares 4.x compat).""" 

90 

91 host: str 

92 port: int 

93 priority: int 

94 weight: int 

95 ttl: int 

96 

97 

98@dataclass(frozen=True, slots=True) 

99class AresQueryNAPTRResult: 

100 """NAPTR record result (pycares 4.x compat).""" 

101 

102 order: int 

103 preference: int 

104 flags: str 

105 service: str 

106 regex: str 

107 replacement: str 

108 ttl: int 

109 

110 

111@dataclass(frozen=True, slots=True) 

112class AresQueryCAAResult: 

113 """CAA record result (pycares 4.x compat).""" 

114 

115 critical: int 

116 property: str 

117 value: str 

118 ttl: int 

119 

120 

121@dataclass(frozen=True, slots=True) 

122class AresQueryPTRResult: 

123 """PTR record result (pycares 4.x compat).""" 

124 

125 name: str 

126 ttl: int 

127 aliases: list[str] 

128 

129 

130@dataclass(frozen=True, slots=True) 

131class AresHostResult: 

132 """Host result (compatible with pycares 4.x ares_host_result).""" 

133 

134 name: str 

135 aliases: list[str] 

136 addresses: list[str] 

137 

138 

139# Type alias for a single converted record 

140ConvertedRecord = Union[ 

141 AresQueryAResult, 

142 AresQueryAAAAResult, 

143 AresQueryCNAMEResult, 

144 AresQueryMXResult, 

145 AresQueryNSResult, 

146 AresQueryTXTResult, 

147 AresQuerySOAResult, 

148 AresQuerySRVResult, 

149 AresQueryNAPTRResult, 

150 AresQueryCAAResult, 

151 AresQueryPTRResult, 

152 pycares.DNSRecord, # Unknown types returned as-is 

153] 

154 

155# Type alias for query results 

156QueryResult = Union[ 

157 list[AresQueryAResult], 

158 list[AresQueryAAAAResult], 

159 AresQueryCNAMEResult, 

160 list[AresQueryMXResult], 

161 list[AresQueryNSResult], 

162 list[AresQueryTXTResult], 

163 AresQuerySOAResult, 

164 list[AresQuerySRVResult], 

165 list[AresQueryNAPTRResult], 

166 list[AresQueryCAAResult], 

167 AresQueryPTRResult, 

168 list[ConvertedRecord], # For ANY query type 

169] 

170 

171 

172def _convert_record(record: pycares.DNSRecord) -> ConvertedRecord: 

173 """Convert a single DNS record to pycares 4.x compatible format.""" 

174 ttl = record.ttl 

175 record_type = record.type 

176 

177 if record_type == pycares.QUERY_TYPE_A: 

178 a_data = cast(pycares.ARecordData, record.data) 

179 return AresQueryAResult(host=a_data.addr, ttl=ttl) 

180 if record_type == pycares.QUERY_TYPE_AAAA: 

181 aaaa_data = cast(pycares.AAAARecordData, record.data) 

182 return AresQueryAAAAResult(host=aaaa_data.addr, ttl=ttl) 

183 if record_type == pycares.QUERY_TYPE_CNAME: 

184 cname_data = cast(pycares.CNAMERecordData, record.data) 

185 return AresQueryCNAMEResult(cname=cname_data.cname, ttl=ttl) 

186 if record_type == pycares.QUERY_TYPE_MX: 

187 mx_data = cast(pycares.MXRecordData, record.data) 

188 return AresQueryMXResult( 

189 host=mx_data.exchange, priority=mx_data.priority, ttl=ttl 

190 ) 

191 if record_type == pycares.QUERY_TYPE_NS: 

192 ns_data = cast(pycares.NSRecordData, record.data) 

193 return AresQueryNSResult(host=ns_data.nsdname, ttl=ttl) 

194 if record_type == pycares.QUERY_TYPE_TXT: 

195 txt_data = cast(pycares.TXTRecordData, record.data) 

196 return AresQueryTXTResult(text=_maybe_str(txt_data.data), ttl=ttl) 

197 if record_type == pycares.QUERY_TYPE_SOA: 

198 soa_data = cast(pycares.SOARecordData, record.data) 

199 return AresQuerySOAResult( 

200 nsname=soa_data.mname, 

201 hostmaster=soa_data.rname, 

202 serial=soa_data.serial, 

203 refresh=soa_data.refresh, 

204 retry=soa_data.retry, 

205 expires=soa_data.expire, 

206 minttl=soa_data.minimum, 

207 ttl=ttl, 

208 ) 

209 if record_type == pycares.QUERY_TYPE_SRV: 

210 srv_data = cast(pycares.SRVRecordData, record.data) 

211 return AresQuerySRVResult( 

212 host=srv_data.target, 

213 port=srv_data.port, 

214 priority=srv_data.priority, 

215 weight=srv_data.weight, 

216 ttl=ttl, 

217 ) 

218 if record_type == pycares.QUERY_TYPE_NAPTR: 

219 naptr_data = cast(pycares.NAPTRRecordData, record.data) 

220 return AresQueryNAPTRResult( 

221 order=naptr_data.order, 

222 preference=naptr_data.preference, 

223 flags=naptr_data.flags, 

224 service=naptr_data.service, 

225 regex=naptr_data.regexp, 

226 replacement=naptr_data.replacement, 

227 ttl=ttl, 

228 ) 

229 if record_type == pycares.QUERY_TYPE_CAA: 

230 caa_data = cast(pycares.CAARecordData, record.data) 

231 return AresQueryCAAResult( 

232 critical=caa_data.critical, 

233 property=caa_data.tag, 

234 value=caa_data.value, 

235 ttl=ttl, 

236 ) 

237 if record_type == pycares.QUERY_TYPE_PTR: 

238 ptr_data = cast(pycares.PTRRecordData, record.data) 

239 return AresQueryPTRResult(name=ptr_data.dname, ttl=ttl, aliases=[]) 

240 # Return raw record for unknown types 

241 return record 

242 

243 

244def convert_result(dns_result: pycares.DNSResult, qtype: int) -> QueryResult: 

245 """Convert pycares 5.x DNSResult to pycares 4.x compatible format.""" 

246 # For ANY - convert all records and return mixed list 

247 if qtype == pycares.QUERY_TYPE_ANY: 

248 return [_convert_record(record) for record in dns_result.answer] 

249 

250 results: list[ConvertedRecord] = [] 

251 

252 for record in dns_result.answer: 

253 record_type = record.type 

254 

255 # Filter by query type since answer can contain other types 

256 # (e.g., CNAME records when querying for A/AAAA) 

257 if record_type != qtype: 

258 continue 

259 

260 converted = _convert_record(record) 

261 

262 # CNAME, SOA, and PTR return single result, not list 

263 if record_type in ( 

264 pycares.QUERY_TYPE_CNAME, 

265 pycares.QUERY_TYPE_SOA, 

266 pycares.QUERY_TYPE_PTR, 

267 ): 

268 return cast(QueryResult, converted) 

269 

270 results.append(converted) 

271 

272 return results