jpayne@69: import binascii
jpayne@69: import email.charset
jpayne@69: import email.message
jpayne@69: import email.errors
jpayne@69: from email import quoprimime
jpayne@69: 
jpayne@69: class ContentManager:
jpayne@69: 
jpayne@69:     def __init__(self):
jpayne@69:         self.get_handlers = {}
jpayne@69:         self.set_handlers = {}
jpayne@69: 
jpayne@69:     def add_get_handler(self, key, handler):
jpayne@69:         self.get_handlers[key] = handler
jpayne@69: 
jpayne@69:     def get_content(self, msg, *args, **kw):
jpayne@69:         content_type = msg.get_content_type()
jpayne@69:         if content_type in self.get_handlers:
jpayne@69:             return self.get_handlers[content_type](msg, *args, **kw)
jpayne@69:         maintype = msg.get_content_maintype()
jpayne@69:         if maintype in self.get_handlers:
jpayne@69:             return self.get_handlers[maintype](msg, *args, **kw)
jpayne@69:         if '' in self.get_handlers:
jpayne@69:             return self.get_handlers[''](msg, *args, **kw)
jpayne@69:         raise KeyError(content_type)
jpayne@69: 
jpayne@69:     def add_set_handler(self, typekey, handler):
jpayne@69:         self.set_handlers[typekey] = handler
jpayne@69: 
jpayne@69:     def set_content(self, msg, obj, *args, **kw):
jpayne@69:         if msg.get_content_maintype() == 'multipart':
jpayne@69:             # XXX: is this error a good idea or not?  We can remove it later,
jpayne@69:             # but we can't add it later, so do it for now.
jpayne@69:             raise TypeError("set_content not valid on multipart")
jpayne@69:         handler = self._find_set_handler(msg, obj)
jpayne@69:         msg.clear_content()
jpayne@69:         handler(msg, obj, *args, **kw)
jpayne@69: 
jpayne@69:     def _find_set_handler(self, msg, obj):
jpayne@69:         full_path_for_error = None
jpayne@69:         for typ in type(obj).__mro__:
jpayne@69:             if typ in self.set_handlers:
jpayne@69:                 return self.set_handlers[typ]
jpayne@69:             qname = typ.__qualname__
jpayne@69:             modname = getattr(typ, '__module__', '')
jpayne@69:             full_path = '.'.join((modname, qname)) if modname else qname
jpayne@69:             if full_path_for_error is None:
jpayne@69:                 full_path_for_error = full_path
jpayne@69:             if full_path in self.set_handlers:
jpayne@69:                 return self.set_handlers[full_path]
jpayne@69:             if qname in self.set_handlers:
jpayne@69:                 return self.set_handlers[qname]
jpayne@69:             name = typ.__name__
jpayne@69:             if name in self.set_handlers:
jpayne@69:                 return self.set_handlers[name]
jpayne@69:         if None in self.set_handlers:
jpayne@69:             return self.set_handlers[None]
jpayne@69:         raise KeyError(full_path_for_error)
jpayne@69: 
jpayne@69: 
jpayne@69: raw_data_manager = ContentManager()
jpayne@69: 
jpayne@69: 
jpayne@69: def get_text_content(msg, errors='replace'):
jpayne@69:     content = msg.get_payload(decode=True)
jpayne@69:     charset = msg.get_param('charset', 'ASCII')
jpayne@69:     return content.decode(charset, errors=errors)
jpayne@69: raw_data_manager.add_get_handler('text', get_text_content)
jpayne@69: 
jpayne@69: 
jpayne@69: def get_non_text_content(msg):
jpayne@69:     return msg.get_payload(decode=True)
jpayne@69: for maintype in 'audio image video application'.split():
jpayne@69:     raw_data_manager.add_get_handler(maintype, get_non_text_content)
jpayne@69: 
jpayne@69: 
jpayne@69: def get_message_content(msg):
jpayne@69:     return msg.get_payload(0)
jpayne@69: for subtype in 'rfc822 external-body'.split():
jpayne@69:     raw_data_manager.add_get_handler('message/'+subtype, get_message_content)
jpayne@69: 
jpayne@69: 
jpayne@69: def get_and_fixup_unknown_message_content(msg):
jpayne@69:     # If we don't understand a message subtype, we are supposed to treat it as
jpayne@69:     # if it were application/octet-stream, per
jpayne@69:     # tools.ietf.org/html/rfc2046#section-5.2.4.  Feedparser doesn't do that,
jpayne@69:     # so do our best to fix things up.  Note that it is *not* appropriate to
jpayne@69:     # model message/partial content as Message objects, so they are handled
jpayne@69:     # here as well.  (How to reassemble them is out of scope for this comment :)
jpayne@69:     return bytes(msg.get_payload(0))
jpayne@69: raw_data_manager.add_get_handler('message',
jpayne@69:                                  get_and_fixup_unknown_message_content)
jpayne@69: 
jpayne@69: 
jpayne@69: def _prepare_set(msg, maintype, subtype, headers):
jpayne@69:     msg['Content-Type'] = '/'.join((maintype, subtype))
jpayne@69:     if headers:
jpayne@69:         if not hasattr(headers[0], 'name'):
jpayne@69:             mp = msg.policy
jpayne@69:             headers = [mp.header_factory(*mp.header_source_parse([header]))
jpayne@69:                        for header in headers]
jpayne@69:         try:
jpayne@69:             for header in headers:
jpayne@69:                 if header.defects:
jpayne@69:                     raise header.defects[0]
jpayne@69:                 msg[header.name] = header
jpayne@69:         except email.errors.HeaderDefect as exc:
jpayne@69:             raise ValueError("Invalid header: {}".format(
jpayne@69:                                 header.fold(policy=msg.policy))) from exc
jpayne@69: 
jpayne@69: 
jpayne@69: def _finalize_set(msg, disposition, filename, cid, params):
jpayne@69:     if disposition is None and filename is not None:
jpayne@69:         disposition = 'attachment'
jpayne@69:     if disposition is not None:
jpayne@69:         msg['Content-Disposition'] = disposition
jpayne@69:     if filename is not None:
jpayne@69:         msg.set_param('filename',
jpayne@69:                       filename,
jpayne@69:                       header='Content-Disposition',
jpayne@69:                       replace=True)
jpayne@69:     if cid is not None:
jpayne@69:         msg['Content-ID'] = cid
jpayne@69:     if params is not None:
jpayne@69:         for key, value in params.items():
jpayne@69:             msg.set_param(key, value)
jpayne@69: 
jpayne@69: 
jpayne@69: # XXX: This is a cleaned-up version of base64mime.body_encode (including a bug
jpayne@69: # fix in the calculation of unencoded_bytes_per_line).  It would be nice to
jpayne@69: # drop both this and quoprimime.body_encode in favor of enhanced binascii
jpayne@69: # routines that accepted a max_line_length parameter.
jpayne@69: def _encode_base64(data, max_line_length):
jpayne@69:     encoded_lines = []
jpayne@69:     unencoded_bytes_per_line = max_line_length // 4 * 3
jpayne@69:     for i in range(0, len(data), unencoded_bytes_per_line):
jpayne@69:         thisline = data[i:i+unencoded_bytes_per_line]
jpayne@69:         encoded_lines.append(binascii.b2a_base64(thisline).decode('ascii'))
jpayne@69:     return ''.join(encoded_lines)
jpayne@69: 
jpayne@69: 
jpayne@69: def _encode_text(string, charset, cte, policy):
jpayne@69:     lines = string.encode(charset).splitlines()
jpayne@69:     linesep = policy.linesep.encode('ascii')
jpayne@69:     def embedded_body(lines): return linesep.join(lines) + linesep
jpayne@69:     def normal_body(lines): return b'\n'.join(lines) + b'\n'
jpayne@69:     if cte==None:
jpayne@69:         # Use heuristics to decide on the "best" encoding.
jpayne@69:         try:
jpayne@69:             return '7bit', normal_body(lines).decode('ascii')
jpayne@69:         except UnicodeDecodeError:
jpayne@69:             pass
jpayne@69:         if (policy.cte_type == '8bit' and
jpayne@69:                 max(len(x) for x in lines) <= policy.max_line_length):
jpayne@69:             return '8bit', normal_body(lines).decode('ascii', 'surrogateescape')
jpayne@69:         sniff = embedded_body(lines[:10])
jpayne@69:         sniff_qp = quoprimime.body_encode(sniff.decode('latin-1'),
jpayne@69:                                           policy.max_line_length)
jpayne@69:         sniff_base64 = binascii.b2a_base64(sniff)
jpayne@69:         # This is a little unfair to qp; it includes lineseps, base64 doesn't.
jpayne@69:         if len(sniff_qp) > len(sniff_base64):
jpayne@69:             cte = 'base64'
jpayne@69:         else:
jpayne@69:             cte = 'quoted-printable'
jpayne@69:             if len(lines) <= 10:
jpayne@69:                 return cte, sniff_qp
jpayne@69:     if cte == '7bit':
jpayne@69:         data = normal_body(lines).decode('ascii')
jpayne@69:     elif cte == '8bit':
jpayne@69:         data = normal_body(lines).decode('ascii', 'surrogateescape')
jpayne@69:     elif cte == 'quoted-printable':
jpayne@69:         data = quoprimime.body_encode(normal_body(lines).decode('latin-1'),
jpayne@69:                                       policy.max_line_length)
jpayne@69:     elif cte == 'base64':
jpayne@69:         data = _encode_base64(embedded_body(lines), policy.max_line_length)
jpayne@69:     else:
jpayne@69:         raise ValueError("Unknown content transfer encoding {}".format(cte))
jpayne@69:     return cte, data
jpayne@69: 
jpayne@69: 
jpayne@69: def set_text_content(msg, string, subtype="plain", charset='utf-8', cte=None,
jpayne@69:                      disposition=None, filename=None, cid=None,
jpayne@69:                      params=None, headers=None):
jpayne@69:     _prepare_set(msg, 'text', subtype, headers)
jpayne@69:     cte, payload = _encode_text(string, charset, cte, msg.policy)
jpayne@69:     msg.set_payload(payload)
jpayne@69:     msg.set_param('charset',
jpayne@69:                   email.charset.ALIASES.get(charset, charset),
jpayne@69:                   replace=True)
jpayne@69:     msg['Content-Transfer-Encoding'] = cte
jpayne@69:     _finalize_set(msg, disposition, filename, cid, params)
jpayne@69: raw_data_manager.add_set_handler(str, set_text_content)
jpayne@69: 
jpayne@69: 
jpayne@69: def set_message_content(msg, message, subtype="rfc822", cte=None,
jpayne@69:                        disposition=None, filename=None, cid=None,
jpayne@69:                        params=None, headers=None):
jpayne@69:     if subtype == 'partial':
jpayne@69:         raise ValueError("message/partial is not supported for Message objects")
jpayne@69:     if subtype == 'rfc822':
jpayne@69:         if cte not in (None, '7bit', '8bit', 'binary'):
jpayne@69:             # http://tools.ietf.org/html/rfc2046#section-5.2.1 mandate.
jpayne@69:             raise ValueError(
jpayne@69:                 "message/rfc822 parts do not support cte={}".format(cte))
jpayne@69:         # 8bit will get coerced on serialization if policy.cte_type='7bit'.  We
jpayne@69:         # may end up claiming 8bit when it isn't needed, but the only negative
jpayne@69:         # result of that should be a gateway that needs to coerce to 7bit
jpayne@69:         # having to look through the whole embedded message to discover whether
jpayne@69:         # or not it actually has to do anything.
jpayne@69:         cte = '8bit' if cte is None else cte
jpayne@69:     elif subtype == 'external-body':
jpayne@69:         if cte not in (None, '7bit'):
jpayne@69:             # http://tools.ietf.org/html/rfc2046#section-5.2.3 mandate.
jpayne@69:             raise ValueError(
jpayne@69:                 "message/external-body parts do not support cte={}".format(cte))
jpayne@69:         cte = '7bit'
jpayne@69:     elif cte is None:
jpayne@69:         # http://tools.ietf.org/html/rfc2046#section-5.2.4 says all future
jpayne@69:         # subtypes should be restricted to 7bit, so assume that.
jpayne@69:         cte = '7bit'
jpayne@69:     _prepare_set(msg, 'message', subtype, headers)
jpayne@69:     msg.set_payload([message])
jpayne@69:     msg['Content-Transfer-Encoding'] = cte
jpayne@69:     _finalize_set(msg, disposition, filename, cid, params)
jpayne@69: raw_data_manager.add_set_handler(email.message.Message, set_message_content)
jpayne@69: 
jpayne@69: 
jpayne@69: def set_bytes_content(msg, data, maintype, subtype, cte='base64',
jpayne@69:                      disposition=None, filename=None, cid=None,
jpayne@69:                      params=None, headers=None):
jpayne@69:     _prepare_set(msg, maintype, subtype, headers)
jpayne@69:     if cte == 'base64':
jpayne@69:         data = _encode_base64(data, max_line_length=msg.policy.max_line_length)
jpayne@69:     elif cte == 'quoted-printable':
jpayne@69:         # XXX: quoprimime.body_encode won't encode newline characters in data,
jpayne@69:         # so we can't use it.  This means max_line_length is ignored.  Another
jpayne@69:         # bug to fix later.  (Note: encoders.quopri is broken on line ends.)
jpayne@69:         data = binascii.b2a_qp(data, istext=False, header=False, quotetabs=True)
jpayne@69:         data = data.decode('ascii')
jpayne@69:     elif cte == '7bit':
jpayne@69:         # Make sure it really is only ASCII.  The early warning here seems
jpayne@69:         # worth the overhead...if you care write your own content manager :).
jpayne@69:         data.encode('ascii')
jpayne@69:     elif cte in ('8bit', 'binary'):
jpayne@69:         data = data.decode('ascii', 'surrogateescape')
jpayne@69:     msg.set_payload(data)
jpayne@69:     msg['Content-Transfer-Encoding'] = cte
jpayne@69:     _finalize_set(msg, disposition, filename, cid, params)
jpayne@69: for typ in (bytes, bytearray, memoryview):
jpayne@69:     raw_data_manager.add_set_handler(typ, set_bytes_content)