#!/bin/bash
#
# mkinitcpio - modular tool for building an initramfs images
#

declare -r version=19

shopt -s extglob

### globals within mkinitcpio, but not intended to be used by hooks

# needed files/directories
_f_functions=/usr/lib/initcpio/functions
_f_config=/etc/mkinitcpio.conf
_d_hooks=/etc/initcpio/hooks:/usr/lib/initcpio/hooks
_d_install=/etc/initcpio/install:/usr/lib/initcpio/install
_d_firmware=({/usr,}/lib/firmware/updates {/usr,}/lib/firmware)
_d_presets=/etc/mkinitcpio.d

# options and runtime data
_optmoduleroot= _optgenimg=
_optcompress= _opttargetdir=
_optshowautomods=0 _optsavetree=0 _optshowmods=0
_optquiet=1 _optcolor=1
_optskiphooks=() _optaddhooks=() _hooks=()  _optpreset=()
declare -A _runhooks _addedmodules _modpaths _autodetect_cache

# export a sane PATH
export PATH='/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin'

# Sanitize environment further
# GREP_OPTIONS="--color=always" will break everything
# CDPATH can affect cd and pushd
# LIBMOUNT_* options can affect findmnt and other tools
unset GREP_OPTIONS CDPATH "${!LIBMOUNT_@}"

usage() {
    cat <<EOF
mkinitcpio $version
usage: ${0##*/} [options]

  Options:
   -A, --addhooks <hooks>       Add specified hooks, comma separated, to image
   -c, --config <config>        Use alternate config file. (default: /etc/mkinitcpio.conf)
   -g, --generate <path>        Generate cpio image and write to specified path
   -H, --hookhelp <hookname>    Display help for given hook and exit
   -h, --help                   Display this message and exit
   -k, --kernel <kernelver>     Use specified kernel version (default: $(uname -r))
   -L, --listhooks              List all available hooks
   -M, --automods               Display modules found via autodetection
   -n, --nocolor                Disable colorized output messages
   -p, --preset <file>          Build specified preset from /etc/mkinitcpio.d
   -r, --moduleroot <dir>       Root directory for modules (default: /)
   -S, --skiphooks <hooks>      Skip specified hooks, comma-separated, during build
   -s, --save                   Save build directory. (default: no)
   -d, --generatedir <dir>      Write generated image into <dir>
   -t, --builddir <dir>         Use DIR as the temporary build directory
   -V, --version                Display version information and exit
   -v, --verbose                Verbose output (default: no)
   -z, --compress <program>     Use an alternate compressor on the image

EOF
}

version() {
    cat <<EOF
mkinitcpio $version
EOF
}

cleanup() {
    local err=${1:-$?}

    if [[ $_d_workdir ]]; then
        # when _optpreset is set, we're in the main loop, not a worker process
        if (( _optsavetree )) && [[ -z $_optpreset ]]; then
            printf '%s\n' "${!_autodetect_cache[@]}" > "$_d_workdir/autodetect_modules"
            msg "build directory saved in %s" "$_d_workdir"
        else
            rm -rf "$_d_workdir"
        fi
    fi

    exit $err
}

resolve_kernver() {
    local kernel=$1 offset kver

    if [[ -z $kernel ]]; then
        uname -r
        return 0
    fi

    if [[ ${kernel:0:1} != / ]]; then
        echo "$kernel"
        return 0
    fi

    if [[ ! -e $kernel ]]; then
        error "specified kernel image does not exist: \`%s'" "$kernel"
        return 1
    fi

    kver "$kernel" && return

    error "invalid kernel specified: \`%s'" "$1"

    return 1
}

hook_help() {
    local resolved script=$(PATH=$_d_install type -p "$1")

    # this will be true for broken symlinks as well
    if [[ -z $script ]]; then
        error "Hook '%s' not found" "$1"
        return 1
    fi

    if [[ -L $script ]]; then
        resolved=$(readlink -e "$script")
        msg "This hook is deprecated. See the '%s' hook" "${resolved##*/}"
        return 0
    fi

    . "$script"
    if ! declare -f help >/dev/null; then
        error "No help for hook $1"
        return 1
    fi

    msg "Help for hook '$1':"
    help

    list_hookpoints "$1"
}

hook_list() {
    local p hook resolved
    local -a paths hooklist depr
    local ss_ordinals=(¹ ² ³ ⁴ ⁵ ⁶ ⁷ ⁸ ⁹)

    IFS=: read -ra paths <<<"$_d_install"

    for path in "${paths[@]}"; do
        for hook in "$path"/*; do
            [[ -e $hook || -L $hook ]] || continue

            # handle deprecated hooks and point to replacement
            if [[ -L $hook ]]; then
                resolved=$(readlink -e "$hook")

                if [[ -z $resolved ]]; then
                    error "found broken symlink '%s'" "$hook"
                    continue
                fi

                resolved=${resolved##*/}

                if ! index_of "$resolved" "${depr[@]}"; then
                    # deprecated hook
                    depr+=("$resolved")
                    _idx=$(( ${#depr[*]} - 1 ))
                fi

                hook+=${ss_ordinals[_idx]}
            fi

            hooklist+=("${hook##*/}")
        done
    done

    msg "Available hooks"
    printf '%s\n' "${hooklist[@]}" | sort -u | column -c$(tput cols)

    if (( ${#depr[*]} )); then
        echo
        for p in "${!depr[@]}"; do
            printf "%s This hook is deprecated in favor of '%s'\n" \
                "${ss_ordinals[p]}" "${depr[p]}"
        done
    fi
}

compute_hookset() {
    local h

    for h in $HOOKS "${_optaddhooks[@]}"; do
        in_array "$h" "${_optskiphooks[@]}" && continue
        _hooks+=("$h")
    done
}

build_image() {
    local out=$1 compress=$2 errmsg=
    local -a pipesave cpio_opts

    if [[ $compress = cat ]]; then
        msg "Creating uncompressed initcpio image: %s" "$out"
    else
        msg "Creating %s-compressed initcpio image: %s" "$compress" "$out"
    fi

    case $compress in
        cat)
            unset COMPRESSION_OPTIONS
            ;;
        xz)
            COMPRESSION_OPTIONS+=' --check=crc32'
            ;;
        lz4)
            COMPRESSION_OPTIONS+=' -l'
            ;;
    esac

    cpio_opts=('-0' '-o' '-H' 'newc')
    (( _optquiet )) && cpio_opts+=('--quiet')
    if (( EUID != 0 )); then
        warning 'Not building as root, ownership cannot be preserved'
        cpio_opts+=('-R' '0:0')
    fi

    pushd "$BUILDROOT" >/dev/null
    find -mindepth 1 -printf '%P\0' |
            LANG=C bsdcpio "${cpio_opts[@]}" |
            $compress $COMPRESSION_OPTIONS > "$out"
    pipesave=("${PIPESTATUS[@]}") # save immediately
    popd >/dev/null

    if (( pipesave[0] )); then
        errmsg="find reported an error"
    elif (( pipesave[1] )); then
        errmsg="bsdcpio reported an error"
    elif (( pipesave[2] )); then
        errmsg="$compress reported an error"
    fi

    if (( _builderrors )); then
        warning "errors were encountered during the build. The image may not be complete."
    fi

    if [[ $errmsg ]]; then
        error "Image generation FAILED: %s" "$errmsg"
    elif (( _builderrors == 0 )); then
        msg "Image generation successful"
    fi
}

process_preset() (
    local preset=$1 preset_image= preset_options=
    local -a preset_mkopts preset_cmd

    if (( MKINITCPIO_PROCESS_PRESET )); then
        error "You appear to be calling a preset from a preset. This is a configuration error."
        cleanup 1
    fi

    # allow path to preset file, else resolve it in $_d_presets
    if [[ $preset != */* ]]; then
        printf -v preset '%s/%s.preset' "$_d_presets" "$preset"
    fi

    . "$preset" || die "Failed to load preset: \`%s'" "$preset"

    # Use -m and -v options specified earlier
    (( _optquiet )) || preset_mkopts+=(-v)
    (( _optcolor )) || preset_mkopts+=(-n)

    ret=0
    for p in "${PRESETS[@]}"; do
        msg "Building image from preset: $preset: '$p'"
        preset_cmd=("${preset_mkopts[@]}")

        preset_kver=${p}_kver
        if [[ ${!preset_kver:-$ALL_kver} ]]; then
            preset_cmd+=(-k "${!preset_kver:-$ALL_kver}")
        else
            warning "No kernel version specified. Skipping image \`%s'" "$p"
            continue
        fi

        preset_config=${p}_config
        if [[ ${!preset_config:-$ALL_config} ]]; then
            preset_cmd+=(-c "${!preset_config:-$ALL_config}")
        else
            warning "No configuration file specified. Skipping image \`%s'" "$p"
            continue
        fi

        preset_image=${p}_image
        if [[ ${!preset_image} ]]; then
            preset_cmd+=(-g "${!preset_image}")
        else
            warning "No image file specified. Skipping image \`%s'" "$p"
            continue
        fi

        preset_options=${p}_options
        if [[ ${!preset_options} ]]; then
            preset_cmd+=(${!preset_options}) # intentional word splitting
        fi

        msg2 "${preset_cmd[*]}"
        MKINITCPIO_PROCESS_PRESET=1 "$0" "${preset_cmd[@]}"
        (( $? )) && ret=1
    done

    exit $ret
)

. "$_f_functions"

trap 'cleanup 130' INT
trap 'cleanup 143' TERM

_opt_short='A:c:g:H:hk:nLMPp:r:S:sd:t:Vvz:'
_opt_long=('add:' 'addhooks:' 'config:' 'generate:' 'hookhelp:' 'help'
          'kernel:' 'listhooks' 'automods' 'moduleroot:' 'nocolor' 'allpresets'
          'preset:' 'skiphooks:' 'save' 'generatedir:' 'builddir:' 'version' 'verbose' 'compress:')

parseopts "$_opt_short" "${_opt_long[@]}" -- "$@" || exit 1
set -- "${OPTRET[@]}"
unset _opt_short _opt_long OPTRET

while :; do
    case $1 in
        # --add remains for backwards compat
        -A|--add|--addhooks)
            shift
            IFS=, read -r -a add <<< "$1"
            _optaddhooks+=("${add[@]}")
            unset add
            ;;
        -c|--config)
            shift
            _f_config=$1
            ;;
        -k|--kernel)
            shift
            KERNELVERSION=$1
            ;;
        -s|--save)
            _optsavetree=1
            ;;
        -d|--generatedir)
            shift
            _opttargetdir=$1
            ;;
        -g|--generate)
            shift
            [[ -d $1 ]] && die "Invalid image path -- must not be a directory"
            if ! _optgenimg=$(readlink -f "$1") || [[ ! -e ${_optgenimg%/*} ]]; then
                die "Unable to write to path: \`%s'" "$1"
            fi
            ;;
        -h|--help)
            usage
            cleanup 0
            ;;
        -V|--version)
            version
            cleanup 0
            ;;
        -p|--preset)
            shift
            _optpreset+=("$1")
            ;;
        -n|--nocolor)
            _optcolor=0
            ;;
        -v|--verbose)
            _optquiet=0
            ;;
        -S|--skiphooks)
            shift
            IFS=, read -r -a skip <<< "$1"
            _optskiphooks+=("${skip[@]}")
            unset skip
            ;;
        -H|--hookhelp)
            shift
            hook_help "$1"
            exit
            ;;
        -L|--listhooks)
            hook_list
            exit 0
            ;;
        -M|--automods)
            _optshowautomods=1
            ;;
        -P|--allpresets)
            _optpreset=("$_d_presets"/*.preset)
            [[ -e ${_optpreset[0]} ]] || die "No presets found in $_d_presets"
            ;;
        -t|--builddir)
            shift
            export TMPDIR=$1
            ;;
        -z|--compress)
            shift
            _optcompress=$1
            ;;
        -r|--moduleroot)
            shift
            _optmoduleroot=$1
            ;;
        --)
            shift
            break 2
            ;;
    esac
    shift
done

if [[ -t 1 ]] && (( _optcolor )); then
    try_enable_color
fi

# insist that /proc and /dev be mounted (important for chroots)
# NOTE: avoid using mountpoint for this -- look for the paths that we actually
# use in mkinitcpio. Avoids issues like FS#26344.
[[ -e /proc/self/mountinfo ]] || die "/proc must be mounted!"
[[ -e /dev/fd ]] || die "/dev must be mounted!"

# use preset $_optpreset (exits after processing)
if (( ${#_optpreset[*]} )); then
    map process_preset "${_optpreset[@]}"
    exit
fi

if [[ $KERNELVERSION != 'none' ]]; then
    KERNELVERSION=$(resolve_kernver "$KERNELVERSION") || cleanup 1
    _d_kmoduledir=$_optmoduleroot/lib/modules/$KERNELVERSION
    [[ -d $_d_kmoduledir ]] || die "'$_d_kmoduledir' is not a valid kernel module directory"
fi

_d_workdir=$(initialize_buildroot "$KERNELVERSION" $_opttargetdir) || cleanup 1
BUILDROOT=${_opttargetdir:-$_d_workdir/root}

. "$_f_config" || die "Failed to read configuration \`%s'" "$_f_config"

# after returning, hooks are populated into the array '_hooks'
# HOOKS should not be referenced from here on
compute_hookset

if (( ${#_hooks[*]} == 0 )); then
    die "Invalid config: No hooks found"
fi

if (( _optshowautomods )); then
    msg "Modules autodetected"
    PATH=$_d_install . 'autodetect'
    build
    printf '%s\n' "${!_autodetect_cache[@]}" | sort
    cleanup 0
fi

if [[ $_optgenimg ]]; then
    # check for permissions. if the image doesn't already exist,
    # then check the directory
    if [[ ( -e $_optgenimg && ! -w $_optgenimg ) ||
            ( ! -d ${_optgenimg%/*} || ! -w ${_optgenimg%/*} ) ]]; then
        die 'Unable to write to %s' "$_optgenimg"
    fi

    _optcompress=${_optcompress:-${COMPRESSION:-gzip}}
    if ! type -P "$_optcompress" >/dev/null; then
        warning "Unable to locate compression method: %s" "$_optcompress"
        _optcompress=cat
    fi

    msg "Starting build: %s" "$KERNELVERSION"
elif [[ $_opttargetdir ]]; then
    msg "Starting build: %s" "$KERNELVERSION"
else
    msg "Starting dry run: %s" "$KERNELVERSION"
fi

# set functrace and trap to catch errors in add_* functions
declare -i _builderrors=0
set -o functrace
trap '(( $? )) && [[ $FUNCNAME = add_* ]] && (( ++_builderrors ))' RETURN

# prime the _addedmodules list with the builtins for this kernel
if [[ -r $_d_kmoduledir/modules.builtin ]]; then
    while IFS=/ read -a path; do
        modname=${path[-1]%.ko}
        _addedmodules["${modname//-/_}"]=2
    done <"$_d_kmoduledir/modules.builtin"
    unset modname path
fi

map run_build_hook "${_hooks[@]}" || (( ++_builderrors ))

# process config file
parse_config "$_f_config"

# switch out the error handler to catch all errors
trap -- RETURN
trap '(( ++_builderrors ))' ERR
set -o errtrace

install_modules "${!_modpaths[@]}"

# unset errtrace and trap
set +o functrace
set +o errtrace
trap -- ERR

# this is simply a nice-to-have -- it doesn't matter if it fails.
ldconfig -r "$BUILDROOT" &>/dev/null

if [[ $_optgenimg ]]; then
    build_image "$_optgenimg" "$_optcompress"
elif [[ $_opttargetdir ]]; then
    msg "Build complete."
else
    msg "Dry run complete, use -g IMAGE to generate a real image"
fi

cleanup $(( !!_builderrors ))

# vim: set ft=sh ts=4 sw=4 et:
