renderer.py 12 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329
  1. # Copyright (C) 2001-2007, 2009-2011 Nominum, Inc.
  2. #
  3. # Permission to use, copy, modify, and distribute this software and its
  4. # documentation for any purpose with or without fee is hereby granted,
  5. # provided that the above copyright notice and this permission notice
  6. # appear in all copies.
  7. #
  8. # THE SOFTWARE IS PROVIDED "AS IS" AND NOMINUM DISCLAIMS ALL WARRANTIES
  9. # WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
  10. # MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL NOMINUM BE LIABLE FOR
  11. # ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
  12. # WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
  13. # ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT
  14. # OF OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
  15. """Help for building DNS wire format messages"""
  16. from io import BytesIO
  17. import struct
  18. import random
  19. import time
  20. import dns.exception
  21. import dns.tsig
  22. from ._compat import long
  23. QUESTION = 0
  24. ANSWER = 1
  25. AUTHORITY = 2
  26. ADDITIONAL = 3
  27. class Renderer(object):
  28. """Helper class for building DNS wire-format messages.
  29. Most applications can use the higher-level L{dns.message.Message}
  30. class and its to_wire() method to generate wire-format messages.
  31. This class is for those applications which need finer control
  32. over the generation of messages.
  33. Typical use::
  34. r = dns.renderer.Renderer(id=1, flags=0x80, max_size=512)
  35. r.add_question(qname, qtype, qclass)
  36. r.add_rrset(dns.renderer.ANSWER, rrset_1)
  37. r.add_rrset(dns.renderer.ANSWER, rrset_2)
  38. r.add_rrset(dns.renderer.AUTHORITY, ns_rrset)
  39. r.add_edns(0, 0, 4096)
  40. r.add_rrset(dns.renderer.ADDTIONAL, ad_rrset_1)
  41. r.add_rrset(dns.renderer.ADDTIONAL, ad_rrset_2)
  42. r.write_header()
  43. r.add_tsig(keyname, secret, 300, 1, 0, '', request_mac)
  44. wire = r.get_wire()
  45. @ivar output: where rendering is written
  46. @type output: BytesIO object
  47. @ivar id: the message id
  48. @type id: int
  49. @ivar flags: the message flags
  50. @type flags: int
  51. @ivar max_size: the maximum size of the message
  52. @type max_size: int
  53. @ivar origin: the origin to use when rendering relative names
  54. @type origin: dns.name.Name object
  55. @ivar compress: the compression table
  56. @type compress: dict
  57. @ivar section: the section currently being rendered
  58. @type section: int (dns.renderer.QUESTION, dns.renderer.ANSWER,
  59. dns.renderer.AUTHORITY, or dns.renderer.ADDITIONAL)
  60. @ivar counts: list of the number of RRs in each section
  61. @type counts: int list of length 4
  62. @ivar mac: the MAC of the rendered message (if TSIG was used)
  63. @type mac: string
  64. """
  65. def __init__(self, id=None, flags=0, max_size=65535, origin=None):
  66. """Initialize a new renderer.
  67. @param id: the message id
  68. @type id: int
  69. @param flags: the DNS message flags
  70. @type flags: int
  71. @param max_size: the maximum message size; the default is 65535.
  72. If rendering results in a message greater than I{max_size},
  73. then L{dns.exception.TooBig} will be raised.
  74. @type max_size: int
  75. @param origin: the origin to use when rendering relative names
  76. @type origin: dns.name.Name or None.
  77. """
  78. self.output = BytesIO()
  79. if id is None:
  80. self.id = random.randint(0, 65535)
  81. else:
  82. self.id = id
  83. self.flags = flags
  84. self.max_size = max_size
  85. self.origin = origin
  86. self.compress = {}
  87. self.section = QUESTION
  88. self.counts = [0, 0, 0, 0]
  89. self.output.write(b'\x00' * 12)
  90. self.mac = ''
  91. def _rollback(self, where):
  92. """Truncate the output buffer at offset I{where}, and remove any
  93. compression table entries that pointed beyond the truncation
  94. point.
  95. @param where: the offset
  96. @type where: int
  97. """
  98. self.output.seek(where)
  99. self.output.truncate()
  100. keys_to_delete = []
  101. for k, v in self.compress.items():
  102. if v >= where:
  103. keys_to_delete.append(k)
  104. for k in keys_to_delete:
  105. del self.compress[k]
  106. def _set_section(self, section):
  107. """Set the renderer's current section.
  108. Sections must be rendered order: QUESTION, ANSWER, AUTHORITY,
  109. ADDITIONAL. Sections may be empty.
  110. @param section: the section
  111. @type section: int
  112. @raises dns.exception.FormError: an attempt was made to set
  113. a section value less than the current section.
  114. """
  115. if self.section != section:
  116. if self.section > section:
  117. raise dns.exception.FormError
  118. self.section = section
  119. def add_question(self, qname, rdtype, rdclass=dns.rdataclass.IN):
  120. """Add a question to the message.
  121. @param qname: the question name
  122. @type qname: dns.name.Name
  123. @param rdtype: the question rdata type
  124. @type rdtype: int
  125. @param rdclass: the question rdata class
  126. @type rdclass: int
  127. """
  128. self._set_section(QUESTION)
  129. before = self.output.tell()
  130. qname.to_wire(self.output, self.compress, self.origin)
  131. self.output.write(struct.pack("!HH", rdtype, rdclass))
  132. after = self.output.tell()
  133. if after >= self.max_size:
  134. self._rollback(before)
  135. raise dns.exception.TooBig
  136. self.counts[QUESTION] += 1
  137. def add_rrset(self, section, rrset, **kw):
  138. """Add the rrset to the specified section.
  139. Any keyword arguments are passed on to the rdataset's to_wire()
  140. routine.
  141. @param section: the section
  142. @type section: int
  143. @param rrset: the rrset
  144. @type rrset: dns.rrset.RRset object
  145. """
  146. self._set_section(section)
  147. before = self.output.tell()
  148. n = rrset.to_wire(self.output, self.compress, self.origin, **kw)
  149. after = self.output.tell()
  150. if after >= self.max_size:
  151. self._rollback(before)
  152. raise dns.exception.TooBig
  153. self.counts[section] += n
  154. def add_rdataset(self, section, name, rdataset, **kw):
  155. """Add the rdataset to the specified section, using the specified
  156. name as the owner name.
  157. Any keyword arguments are passed on to the rdataset's to_wire()
  158. routine.
  159. @param section: the section
  160. @type section: int
  161. @param name: the owner name
  162. @type name: dns.name.Name object
  163. @param rdataset: the rdataset
  164. @type rdataset: dns.rdataset.Rdataset object
  165. """
  166. self._set_section(section)
  167. before = self.output.tell()
  168. n = rdataset.to_wire(name, self.output, self.compress, self.origin,
  169. **kw)
  170. after = self.output.tell()
  171. if after >= self.max_size:
  172. self._rollback(before)
  173. raise dns.exception.TooBig
  174. self.counts[section] += n
  175. def add_edns(self, edns, ednsflags, payload, options=None):
  176. """Add an EDNS OPT record to the message.
  177. @param edns: The EDNS level to use.
  178. @type edns: int
  179. @param ednsflags: EDNS flag values.
  180. @type ednsflags: int
  181. @param payload: The EDNS sender's payload field, which is the maximum
  182. size of UDP datagram the sender can handle.
  183. @type payload: int
  184. @param options: The EDNS options list
  185. @type options: list of dns.edns.Option instances
  186. @see: RFC 2671
  187. """
  188. # make sure the EDNS version in ednsflags agrees with edns
  189. ednsflags &= long(0xFF00FFFF)
  190. ednsflags |= (edns << 16)
  191. self._set_section(ADDITIONAL)
  192. before = self.output.tell()
  193. self.output.write(struct.pack('!BHHIH', 0, dns.rdatatype.OPT, payload,
  194. ednsflags, 0))
  195. if options is not None:
  196. lstart = self.output.tell()
  197. for opt in options:
  198. stuff = struct.pack("!HH", opt.otype, 0)
  199. self.output.write(stuff)
  200. start = self.output.tell()
  201. opt.to_wire(self.output)
  202. end = self.output.tell()
  203. assert end - start < 65536
  204. self.output.seek(start - 2)
  205. stuff = struct.pack("!H", end - start)
  206. self.output.write(stuff)
  207. self.output.seek(0, 2)
  208. lend = self.output.tell()
  209. assert lend - lstart < 65536
  210. self.output.seek(lstart - 2)
  211. stuff = struct.pack("!H", lend - lstart)
  212. self.output.write(stuff)
  213. self.output.seek(0, 2)
  214. after = self.output.tell()
  215. if after >= self.max_size:
  216. self._rollback(before)
  217. raise dns.exception.TooBig
  218. self.counts[ADDITIONAL] += 1
  219. def add_tsig(self, keyname, secret, fudge, id, tsig_error, other_data,
  220. request_mac, algorithm=dns.tsig.default_algorithm):
  221. """Add a TSIG signature to the message.
  222. @param keyname: the TSIG key name
  223. @type keyname: dns.name.Name object
  224. @param secret: the secret to use
  225. @type secret: string
  226. @param fudge: TSIG time fudge
  227. @type fudge: int
  228. @param id: the message id to encode in the tsig signature
  229. @type id: int
  230. @param tsig_error: TSIG error code; default is 0.
  231. @type tsig_error: int
  232. @param other_data: TSIG other data.
  233. @type other_data: string
  234. @param request_mac: This message is a response to the request which
  235. had the specified MAC.
  236. @type request_mac: string
  237. @param algorithm: the TSIG algorithm to use
  238. @type algorithm: dns.name.Name object
  239. """
  240. self._set_section(ADDITIONAL)
  241. before = self.output.tell()
  242. s = self.output.getvalue()
  243. (tsig_rdata, self.mac, ctx) = dns.tsig.sign(s,
  244. keyname,
  245. secret,
  246. int(time.time()),
  247. fudge,
  248. id,
  249. tsig_error,
  250. other_data,
  251. request_mac,
  252. algorithm=algorithm)
  253. keyname.to_wire(self.output, self.compress, self.origin)
  254. self.output.write(struct.pack('!HHIH', dns.rdatatype.TSIG,
  255. dns.rdataclass.ANY, 0, 0))
  256. rdata_start = self.output.tell()
  257. self.output.write(tsig_rdata)
  258. after = self.output.tell()
  259. assert after - rdata_start < 65536
  260. if after >= self.max_size:
  261. self._rollback(before)
  262. raise dns.exception.TooBig
  263. self.output.seek(rdata_start - 2)
  264. self.output.write(struct.pack('!H', after - rdata_start))
  265. self.counts[ADDITIONAL] += 1
  266. self.output.seek(10)
  267. self.output.write(struct.pack('!H', self.counts[ADDITIONAL]))
  268. self.output.seek(0, 2)
  269. def write_header(self):
  270. """Write the DNS message header.
  271. Writing the DNS message header is done after all sections
  272. have been rendered, but before the optional TSIG signature
  273. is added.
  274. """
  275. self.output.seek(0)
  276. self.output.write(struct.pack('!HHHHHH', self.id, self.flags,
  277. self.counts[0], self.counts[1],
  278. self.counts[2], self.counts[3]))
  279. self.output.seek(0, 2)
  280. def get_wire(self):
  281. """Return the wire format message.
  282. @rtype: string
  283. """
  284. return self.output.getvalue()