#!/usr/bin/env cctools_python
# CCTOOLS_PYTHON_VERSION 2.7 2.6 2.5 2.4 3

# 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
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() {
    $SFX_DIR/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 $0`

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
%s "$@"
'''

# 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)
        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 map(lambda s: s.split(':'), 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('run.sh')
    run_info_data  = RUN_SH % command
    run_info.mode  = int('755', 8)
    run_info.mtime = time()
    run_info.size  = len(run_info_data)

    logging.debug('adding run.sh...')
    archive.addfile(run_info, BytesIO(run_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)
    raise StopIteration


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 get(self, section, name, default = None):
        try:
            return ConfigParser.get(self, 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.get('starch', 'executables', '').split())
        options.libraries.extend(config.get('starch', 'libraries', '').split())
        options.data.extend(config.get('starch', 'data', '').split())
        options.environments.extend(config.get('starch', 'environments', '').split())
        if not options.command:
            options.command = config.get('starch', 'command', '')

    if not options.executables:
        logging.error('no executables specified')

    if not options.command:
        options.command = os.path.basename(options.executables[0])
        logging.warn('no command specified, so 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:
