contentmanager.py 10 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248
  1. import binascii
  2. import email.charset
  3. import email.message
  4. import email.errors
  5. from email import quoprimime
  6. class ContentManager:
  7. def __init__(self):
  8. self.get_handlers = {}
  9. self.set_handlers = {}
  10. def add_get_handler(self, key, handler):
  11. self.get_handlers[key] = handler
  12. def get_content(self, msg, *args, **kw):
  13. content_type = msg.get_content_type()
  14. if content_type in self.get_handlers:
  15. return self.get_handlers[content_type](msg, *args, **kw)
  16. maintype = msg.get_content_maintype()
  17. if maintype in self.get_handlers:
  18. return self.get_handlers[maintype](msg, *args, **kw)
  19. if '' in self.get_handlers:
  20. return self.get_handlers[''](msg, *args, **kw)
  21. raise KeyError(content_type)
  22. def add_set_handler(self, typekey, handler):
  23. self.set_handlers[typekey] = handler
  24. def set_content(self, msg, obj, *args, **kw):
  25. if msg.get_content_maintype() == 'multipart':
  26. # XXX: is this error a good idea or not? We can remove it later,
  27. # but we can't add it later, so do it for now.
  28. raise TypeError("set_content not valid on multipart")
  29. handler = self._find_set_handler(msg, obj)
  30. msg.clear_content()
  31. handler(msg, obj, *args, **kw)
  32. def _find_set_handler(self, msg, obj):
  33. full_path_for_error = None
  34. for typ in type(obj).__mro__:
  35. if typ in self.set_handlers:
  36. return self.set_handlers[typ]
  37. qname = typ.__qualname__
  38. modname = getattr(typ, '__module__', '')
  39. full_path = '.'.join((modname, qname)) if modname else qname
  40. if full_path_for_error is None:
  41. full_path_for_error = full_path
  42. if full_path in self.set_handlers:
  43. return self.set_handlers[full_path]
  44. if qname in self.set_handlers:
  45. return self.set_handlers[qname]
  46. name = typ.__name__
  47. if name in self.set_handlers:
  48. return self.set_handlers[name]
  49. if None in self.set_handlers:
  50. return self.set_handlers[None]
  51. raise KeyError(full_path_for_error)
  52. raw_data_manager = ContentManager()
  53. def get_text_content(msg, errors='replace'):
  54. content = msg.get_payload(decode=True)
  55. charset = msg.get_param('charset', 'ASCII')
  56. return content.decode(charset, errors=errors)
  57. raw_data_manager.add_get_handler('text', get_text_content)
  58. def get_non_text_content(msg):
  59. return msg.get_payload(decode=True)
  60. for maintype in 'audio image video application'.split():
  61. raw_data_manager.add_get_handler(maintype, get_non_text_content)
  62. def get_message_content(msg):
  63. return msg.get_payload(0)
  64. for subtype in 'rfc822 external-body'.split():
  65. raw_data_manager.add_get_handler('message/'+subtype, get_message_content)
  66. def get_and_fixup_unknown_message_content(msg):
  67. # If we don't understand a message subtype, we are supposed to treat it as
  68. # if it were application/octet-stream, per
  69. # tools.ietf.org/html/rfc2046#section-5.2.4. Feedparser doesn't do that,
  70. # so do our best to fix things up. Note that it is *not* appropriate to
  71. # model message/partial content as Message objects, so they are handled
  72. # here as well. (How to reassemble them is out of scope for this comment :)
  73. return bytes(msg.get_payload(0))
  74. raw_data_manager.add_get_handler('message',
  75. get_and_fixup_unknown_message_content)
  76. def _prepare_set(msg, maintype, subtype, headers):
  77. msg['Content-Type'] = '/'.join((maintype, subtype))
  78. if headers:
  79. if not hasattr(headers[0], 'name'):
  80. mp = msg.policy
  81. headers = [mp.header_factory(*mp.header_source_parse([header]))
  82. for header in headers]
  83. try:
  84. for header in headers:
  85. if header.defects:
  86. raise header.defects[0]
  87. msg[header.name] = header
  88. except email.errors.HeaderDefect as exc:
  89. raise ValueError("Invalid header: {}".format(
  90. header.fold(policy=msg.policy))) from exc
  91. def _finalize_set(msg, disposition, filename, cid, params):
  92. if disposition is None and filename is not None:
  93. disposition = 'attachment'
  94. if disposition is not None:
  95. msg['Content-Disposition'] = disposition
  96. if filename is not None:
  97. msg.set_param('filename',
  98. filename,
  99. header='Content-Disposition',
  100. replace=True)
  101. if cid is not None:
  102. msg['Content-ID'] = cid
  103. if params is not None:
  104. for key, value in params.items():
  105. msg.set_param(key, value)
  106. # XXX: This is a cleaned-up version of base64mime.body_encode (including a bug
  107. # fix in the calculation of unencoded_bytes_per_line). It would be nice to
  108. # drop both this and quoprimime.body_encode in favor of enhanced binascii
  109. # routines that accepted a max_line_length parameter.
  110. def _encode_base64(data, max_line_length):
  111. encoded_lines = []
  112. unencoded_bytes_per_line = max_line_length // 4 * 3
  113. for i in range(0, len(data), unencoded_bytes_per_line):
  114. thisline = data[i:i+unencoded_bytes_per_line]
  115. encoded_lines.append(binascii.b2a_base64(thisline).decode('ascii'))
  116. return ''.join(encoded_lines)
  117. def _encode_text(string, charset, cte, policy):
  118. lines = string.encode(charset).splitlines()
  119. linesep = policy.linesep.encode('ascii')
  120. def embedded_body(lines): return linesep.join(lines) + linesep
  121. def normal_body(lines): return b'\n'.join(lines) + b'\n'
  122. if cte is None:
  123. # Use heuristics to decide on the "best" encoding.
  124. if max((len(x) for x in lines), default=0) <= policy.max_line_length:
  125. try:
  126. return '7bit', normal_body(lines).decode('ascii')
  127. except UnicodeDecodeError:
  128. pass
  129. if policy.cte_type == '8bit':
  130. return '8bit', normal_body(lines).decode('ascii', 'surrogateescape')
  131. sniff = embedded_body(lines[:10])
  132. sniff_qp = quoprimime.body_encode(sniff.decode('latin-1'),
  133. policy.max_line_length)
  134. sniff_base64 = binascii.b2a_base64(sniff)
  135. # This is a little unfair to qp; it includes lineseps, base64 doesn't.
  136. if len(sniff_qp) > len(sniff_base64):
  137. cte = 'base64'
  138. else:
  139. cte = 'quoted-printable'
  140. if len(lines) <= 10:
  141. return cte, sniff_qp
  142. if cte == '7bit':
  143. data = normal_body(lines).decode('ascii')
  144. elif cte == '8bit':
  145. data = normal_body(lines).decode('ascii', 'surrogateescape')
  146. elif cte == 'quoted-printable':
  147. data = quoprimime.body_encode(normal_body(lines).decode('latin-1'),
  148. policy.max_line_length)
  149. elif cte == 'base64':
  150. data = _encode_base64(embedded_body(lines), policy.max_line_length)
  151. else:
  152. raise ValueError("Unknown content transfer encoding {}".format(cte))
  153. return cte, data
  154. def set_text_content(msg, string, subtype="plain", charset='utf-8', cte=None,
  155. disposition=None, filename=None, cid=None,
  156. params=None, headers=None):
  157. _prepare_set(msg, 'text', subtype, headers)
  158. cte, payload = _encode_text(string, charset, cte, msg.policy)
  159. msg.set_payload(payload)
  160. msg.set_param('charset',
  161. email.charset.ALIASES.get(charset, charset),
  162. replace=True)
  163. msg['Content-Transfer-Encoding'] = cte
  164. _finalize_set(msg, disposition, filename, cid, params)
  165. raw_data_manager.add_set_handler(str, set_text_content)
  166. def set_message_content(msg, message, subtype="rfc822", cte=None,
  167. disposition=None, filename=None, cid=None,
  168. params=None, headers=None):
  169. if subtype == 'partial':
  170. raise ValueError("message/partial is not supported for Message objects")
  171. if subtype == 'rfc822':
  172. if cte not in (None, '7bit', '8bit', 'binary'):
  173. # http://tools.ietf.org/html/rfc2046#section-5.2.1 mandate.
  174. raise ValueError(
  175. "message/rfc822 parts do not support cte={}".format(cte))
  176. # 8bit will get coerced on serialization if policy.cte_type='7bit'. We
  177. # may end up claiming 8bit when it isn't needed, but the only negative
  178. # result of that should be a gateway that needs to coerce to 7bit
  179. # having to look through the whole embedded message to discover whether
  180. # or not it actually has to do anything.
  181. cte = '8bit' if cte is None else cte
  182. elif subtype == 'external-body':
  183. if cte not in (None, '7bit'):
  184. # http://tools.ietf.org/html/rfc2046#section-5.2.3 mandate.
  185. raise ValueError(
  186. "message/external-body parts do not support cte={}".format(cte))
  187. cte = '7bit'
  188. elif cte is None:
  189. # http://tools.ietf.org/html/rfc2046#section-5.2.4 says all future
  190. # subtypes should be restricted to 7bit, so assume that.
  191. cte = '7bit'
  192. _prepare_set(msg, 'message', subtype, headers)
  193. msg.set_payload([message])
  194. msg['Content-Transfer-Encoding'] = cte
  195. _finalize_set(msg, disposition, filename, cid, params)
  196. raw_data_manager.add_set_handler(email.message.Message, set_message_content)
  197. def set_bytes_content(msg, data, maintype, subtype, cte='base64',
  198. disposition=None, filename=None, cid=None,
  199. params=None, headers=None):
  200. _prepare_set(msg, maintype, subtype, headers)
  201. if cte == 'base64':
  202. data = _encode_base64(data, max_line_length=msg.policy.max_line_length)
  203. elif cte == 'quoted-printable':
  204. # XXX: quoprimime.body_encode won't encode newline characters in data,
  205. # so we can't use it. This means max_line_length is ignored. Another
  206. # bug to fix later. (Note: encoders.quopri is broken on line ends.)
  207. data = binascii.b2a_qp(data, istext=False, header=False, quotetabs=True)
  208. data = data.decode('ascii')
  209. elif cte == '7bit':
  210. data = data.decode('ascii')
  211. elif cte in ('8bit', 'binary'):
  212. data = data.decode('ascii', 'surrogateescape')
  213. msg.set_payload(data)
  214. msg['Content-Transfer-Encoding'] = cte
  215. _finalize_set(msg, disposition, filename, cid, params)
  216. for typ in (bytes, bytearray, memoryview):
  217. raw_data_manager.add_set_handler(typ, set_bytes_content)