Generate 'multipart/related' content

multipart-related.py

#!/usr/bin/python3

# [MIT License]
# 
# Copyright (C) 2020  Stefan Vargyas
# 
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
# 
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
# 
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

import sys, os

program = os.path.basename(sys.argv[0])
verdate = '0.1 2020-09-04 14:28' # $ date +'%F %R'

def error(msg, *args):
    if len(args):
        msg = msg % args
    sys.stderr.write("%s: error: %s\n" % (program, msg))
    sys.exit(1)

# stev: the code below was extracted from:
# google-api-python-client/googleapiclient/discovery.py

from six import BytesIO

try:
    from email.generator import BytesGenerator
except ImportError:
    from email.generator import Generator as BytesGenerator
from email.mime.multipart import MIMEMultipart
from email.mime.nonmultipart import MIMENonMultipart

from googleapiclient.http import MediaUpload, MediaFileUpload

# patch _write_lines to avoid munging '\r' into '\n'
# ( https://bugs.python.org/issue18886 https://bugs.python.org/issue19003 )
class _BytesGenerator(BytesGenerator):
    _write_lines = BytesGenerator.write

def multipart_related(body, content_type, media_upload):
    assert isinstance(media_upload, MediaUpload)
    assert not media_upload.resumable()

    # This is a multipart/related upload.
    msgRoot = MIMEMultipart("related")
    # msgRoot should not write out it's own headers
    setattr(msgRoot, "_write_headers", lambda self: None)

    # attach the body as one part
    msg = MIMENonMultipart(*content_type.split("/"))
    msg.set_payload(body)
    msgRoot.attach(msg)

    # attach the media as the second part
    msg = MIMENonMultipart(*media_upload.mimetype().split("/"))
    msg["Content-Transfer-Encoding"] = "binary"

    payload = media_upload.getbytes(0, media_upload.size())
    msg.set_payload(payload)
    msgRoot.attach(msg)
    # encode the body: note that we can't use `as_string`, because
    # it plays games with `From ` lines.
    fp = BytesIO()
    g = _BytesGenerator(fp, mangle_from_=False)
    g.flatten(msgRoot, unixfrom=False)
    return \
        'Content-Type: multipart/related; boundary="%s"' % \
        (msgRoot.get_boundary()), \
        fp.getvalue()

def options():
    from argparse import ArgumentParser, HelpFormatter

    class Formatter(HelpFormatter):

        def _format_action_invocation(self, act):
            # https://stackoverflow.com/a/31124505/8327971
            meta = self._get_default_metavar_for_optional(act)
            return '|'.join(act.option_strings) + ' ' + \
                self._format_args(act, meta)

    p = ArgumentParser(
            formatter_class = Formatter,
            add_help = False)
    p.error = error

    STR = 'STR'
    APJ = 'application/json'
    OUT = 'multipart-related.output'

    # stev: main options:
    p.add_argument('-b', '--body-file',
        help = 'the name of the file containing the HTTP body text',
        metavar = STR, required = True)
    p.add_argument('-c', '--content-type',
        help = 'the `Content-Type\' header corresponding to the '
            'specified body (default: %r)' % (APJ),
        metavar = STR, default = APJ)
    p.add_argument('-m', '--media-file',
        help = 'the name of the file containing the media content to be attached',
        metavar = STR, required = True)
    p.add_argument('-o', '--output-file',
        help = 'the name of the generated file (default: %r)' % (OUT),
        metavar = STR, default = OUT)

    # stev: info options:
    p.add_argument('-v', '--version',
        action = 'version', version = '%(prog)s: version ' + verdate,
        help = 'print version numbers and exit')
    p.add_argument('-h', '--help',
        help = 'display this help info and exit',
        action = 'help')

    return p.parse_args()

def main():
    opt = options()

    with open(opt.body_file, 'rb') as file:
        body = file.read()

    header, bytes = multipart_related(
        content_type = opt.content_type,
        media_upload = MediaFileUpload(
            opt.media_file),
        body = body
    )

    with open(opt.output_file, 'wb') as file:
        file.write(bytes)
    print(header)

if __name__ == '__main__':
    main()