#!/usr/bin/python3

import sys
import os
import zipfile
import tempfile
from xml.etree import ElementTree
from shutil import copyfile

def stuffer(py_file, doc_file):
    # <Override PartName="__main__.py" ContentType="text/plain"/>
    # Create temp files for new doc/zip and content_types.xml
    tmpdz, temp_doczip = tempfile.mkstemp(dir=os.path.dirname(doc_file))
    os.close(tmpdz)
    
    with zipfile.ZipFile(doc_file, "r") as zipread:
        # read and patch the [Content_Types].xml
        orig_content = zipread.read("[Content_Types].xml")
        root = ElementTree.fromstring(orig_content)
        new_element = ElementTree.Element('ns0:Override PartName="/__main__.py" ContentType="text/plain"')
        root.insert(1, new_element)

        # Clone original doc without the [Content_Types].xml file
        with zipfile.ZipFile(temp_doczip, "w") as newdoc:
            for item in zipread.infolist():
                if item.filename != "[Content_Types].xml":
                    newdoc.writestr(item, zipread.read(item.filename))
        
        # add __main__.py and new [Content_Types].xml
        with zipfile.ZipFile(temp_doczip, mode="a") as zd:
            copyfile(py_file, "__main__.py")
            zd.write("__main__.py")
            os.remove("__main__.py")
            zd.writestr("[Content_Type].xml", ElementTree.tostring(root))

    # Nuke original doc file and replace with newly stuffed doc
    os.remove(doc_file)
    os.rename(temp_doczip, doc_file)


if __name__ == '__main__':
    if len(sys.argv) < 2:
        print("Usage:\r\n\t{} /path/to/script.py /path/to/word.doc /optional/output/name.docx".format(sys.argv[0]))
    else:
        if os.path.exists(sys.argv[1]) and os.path.exists(sys.argv[2]): # both files are accessible
            # name the output file
            if not len(sys.argv) == 3:
                out_name = sys.argv[3]     
            else:
                out_name = "pydocexec_{}".format(os.path.basename(sys.argv[2]))
            out_file = os.path.join(os.path.dirname(sys.argv[2]), out_name)

            # Clone the doc to doc directory
            copyfile(sys.argv[2], out_file)

            # Stuff it!
            stuffer(sys.argv[1], out_file)

        else:
            print("Check the file paths. One of the files was not found.")