#!/bin/ash

#
# Copyright (c) 2009 Charles Wilson
# Based on rebaseall by Jason Tishler
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; either version 2 of the License, or
# (at your option) any later version.
#
# A copy of the GNU General Public License can be found at
# http://www.gnu.org/
#

# Define constants
PATH=/bin
ProgramName=`basename $0`
ProgramOptions='p:d:t:s:T:knvh'
DefaultExtraArgs=
DefaultTSAware=1
DefaultDynBase=1
DefaultKeep=
DefaultDoNothing=
DefaultVerbose=
DefaultFileList=
DefaultSuffixes='exe|dll|so'

# Define functions
usage_string="usage: $ProgramName [-p extra_args] [-d bool] [-t bool] [-s suffix] [-T FileList | -] [-vnh]"
short_usage ()
{
    echo "$usage_string" 1>&2
    exit 1
}

long_help ()
{
    echo "$usage_string"
    echo "When invoked with no arguments, $ProgramName modifies every cygwin $DefaultSuffixes"
    echo "on the system: .exe files have their tsaware flag set, while .dll and .so files"
    echo "have their dynamicbase flag set. However, if any of [-d|-t|-s] are specified"
    echo "then ONLY the actions so specified will occur."
    echo "   -p extra_args   pass extra_args to peflags.exe"  
    echo "   -d bool         set the dynamicbase flag to 'bool' on all specified files"
    echo "   -t bool         set the tsaware flag to 'bool' on all specified files"
    echo "   -s suffix       search for all files with the specified suffix(es)"
    echo "                     default: 'exe|dll|so'"
    echo "   -T FileList     in addition to files located via the normal search and"
    echo "                   suffix list, also operate on files listed in FileList"
    echo "                   May be '-' which indicates stdin"
    echo "   -n              do not modify any files, but display the peflags commands"
    echo "   -k              keep all temporary files"
    echo "   -v              verbose mode"
    echo "   -h              show this help"
    echo "bool may be '0', '1', 'true', 'false', 'yes', or 'no'"
    echo "Note: peflagsall will NOT set the dynamicbase flag on executables, nor will"
    echo "      it set the tsaware flag on dlls. If you must do this, use peflags itself"
    exit 0
}

check_args_for_help()
{
    # we do this early, so that we can print help before
    # the process-check.
    for a in "$@"
    do
        case "$a" in
        "--help"|"-help"|"-h") long_help ;;
        esac
    done
}

to_bool_result=
to_bool()
{
    case "$1" in
    [Tt][Rr][Uu][Ee]     ) to_bool_result=1;;
    [Ff][Aa][Ll][Ss][Ee] ) to_bool_result=0;;
    [Yy][Ee][Ss]         ) to_bool_result=1;;
    [Nn][Oo]             ) to_bool_result=0;;
    [Yy]                 ) to_bool_result=1;;
    [Nn]                 ) to_bool_result=0;;
    1|0                  ) to_bool_result=$1;;
    *) echo "Bad bool argument '$1'" 1>&2
       short_usage
       exit 1
    esac
}

cleanup()
{
    if test -z "$Keep"
    then
    	rm -f "$TmpFile" "$ExeFile" "$DllFile"
    else
        echo "Saving temp files '$TmpDir/peflags_*' (may not exist)" 1>&2
    fi
    exit $ExitCode
}

verbose_only()
{
    if test -n "$Verbose"
    then
        echo "$@" 1>&2
    fi
}

# Set traps
trap cleanup 1 2 15

# Set defaults
ExtraArgs=$DefaultExtraArgs
TSAware=$DefaultTSAware
DynBase=$DefaultDynBase
Keep=$DefaultKeep
DoNothing=$DefaultDoNothing
Verbose=$DefaultVerbose
FileList=$DefaultFileList
Suffixes=$DefaultSuffixes

# Argument accumulators
ArgSuffixes=
ArgTSAware=
ArgDynBase=

# First see if caller requested help
check_args_for_help "$@"

# Verify only ash processes are running
grep -E -q -i -v '/ash(.exe)?$' /proc/[0-9]*/exename
if [ $? -eq 0 -a -z "$RebaseDebug" ]
then
    echo "$ProgramName: only ash processes are allowed during this process."
    echo "    Exit all Cygwin processes and stop all Cygwin services."
    echo "    Execute ash from Start/Run... or a cmd or command window."
    echo "    Execute '/bin/peflagsall' from ash."
    exit 2
fi

# Parse command line arguments
while getopts $ProgramOptions Option "$@"
do
    case $Option in
    p)  ExtraArgs="$OPTARG";;
    d)  to_bool "$OPTARG"; ArgDynBase=$to_bool_result;;
    t)  to_bool "$OPTARG"; ArgTSAware=$to_bool_result;;
    s)
	if test -z "$ArgSuffixes"
        then
		ArgSuffixes="$OPTARG"
	else
		ArgSuffixes="$ArgSuffixes|$OPTARG"
	fi
	;;
    n)  DoNothing=:;;
    k)  Keep=1;;
    T)  FileList="$OPTARG";;
    v)  Verbose=-v;;
    h)  long_help;;
    \?) short_usage;;
    esac
done

# Evaluate command line arguments
if test -n "$ArgSuffixes" || test -n "$ArgDynBase" || test -n "$ArgTSAware"
then
    echo "Using custom arguments; all default behavior suppressed" 1>&2
    Suffixes="$ArgSuffixes"
    DynBase="$ArgDynBase"
    TSAware="$ArgTSAware"
fi
if test -z "$Suffixes"
then
    echo "Error: must specify suffixes (-s) when default behavior is suppressed" 1>&2
    short_usage
fi


# Set temp directory
TmpDir="${TMP:-${TEMP:-/tmp}}"

# Validate temp directory
if [ ! -d "$TmpDir" ]
then
    echo "$ProgramName: '$TmpDir' is not a directory" 1>&2
    exit 2
fi
if [ ! -w "$TmpDir" ]
then
    echo "$ProgramName: '$TmpDir' is not writable" 1>&2
    exit 2
fi

# Validate user supplied file list, if necessary
if [ -n "$FileList" -a ! -r "$FileList" -a "$FileList" != - ]
then
    echo "$ProgramName: '$FileList' is not readable" 1>&2
    exit 2
fi

# Set temp files
TmpFile="$TmpDir/peflags.lst"
ExeFile="$TmpDir/peflags_exe.lst"
DllFile="$TmpDir/peflags_dll.lst"

# Create file list
find /etc/setup -name '*.lst.gz' | xargs gzip -d -c |
    grep -E "\.($Suffixes)\$" |
    sed -e '/cygwin1.dll$/d' -e '/cyglsa64.dll$/d' \
	-e 's/^/\//' -e '/ash.exe$/d' -e '/peflags.exe$/d' >"$TmpFile"

# Append user supplied file list, if any
if [ -n "$FileList" ]
then
    cat "$FileList" >>"$TmpFile"
fi

# Split file list into two groups: exe's, and everything else
# (presumably, DLLs by any name).
sed -n -e '/\.[Ee][Xx][Ee]$/p' < "$TmpFile" > "$ExeFile"
sed    -e '/\.[Ee][Xx][Ee]$/d' < "$TmpFile" > "$DllFile"
NumExes=`cat "$ExeFile" | sed -n '$='`
NumDlls=`cat "$DllFile" | sed -n '$='`

# Prepare to do exes
DoExes=
TSAwareCmd="--tsaware=$TSAware"
if test -n "$NumExes" && test "$NumExes" -gt 0
then
    if test -n "$TSAware" || test -n "$ExtraArgs"
    then
       DoExes=1
       if test -z "$TSAware"
       then
           TSAwareCmd=
       fi
    else
       verbose_only "Not processing $NumExes executables; neither -t nor -p specified"
    fi
else
    verbose_only "No executable files to process"
fi

# Prepare to do dlls
DoDlls=
DynBaseCmd="--dynamicbase=$DynBase"
if test -n "$NumDlls" && test "$NumDlls" -gt 0
then
    if test -n "$DynBase" || test -n "$ExtraArgs"
    then
       DoDlls=1
       if test -z "$DynBase"
       then
           DynBaseCmd=
       fi
    else
       verbose_only "Not processing $NumDlls dlls; neither -d nor -p specified"
    fi
else
    verbose_only "No dlls to process"
fi

if test -n "$DoExes"
then
    if test -n "$DoNothing" || test -n "$Verbose"
    then
        echo "peflags $Verbose $ExtraArgs $TSAwareCmd -T $ExeFile" 1>&2
    fi
    $DoNothing peflags $Verbose $ExtraArgs $TSAwareCmd -T "$ExeFile"
fi
ExitCode=$?

if test "$ExitCode" -eq 0
    then
    if test -n "$DoDlls" ; then
        if test -n "$DoNothing" || test -n "$Verbose"
        then
            echo "peflags $Verbose $ExtraArgs $DynBaseCmd -T $DllFile" 1>&2
        fi
        $DoNothing peflags $Verbose $ExtraArgs $DynBaseCmd -T "$DllFile"
    fi
    ExitCode=$?
fi

# Clean up
cleanup

