#!/usr/bin/env python

# Copyright (c) 2010- The University of Notre Dame.
# This software is distributed under the GNU General Public License.
# See the file COPYING for details.

""" Starch """

import logging
import os
import platform
import sys

from optparse   import OptionParser
from subprocess import Popen, PIPE
from tarfile    import REGTYPE, TarInfo, open as Tar
from time       import time
from tempfile   import NamedTemporaryFile
from shutil     import copyfile, copyfileobj

if sys.version_info[0:2] < (2, 6):
    from cStringIO    import StringIO as BytesIO
else:
    from io           import BytesIO

if sys.version_info[0:2] < (3, 0):
    from ConfigParser import ConfigParser, NoSectionError, NoOptionError
else:
    from configparser import ConfigParser, NoSectionError, NoOptionError

# Todo

"""
1. Keep track of files added to avoid duplication (efficiency).
"""

# Global variables

STARCH_LOGGING_FORMAT = '[%(levelname)s] %(filename)s[%(process)d] %(message)s'
STARCH_AUTODETECT     = True
STARCH_PLATFORM       = 'Linux'

# Shell scripts

SFX_SH='''#!/bin/sh

SFX_FILE=$0
SFX_EXIT_STATUS=0

if [ -z "$SFX_DIR" ]; then
    basename=`basename $SFX_FILE`
    SFX_DIR=$basename.$$.dir
    dir_lock=".${SFX_DIR}.lock"

    if [ -d "$CONDOR_SCRATCH_DIR" ] && [ -r "$CONDOR_SCRATCH_DIR" ] && [ -w "$CONDOR_SCRATCH_DIR" ] && [ -x "$CONDOR_SCRATCH_DIR" ]; then
        SFX_DIR=$CONDOR_SCRATCH_DIR/$SFX_DIR
    elif [ -d "$_CONDOR_SCRATCH_DIR" ] && [ -r "$_CONDOR_SCRATCH_DIR" ] && [ -w "$_CONDOR_SCRATCH_DIR" ] && [ -x "$_CONDOR_SCRATCH_DIR" ]; then
        SFX_DIR=$_CONDOR_SCRATCH_DIR/$SFX_DIR
    else
        SFX_DIR=`pwd`/$SFX_DIR
    fi

    if [ ! -z "$SFX_UNIQUE" ] && [ "$SFX_UNIQUE" = 1 ]; then
        SFX_DIR=`mktemp -d -u $SFX_DIR.XXXXXX`
    fi
else
    dir_lock=".${SFX_DIR}.lock"
fi

extract_real () {
    archive=`awk '/^__ARCHIVE__/ {print NR + 1; exit 0; }' $SFX_FILE`
    tail -n+$archive $SFX_FILE | tar xz -C $SFX_DIR
    SFX_EXIT_STATUS=$?
}

extract() {
    if [ "$SFX_EXTRACT_FORCE" -eq 1 ]; then
        mkdir -p "${SFX_DIR}" 2> /dev/null
        extract_real
    elif [ ! -d "${SFX_DIR}" ]; then
        if { mkdir -p "${SFX_DIR}" && mkdir "${dir_lock}"; } 2> /dev/null; then
            extract_real
            rmdir "${dir_lock}"
        else
            while [ -d "${dir_lock}" ]; do
                sleep 2
            done
        fi
    else
        while [ -d "${dir_lock}" ]; do
            sleep 2
        done
    fi
}

run() {
    # make the extraction directory easily available in bin/run.sh
    export SFX_DIR
    $SFX_DIR/bin/run.sh "$@"
    SFX_EXIT_STATUS=$?
}

if [ -z "$SFX_EXTRACT_ONLY" ]; then
    SFX_EXTRACT_ONLY=0
fi

if [ -z "$SFX_EXTRACT_FORCE" ]; then
    SFX_EXTRACT_FORCE=0
fi

if [ -z "$SFX_KEEP" ]; then
    SFX_KEEP=0
fi

if [ "$SFX_EXTRACT_ONLY" -eq 1 ]; then
    extract
    echo $SFX_DIR
    SFX_KEEP=1
else
    extract && run "$@"
fi

if [ "$SFX_KEEP" -ne 1 ]; then
    rm -fr $SFX_DIR
fi

exit $SFX_EXIT_STATUS

__ARCHIVE__
'''

RUN_SH = '''#!/bin/sh
SFX_DIR=$(dirname $( cd -- "$( dirname -- "$0" )" > /dev/null 2>&1 && pwd ))

if [ -d "$SFX_DIR/env" ]; then
    for f in $SFX_DIR/env/*; do
        . $f
    done
fi

if [ -d "$SFX_DIR/bin" ]; then
    export PATH=$SFX_DIR/bin:$PATH
fi

if [ -d "$SFX_DIR/lib" ]; then
    if [ x"`uname -s`" = x"Darwin" ]; then
        if [ -z "$DYLD_LIBRARY_PATH" ]; then
            export DYLD_LIBRARY_PATH=$SFX_DIR/lib
        else
            export DYLD_LIBRARY_PATH=$SFX_DIR/lib:$DYLD_LIBRARY_PATH
        fi
    else
        if [ -z "$LD_LIBRARY_PATH" ]; then
            export LD_LIBRARY_PATH=$SFX_DIR/lib
        else
            export LD_LIBRARY_PATH=$SFX_DIR/lib:$LD_LIBRARY_PATH
        fi
    fi
fi

if [ "$1" = "--sfx-help" ]
then
    echo starch package options:
    echo executing the sfx file by itself executes the command:
    echo  "{cmd}"   
    echo
    echo environment variables to set packages options:
    echo SFX_EXEC=CMD        Execute CMD instead of the default.
    echo SFX_EXTRACT_ONLY=1  Extract package, but do not run a command.
    echo SFX_KEEP=1          Do not remove extracted package when command finishes.
    echo SFX_DIR=DIR         Extract to DIR, instead of a temporary directory. Implies SFX_KEEP=1
    echo SFX_ENV=1           Execute the arguments given as a command inside the starch environment

    exit 0
fi

if [ -n "$SFX_EXEC" ]
then
    $SFX_EXEC "$@"
elif [ "$SFX_ENV" = 1 ]
then
    "$@"
else
    {cmd} "$@"
fi
'''

RUN_IN_ENV = '''#!/bin/sh
export SFX_ENV=1

SFX_DIR=$(dirname $( cd -- "$( dirname -- "$0" )" > /dev/null 2>&1 && pwd ))

exec $SFX_DIR/bin/run.sh "$@"
'''

# Create SFX
def create_sfx(sfx_path, executables, libraries, data, environments, command):
    """ Create self-extracting executable

    Directory structure:
        bin         Executables
        lib         Libraries
        env         Environment Scripts (static)
        bin/run.sh  Script that contains command
    """

    tmp_file = NamedTemporaryFile()
    tmp_file.close()

    arc_path = tmp_file.name
    archive  = Tar(arc_path, 'w:gz')

    logging.debug('adding executables...')
    executables = find_executables(executables)
    for exe_path, real_path in executables:
        exe_name = os.path.basename(exe_path)
        exe_info = archive.gettarinfo(real_path, os.path.join('bin', exe_name))
        exe_info.mode = int('755', 8)

        logging.debug('    adding executable: %s (%s)' % (exe_name, real_path))
        archive.addfile(exe_info, open(real_path, 'rb'))

    logging.debug('adding libraries...')
    libraries = find_libraries(libraries, executables)
    for lib_path, real_path in libraries:
        lib_name = os.path.basename(lib_path)
        lib_info = archive.gettarinfo(real_path, os.path.join('lib', lib_name))

        logging.debug('    adding library: %s (%s)' % (lib_name, real_path))
        archive.addfile(lib_info, open(real_path, 'rb'))

    logging.debug('adding data...')
    for data_path, real_path in [s.split(':') for s in data]:
        add_data_to_archive(archive, os.path.normpath(data_path), os.path.normpath(real_path))

    logging.debug('adding environment scripts...')
    for env_path, real_path in find_files(environments, 'PWD'):
        env_name = os.path.basename(env_path)
        env_info = archive.gettarinfo(real_path, os.path.join('env', env_name))

        logging.debug('    adding environment script: %s (%s)' % (env_name, real_path))
        archive.addfile(env_info, open(real_path, 'rb'))

    run_info = TarInfo('bin/run.sh')
    run_info_data  = RUN_SH.format(cmd=command)
    run_info.mode  = int('755', 8)
    run_info.mtime = time()
    run_info.size  = len(run_info_data)

    run_in_env_info = TarInfo('bin/run_in_env')
    run_in_env_info_data  = RUN_IN_ENV
    run_in_env_info.mode  = int('755', 8)
    run_in_env_info.mtime = time()
    run_in_env_info.size  = len(run_in_env_info_data)

    logging.debug('adding bin/run.sh...')
    archive.addfile(run_info, BytesIO(run_info_data.encode('utf-8')))

    logging.debug('adding bin/run_in_env...')
    archive.addfile(run_in_env_info, BytesIO(run_in_env_info_data.encode('utf-8')))

    archive.close()

    logging.debug('creating sfx...')
    sfx_file = open(sfx_path, 'wb')
    copyfileobj(BytesIO(SFX_SH.encode('utf-8')), sfx_file)
    copyfileobj(open(arc_path, 'rb'), sfx_file)
    sfx_file.close()

    logging.debug('cleaning up...')
    os.chmod(sfx_path, int('755', 8))
    os.unlink(arc_path)

def add_data_to_archive(archive, data_path, real_path):
    if os.path.isdir(real_path):
        for root, dirs, files in os.walk(real_path):
            for n in files + dirs:
                dp = os.path.join(root.replace(real_path, data_path), n)
                rp = os.path.join(root, n)
                add_data_to_archive(archive, dp, rp)
    else:
        data_info = archive.gettarinfo(os.path.realpath(real_path), data_path)
        logging.debug('    adding data: %s (%s)' % (data_path, real_path))
        archive.addfile(data_info, open(real_path, 'rb'))


# Find file utilities

def find_files(files, env_var, default_paths = None):
    if default_paths:
        paths = default_paths
    else:
        paths = [os.curdir]

    if env_var in os.environ:
        paths.extend(os.environ[env_var].split(':'))

    for file in files:
        is_found = False
        for path in paths:
            file_path = os.path.join(path, file)
            if os.path.exists(file_path):
                is_found = True
                yield file_path, os.path.realpath(file_path)
                break
        if not is_found:
            logging.error('could not find file: %s' % file)

def find_executables(executables):
    exes = []
    for ep, rp in find_files(executables, 'PATH'):
        exes.append((ep, rp))
    return exes


def find_libraries(libraries, executables):
    libs = []
    search = ['/lib', '/lib64', '/usr/lib', '/usr/lib64']

    if platform.system() == 'Linux':
        platlib = '%s-linux-gnu' % platform.machine()
        platsearch = []
        for s in search:
            platsearch.append(os.path.join(s, platlib))
        search.extend(platsearch)

    for lp, rp in find_files(libraries, 'LD_LIBRARY_PATH', search):
        libs.append((lp, rp))

    if STARCH_AUTODETECT:
        for exe_path, real_path in executables:
            if STARCH_PLATFORM == 'Darwin':
                libs.extend(autodetect_libraries_darwin(exe_path))
            else:
                libs.extend(autodetect_libraries_linux(exe_path))

    return libs


def autodetect_libraries_linux(executable):
    libs = []

    try:
        p = Popen(['ldd', executable], stdout=PIPE)
        for line in p.communicate()[0].decode().split('\n'):
            try:
                lib_path = line.split('=>')[-1].strip().split()[0]
                if os.path.exists(lib_path):
                    libs.append((lib_path, os.path.realpath(lib_path)))
            except Exception:
                pass
    except Exception:
        _, e = sys.exc_info()[:2]
        logging.error('could not execute ldd on %s: %s' % (executable, str(e)))

    return libs


def autodetect_libraries_darwin(executable):
    libs = []

    try:
        p = Popen(['otool', '-L', executable], stdout=PIPE)
        for line in p.communicate()[0].decode().split('\n'):
            try:
                lib_path = line.split('(')[0].strip()
                if os.path.exists(lib_path):
                    libs.append((lib_path, os.path.realpath(lib_path)))
            except Exception:
                pass
    except Exception:
        _, e = sys.exc_info()[:2]
        logging.error('could not execute otool on %s: %s' % (executable, str(e)))

    return libs

# Configuration Parser
class StarchConfigParser(ConfigParser):
    def safe_get(self, section, name, default = None):
        try:
            return self.get(section, name)
        except (NoSectionError, NoOptionError):
            return default

# Parse commandline options
def parse_command_line_options():
    global STARCH_AUTODETECT
    global STARCH_PLATFORM

    parser = OptionParser('%prog [options] <sfx_path>')

    parser.add_option('-A', dest = 'autodetect_disable', action = 'store_true',
        help = 'Disable automatic detection of library dependencies', default=False)
    parser.add_option('-C', dest = 'config', action = 'store',
        help = 'Use configuration file.', metavar = 'cfg', default = None)
    parser.add_option('-c', dest = 'command', action = 'store',
        help = 'Specify command to execute.', metavar = 'cmd', default = '')
    parser.add_option('-d', dest = 'data', action = 'append',
        help = 'Add data (new path:old path).', metavar = 'npath:opath', default = [])
    parser.add_option('-e', dest = 'environments', action = 'append',
        help = 'Add environment script.', metavar = 'env', default = [])
    parser.add_option('-l', dest = 'libraries', action = 'append',
        help = 'Add library.', metavar = 'lib', default = [])
    parser.add_option('-x', dest = 'executables', action = 'append',
        help = 'Add executable.', metavar = 'exe', default = [])
    parser.add_option('-v', dest = 'verbose', action = 'store_true',
        help = 'Display verbose messages.', default = False)

    (options, args) = parser.parse_args()

    if len(args) != 1:
        parser.print_help()
        sys.exit(1)

    if options.verbose:
        logging.basicConfig(format=STARCH_LOGGING_FORMAT, level=logging.DEBUG)
    else:
        logging.basicConfig(format=STARCH_LOGGING_FORMAT)

    STARCH_AUTODETECT = not(options.autodetect_disable)
    STARCH_PLATFORM   = os.uname()[0]

    if options.config:
        if not os.path.exists(options.config):
            logging.error('config file \'%s\' does not exist' % options.config)

        config = StarchConfigParser()
        config.read(options.config)

        options.executables.extend(config.safe_get('starch', 'executables', '').split())
        options.libraries.extend(config.safe_get('starch', 'libraries', '').split())
        options.data.extend(config.safe_get('starch', 'data', '').split())
        options.environments.extend(config.safe_get('starch', 'environments', '').split())
        if not options.command:
            options.command = config.safe_get('starch', 'command', '')

    if not options.executables:
        logging.error('At least one executable should be specified with -x')
        parser.print_help()
        sys.exit(1)


    if not options.command:
        options.command = os.path.basename(options.executables[0])
        logging.warn('no command specified with -c, using: %s' % options.command)

    logging.debug('command ... ' + options.command)

    return args[0], options.executables, options.libraries, \
                    options.data, options.environments, options.command

# Main execution
if __name__ == '__main__':
    create_sfx(*parse_command_line_options())

# vim: set sts=4 sw=4 ts=8 expandtab ft=python:
