third_party/WebKit/LayoutTests/mhtml/resources/mhtml_generator.py

   1 #!/usr/bin/python3
   2 # -*- coding: utf-8 -*-
   3
   4 # Copyright (c) 2013, Opera Software ASA. All rights reserved.
   5 #
   6 # Redistribution and use in source and binary forms, with or without
   7 # modification, are permitted provided that the following conditions
   8 # are met:
   9 #
  10 # 1. Redistributions of source code must retain the above copyright
  11 #    notice, this list of conditions and the following disclaimer.
  12 # 2. Redistributions in binary form must reproduce the above copyright
  13 #    notice, this list of conditions and the following disclaimer in the
  14 #    documentation and/or other materials provided with the distribution.
  15 # 3. Neither the name of Opera Software ASA nor the names of its
  16 #    contributors may be used to endorse or promote products derived
  17 #    from this software without specific prior written permission.
  18 #
  19 # THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
  20 # "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
  21 # LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS
  22 # FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE
  23 # COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT,
  24 # INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
  25 # (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
  26 # SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION)
  27 # HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT,
  28 # STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
  29 # ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED
  30 # OF THE POSSIBILITY OF SUCH DAMAGE.
  31
  32 import email
  33 import email.encoders
  34 import email.generator
  35 import email.mime.image
  36 import email.mime.multipart
  37 import email.mime.nonmultipart
  38 import email.mime.text
  39 import mimetypes
  40 import os.path
  41 import quopri
  42 import sys
  43
  44 class ArgumentError(Exception):
  45     pass
  46
  47 def _encode_quopri(msg):
  48     """Own version of quopri isntead of email.encoders.quopri which seems to
  49        be buggy in python3"""
  50     orig = msg.get_payload()
  51     encdata = quopri.encodestring(orig, quotetabs=True)
  52     encdata.replace(b' ', b'=20')
  53     msg.set_payload(encdata.decode('ascii', 'surrogateescape'))
  54     msg['Content-Transfer-Encoding'] = 'quoted-printable'
  55
  56 def _encode_binary(msg):
  57     email.encoders.encode_noop(msg)
  58     msg['Content-Transfer-Encoding'] = 'binary'
  59
  60 TRANSFER_ENCODINGS = {
  61     "8bit": email.encoders.encode_7or8bit,
  62     "7bit": email.encoders.encode_7or8bit,
  63     "base64": email.encoders.encode_base64,
  64     "binary": _encode_binary,
  65     "none": email.encoders.encode_noop,
  66     "quoted-printable": _encode_quopri}
  67
  68 BASE = "http://test/"
  69
  70 def generate_message(parts):
  71     """Generate a mime message from the given parts"""
  72
  73     main = email.mime.multipart.MIMEMultipart("related")
  74     main.add_header("Content-Location", BASE + parts[0]["name"])
  75
  76     for part in parts:
  77         with open(part["name"], 'rb') as payload:
  78             sub = email.mime.text.MIMENonMultipart(*part["mime"].split("/"))
  79             sub.add_header("Content-Location", BASE + part["name"])
  80             sub.set_payload(payload.read())
  81             TRANSFER_ENCODINGS[part["transfer_encoding"]](sub)
  82             main.attach(sub)
  83
  84     return main
  85
  86 def parse_arguments(args):
  87     """Parse arguments to extract file, transfer encoding, mime pairs"""
  88
  89     parts = []
  90     current = {}
  91
  92     for arg in args:
  93         if os.path.isfile(arg):
  94             if current:
  95                 parts.append(current)
  96             current = {}
  97             current["name"] = arg
  98             current["transfer_encoding"] = "binary"
  99             current["mime"] = mimetypes.guess_type(arg)[0]
 100         elif arg.lower() in TRANSFER_ENCODINGS:
 101             current["transfer_encoding"] = arg.lower()
 102         elif "/" in arg:
 103             current["mime"] = arg.lower()
 104         else:
 105             raise ArgumentError("Unknown argument '" + arg + "'")
 106
 107     if current:
 108         parts.append(current)
 109
 110     return parts
 111
 112 def main():
 113     PARTS = parse_arguments(sys.argv[1:])
 114     MESSAGE = generate_message(PARTS)
 115
 116     GENERATOR = email.generator.Generator(sys.stdout, mangle_from_=True,
 117                                           maxheaderlen=1000)
 118     GENERATOR.flatten(MESSAGE, linesep="\r\n")
 119
 120 if __name__ == "__main__":
 121     main()
 122