#! /usr/bin/env python
# -*- coding: utf-8 -*-
#
#############################################################
#                                                           #
#      Copyright @ 2018 -  Dashingsoft corp.                #
#      All rights reserved.                                 #
#                                                           #
#      pyarmor                                              #
#                                                           #
#      Version: 4.3.2 -                                     #
#                                                           #
#############################################################
#
#
#  @File: packer.py
#
#  @Author: Jondy Zhao(jondy.zhao@gmail.com)
#
#  @Create Date: 2018/11/08
#
#  @Description:
#
#   Pack obfuscated Python scripts with any of third party
#   tools: py2exe, py2app, cx_Freeze, PyInstaller
#

'''After the py2exe or cx_Freeze setup script works, this tool could
obfuscate all the python scripts and package them again. The basic
usage:

    python packer.py --type py2exe /path/to/src/hello.py

It will replace all the original python scripts with obfuscated ones
in the compressed archive generated by py2exe or cx_Freeze.

'''

import logging
import os
import shutil
import subprocess
import sys
import time

from distutils.util import get_platform
from glob import glob
from py_compile import compile as compile_file
from zipfile import PyZipFile

try:
    import argparse
except ImportError:
    # argparse is new in version 2.7
    import polyfills.argparse as argparse

# Default output path, library name, command options for setup script
DEFAULT_PACKER = {
    'py2app': ('dist', 'library.zip', ['py2app', '--dist-dir']),
    'py2exe': ('dist', 'library.zip', ['py2exe', '--dist-dir']),
    'pyInstaller': ('dist', 'library.zip', ['build', '--dist-dir']),
    'cx_Freeze': (
        os.path.join(
            'build', 'exe.%s-%s' % (get_platform(), sys.version[0:3])),
        'python%s%s.zip' % sys.version_info[:2],
        ['build', '--build-exe'])
}

def logaction(func):
    def wrap(*args, **kwargs):
        logging.info('')
        logging.info('%s', func.__name__)
        return func(*args, **kwargs)
    return wrap

@logaction
def update_library(obfdist, libzip):
    '''Update compressed library generated by py2exe or cx_Freeze, replace
the original scripts with obfuscated ones.

    '''
    # # It's simple ,but there are duplicated .pyc files
    # with PyZipFile(libzip, 'a') as f:
    #     f.writepy(obfdist)
    filelist = []
    for root, dirs, files in os.walk(obfdist):
        filelist.extend([os.path.join(root, s) for s in files])

    with PyZipFile(libzip, 'r') as f:
        namelist = f.namelist()
        f.extractall(obfdist)

    for s in filelist:
        if s.lower().endswith('.py'):
            compile_file(s, s + 'c')

    with PyZipFile(libzip, 'w') as f:
        for name in namelist:
            f.write(os.path.join(obfdist, name), name)

@logaction
def run_setup_script(src, entry, build, script, packcmd, obfdist):
    '''Update entry script, copy pytransform.py to source path, then run
setup script to build the bundle.

    '''
    obf_entry = os.path.join(obfdist, entry)

    tempfile = '%s.armor.bak' % entry
    shutil.move(os.path.join(src, entry), tempfile)
    shutil.move(obf_entry, src)
    shutil.copy('pytransform.py', src)

    p = subprocess.Popen([sys.executable, script] + packcmd, cwd=build,
                         stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
    stdoutdata, _ = p.communicate()

    shutil.move(tempfile, os.path.join(src, entry))
    os.remove(os.path.join(src, 'pytransform.py'))

    if p.returncode != 0:
        logging.error('\n\n%s\n\n', stdoutdata)
        raise RuntimeError('Run setup script failed')

@logaction
def copy_runtime_files(runtimes, output):
    for s in 'pyshield.key', 'pyshield.lic', 'product.key', 'license.lic':
        shutil.copy(os.path.join(runtimes, s), output)
    for dllname in glob(os.path.join(runtimes, '_pytransform.*')):
        shutil.copy(dllname, output)

def call_armor(args):
    logging.info('')
    logging.info('')
    p = subprocess.Popen([sys.executable, 'pyarmor.py'] + list(args))
    p.wait()
    if p.returncode != 0:
        raise RuntimeError('Call pyarmor failed')

def pathwrapper(func):
    def wrap(*args, **kwargs):
        oldpath = os.getcwd()
        os.chdir(os.path.abspath(os.path.dirname(__file__)))
        logging.info('Change path to %s', os.getcwd())
        try:
            return func(*args, **kwargs)
        finally:
            os.chdir(oldpath)
    return wrap

@pathwrapper
def _packer(src, entry, build, script, packcmd, output, libname):
    project = os.path.join('projects', 'build-for-packer-v0.1')
    obfdist = os.path.join(project, 'dist')

    args = 'init', '-t', 'app', '--src', src, '--entry', entry, project
    call_armor(args)

    filters = ('global-include *.py', 'prune build, prune dist',
               'exclude %s pytransform.py' % entry)
    args = ('config', '--runtime-path', '', '--disable-restrict-mode', '1',
               '--manifest', ','.join(filters), project)
    call_armor(args)

    args = 'build', project
    call_armor(args)

    run_setup_script(src, entry, build, script, packcmd, obfdist)

    update_library(obfdist, os.path.join(output, libname))

    copy_runtime_files(obfdist, output)

    shutil.rmtree(project)

def packer(args):
    src = os.path.abspath(os.path.dirname(args.entry[0]))
    entry = os.path.basename(args.entry[0])

    if args.setup is None:
        build = src
        script = 'setup.py'
    else:
        build = os.path.abspath(os.path.dirname(args.setup))
        script = os.path.basename(args.setup)

    if args.output is None:
        dist = DEFAULT_PACKER[args.type][0]
        output = os.path.normpath(os.path.join(build, dist))
    else:
        output = args.output if os.path.isabs(args.output) \
            else os.path.join(build, args.output)

    libname = DEFAULT_PACKER[args.type][1]
    packcmd = DEFAULT_PACKER[args.type][2] + [output]

    logging.info('Prepare to pack obfuscated scripts with %s', args.type)
    _packer(src, entry, build, script, packcmd, output, libname)

    logging.info('')
    logging.info('Pack obfuscated scripts successfully in the path')
    logging.info('')
    logging.info('\t%s', output)

def add_arguments(parser):
    parser.add_argument('-v', '--version', action='version', version='v0.1')

    parser.add_argument('-t', '--type', default='py2exe', metavar='TYPE',
                        choices=DEFAULT_PACKER.keys(),
                        help=', '.join(DEFAULT_PACKER.keys()))
    # parser.add_argument('-p', '--path',
    #                     help='Base path, default is the path of entry script')
    parser.add_argument('-s', '--setup',
                        help='Setup script, default is setup.py')
    parser.add_argument('-O', '--output',
                        help='Directory to put final built distributions in' \
                        ' (default is output path of setup script)')
    parser.add_argument('entry', metavar='SCRIPT', nargs=1,
                        help='Entry script')

def main(args):
    parser = argparse.ArgumentParser(
        prog='packer.py',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        description='Pack obfuscated scripts',
        epilog=__doc__,
    )
    add_arguments(parser)
    packer(parser.parse_args(args))

if __name__ == '__main__':
    logging.basicConfig(
        level=logging.INFO,
        format='%(levelname)-8s %(message)s',
    )
    main(sys.argv[1:])
