Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.11/site-packages/dns/renderer.py: 28%

1# Copyright (C) Dnspython Contributors, see LICENSE for text of ISC license 

2 

3# Copyright (C) 2001-2017 Nominum, Inc. 

4# 

5# Permission to use, copy, modify, and distribute this software and its 

6# documentation for any purpose with or without fee is hereby granted, 

7# provided that the above copyright notice and this permission notice 

8# appear in all copies. 

9# 

10# THE SOFTWARE IS PROVIDED "AS IS" AND NOMINUM DISCLAIMS ALL WARRANTIES 

11# WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF 

12# MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL NOMINUM BE LIABLE FOR 

13# ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES 

14# WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN 

15# ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT 

16# OF OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE. 

17 

18"""Help for building DNS wire format messages""" 

19 

20import contextlib 

21import io 

22import random 

23import struct 

24import time 

25 

26import dns.edns 

27import dns.exception 

28import dns.name 

29import dns.rdataclass 

30import dns.rdatatype 

31import dns.rdtypes.ANY.OPT 

32import dns.rdtypes.ANY.TSIG 

33import dns.rrset 

34import dns.tsig 

35from dns._render_util import prefixed_length as prefixed_length # pyright: ignore 

36 

37DEFAULT_EDNS_PAYLOAD = 1232 

38 

39QUESTION = 0 

40ANSWER = 1 

41AUTHORITY = 2 

42ADDITIONAL = 3 

43 

44 

45def _make_opt(flags=0, payload=DEFAULT_EDNS_PAYLOAD, options=None): 

46 opt = dns.rdtypes.ANY.OPT.OPT(payload, dns.rdatatype.OPT, options or ()) 

47 return dns.rrset.from_rdata(dns.name.root, int(flags), opt) 

48 

49 

50def _make_tsig(keyname, algorithm, time_signed, fudge, mac, original_id, error, other): 

51 tsig = dns.rdtypes.ANY.TSIG.TSIG( 

52 dns.rdataclass.ANY, 

53 dns.rdatatype.TSIG, 

54 algorithm, 

55 time_signed, 

56 fudge, 

57 mac, 

58 original_id, 

59 error, 

60 other, 

61 ) 

62 return dns.rrset.from_rdata(keyname, 0, tsig) 

63 

64 

65class Renderer: 

66 """Helper class for building DNS wire-format messages. 

67 

68 Most applications can use the higher-level L{dns.message.Message} 

69 class and its to_wire() method to generate wire-format messages. 

70 This class is for those applications which need finer control 

71 over the generation of messages. 

72 

73 Typical use:: 

74 

75 r = dns.renderer.Renderer(id=1, flags=0x80, max_size=512) 

76 r.add_question(qname, qtype, qclass) 

77 r.add_rrset(dns.renderer.ANSWER, rrset_1) 

78 r.add_rrset(dns.renderer.ANSWER, rrset_2) 

79 r.add_rrset(dns.renderer.AUTHORITY, ns_rrset) 

80 r.add_rrset(dns.renderer.ADDITIONAL, ad_rrset_1) 

81 r.add_rrset(dns.renderer.ADDITIONAL, ad_rrset_2) 

82 r.add_edns(0, 0, 4096) 

83 r.write_header() 

84 r.add_tsig(keyname, secret, 300, 1, 0, '', request_mac) 

85 wire = r.get_wire() 

86 

87 If padding is going to be used, then the OPT record MUST be 

88 written after everything else in the additional section except for 

89 the TSIG (if any). 

90 

91 output, an io.BytesIO, where rendering is written 

92 

93 id: the message id 

94 

95 flags: the message flags 

96 

97 max_size: the maximum size of the message 

98 

99 origin: the origin to use when rendering relative names 

100 

101 compress: the compression table 

102 

103 section: an int, the section currently being rendered 

104 

105 counts: list of the number of RRs in each section 

106 

107 mac: the MAC of the rendered message (if TSIG was used) 

108 """ 

109 

110 def __init__(self, id=None, flags=0, max_size=65535, origin=None): 

111 """Initialize a new renderer.""" 

112 

113 self.output = io.BytesIO() 

114 if id is None: 

115 self.id = random.randint(0, 65535) 

116 else: 

117 self.id = id 

118 self.flags = flags 

119 self.max_size = max_size 

120 self.origin = origin 

121 self.compress = {} 

122 self.section = QUESTION 

123 self.counts = [0, 0, 0, 0] 

124 self.output.write(b"\x00" * 12) 

125 self.mac = "" 

126 self.reserved = 0 

127 self.was_padded = False 

128 

129 def _rollback(self, where): 

130 """Truncate the output buffer at offset *where*, and remove any 

131 compression table entries that pointed beyond the truncation 

132 point. 

133 """ 

134 

135 self.output.seek(where) 

136 self.output.truncate() 

137 keys_to_delete = [] 

138 for k, v in self.compress.items(): 

139 if v >= where: 

140 keys_to_delete.append(k) 

141 for k in keys_to_delete: 

142 del self.compress[k] 

143 

144 def _set_section(self, section): 

145 """Set the renderer's current section. 

146 

147 Sections must be rendered order: QUESTION, ANSWER, AUTHORITY, 

148 ADDITIONAL. Sections may be empty. 

149 

150 Raises dns.exception.FormError if an attempt was made to set 

151 a section value less than the current section. 

152 """ 

153 

154 if self.section != section: 

155 if self.section > section: 

156 raise dns.exception.FormError 

157 self.section = section 

158 

159 @contextlib.contextmanager 

160 def _track_size(self): 

161 start = self.output.tell() 

162 yield start 

163 if self.output.tell() > self.max_size: 

164 self._rollback(start) 

165 raise dns.exception.TooBig 

166 

167 @contextlib.contextmanager 

168 def _temporarily_seek_to(self, where): 

169 current = self.output.tell() 

170 try: 

171 self.output.seek(where) 

172 yield 

173 finally: 

174 self.output.seek(current) 

175 

176 def add_question(self, qname, rdtype, rdclass=dns.rdataclass.IN): 

177 """Add a question to the message.""" 

178 

179 self._set_section(QUESTION) 

180 with self._track_size(): 

181 qname.to_wire(self.output, self.compress, self.origin) 

182 self.output.write(struct.pack("!HH", rdtype, rdclass)) 

183 self.counts[QUESTION] += 1 

184 

185 def add_rrset(self, section, rrset, **kw): 

186 """Add the rrset to the specified section. 

187 

188 Any keyword arguments are passed on to the rdataset's to_wire() 

189 routine. 

190 """ 

191 

192 self._set_section(section) 

193 with self._track_size(): 

194 n = rrset.to_wire(self.output, self.compress, self.origin, **kw) 

195 self.counts[section] += n 

196 

197 def add_rdataset(self, section, name, rdataset, **kw): 

198 """Add the rdataset to the specified section, using the specified 

199 name as the owner name. 

200 

201 Any keyword arguments are passed on to the rdataset's to_wire() 

202 routine. 

203 """ 

204 

205 self._set_section(section) 

206 with self._track_size(): 

207 n = rdataset.to_wire(name, self.output, self.compress, self.origin, **kw) 

208 self.counts[section] += n 

209 

210 def add_opt(self, opt, pad=0, opt_size=0, tsig_size=0): 

211 """Add *opt* to the additional section, applying padding if desired. The 

212 padding will take the specified precomputed OPT size and TSIG size into 

213 account. 

214 

215 Note that we don't have reliable way of knowing how big a GSS-TSIG digest 

216 might be, so we we might not get an even multiple of the pad in that case.""" 

217 if pad: 

218 ttl = opt.ttl 

219 assert opt_size >= 11 

220 opt_rdata = opt[0] 

221 size_without_padding = self.output.tell() + opt_size + tsig_size 

222 remainder = size_without_padding % pad 

223 if remainder: 

224 pad = b"\x00" * (pad - remainder) 

225 else: 

226 pad = b"" 

227 options = list(opt_rdata.options) 

228 options.append(dns.edns.GenericOption(dns.edns.OptionType.PADDING, pad)) 

229 opt = _make_opt(ttl, opt_rdata.rdclass, options) # pyright: ignore 

230 self.was_padded = True 

231 self.add_rrset(ADDITIONAL, opt) 

232 

233 def add_edns(self, edns, ednsflags, payload, options=None): 

234 """Add an EDNS OPT record to the message.""" 

235 

236 # make sure the EDNS version in ednsflags agrees with edns 

237 ednsflags &= 0xFF00FFFF 

238 ednsflags |= edns << 16 

239 opt = _make_opt(ednsflags, payload, options) # pyright: ignore 

240 self.add_opt(opt) 

241 

242 def add_tsig( 

243 self, 

244 keyname, 

245 secret, 

246 fudge, 

247 id, 

248 tsig_error, 

249 other_data, 

250 request_mac, 

251 algorithm=dns.tsig.default_algorithm, 

252 ): 

253 """Add a TSIG signature to the message.""" 

254 

255 s = self.output.getvalue() 

256 

257 if isinstance(secret, dns.tsig.Key): 

258 key = secret 

259 else: 

260 key = dns.tsig.Key(keyname, secret, algorithm) 

261 tsig = _make_tsig( # pyright: ignore 

262 keyname, algorithm, 0, fudge, b"", id, tsig_error, other_data 

263 ) 

264 tsig, _ = dns.tsig.sign(s, key, tsig[0], int(time.time()), request_mac) 

265 self._write_tsig(tsig, keyname) 

266 

267 def add_multi_tsig( 

268 self, 

269 ctx, 

270 keyname, 

271 secret, 

272 fudge, 

273 id, 

274 tsig_error, 

275 other_data, 

276 request_mac, 

277 algorithm=dns.tsig.default_algorithm, 

278 ): 

279 """Add a TSIG signature to the message. Unlike add_tsig(), this can be 

280 used for a series of consecutive DNS envelopes, e.g. for a zone 

281 transfer over TCP [RFC2845, 4.4]. 

282 

283 For the first message in the sequence, give ctx=None. For each 

284 subsequent message, give the ctx that was returned from the 

285 add_multi_tsig() call for the previous message.""" 

286 

287 s = self.output.getvalue() 

288 

289 if isinstance(secret, dns.tsig.Key): 

290 key = secret 

291 else: 

292 key = dns.tsig.Key(keyname, secret, algorithm) 

293 tsig = _make_tsig( # pyright: ignore 

294 keyname, algorithm, 0, fudge, b"", id, tsig_error, other_data 

295 ) 

296 tsig, ctx = dns.tsig.sign( 

297 s, key, tsig[0], int(time.time()), request_mac, ctx, True 

298 ) 

299 self._write_tsig(tsig, keyname) 

300 return ctx 

301 

302 def _write_tsig(self, tsig, keyname): 

303 if self.was_padded: 

304 compress = None 

305 else: 

306 compress = self.compress 

307 self._set_section(ADDITIONAL) 

308 with self._track_size(): 

309 keyname.to_wire(self.output, compress, self.origin) 

310 self.output.write( 

311 struct.pack("!HHI", dns.rdatatype.TSIG, dns.rdataclass.ANY, 0) 

312 ) 

313 with prefixed_length(self.output, 2): 

314 tsig.to_wire(self.output) 

315 

316 self.counts[ADDITIONAL] += 1 

317 with self._temporarily_seek_to(10): 

318 self.output.write(struct.pack("!H", self.counts[ADDITIONAL])) 

319 

320 def write_header(self): 

321 """Write the DNS message header. 

322 

323 Writing the DNS message header is done after all sections 

324 have been rendered, but before the optional TSIG signature 

325 is added. 

326 """ 

327 

328 with self._temporarily_seek_to(0): 

329 self.output.write( 

330 struct.pack( 

331 "!HHHHHH", 

332 self.id, 

333 self.flags, 

334 self.counts[0], 

335 self.counts[1], 

336 self.counts[2], 

337 self.counts[3], 

338 ) 

339 ) 

340 

341 def get_wire(self): 

342 """Return the wire format message.""" 

343 

344 return self.output.getvalue() 

345 

346 def reserve(self, size: int) -> None: 

347 """Reserve *size* bytes.""" 

348 if size < 0: 

349 raise ValueError("reserved amount must be non-negative") 

350 if size > self.max_size: 

351 raise ValueError("cannot reserve more than the maximum size") 

352 self.reserved += size 

353 self.max_size -= size 

354 

355 def release_reserved(self) -> None: 

356 """Release the reserved bytes.""" 

357 self.max_size += self.reserved 

358 self.reserved = 0