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

declare -r version=%VERSION%

shopt -s extglob

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

# needed files/directories
_f_functions=functions
_f_config=mkinitcpio.conf
_d_hooks=(hooks /usr/lib/initcpio/hooks /lib/initcpio/hooks)
_d_install=(install /usr/lib/initcpio/install /lib/initcpio/install)
_d_presets=mkinitcpio.d

# options and runtime data
_optmoduleroot= _optkver= _optgenimg= _optpreset=
_optcompress=
_optshowautomods=0 _optsavetree=0 _optshowmods=0
_optquiet=1 _optcolor=1
_optskiphooks=() _optaddhooks=() _modpaths=() _hooks=()
declare -A _runhooks=() _addedmodules=() _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
unset GREP_OPTIONS CDPATH

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)
   -t, --builddir <dir>         Use DIR as the temporary build directory
   -v, --verbose                Verbose output (default: no)
   -z, --compress <program>     Use an alternate compressor on the image

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

    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

    if file -Lb "$kernel" | grep -oP '(?<=version )[^ ]+'; then
        return 0
    fi

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

    return 1
}

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

    msg "Creating $compress initcpio image: %s" "$out"

    case $compress in
        xz)
            COMPRESSION_OPTIONS+=' --check=crc32'
            ;;
    esac

    cpio_opts=('-R' '0:0' '-0' '-o' '-H' 'newc')
    (( _optquiet )) && cpio_opts+=('--quiet')

    # write version stamp
    printf '%s' "$version" > "$BUILDROOT/VERSION"

    pushd "$BUILDROOT" >/dev/null
    find . -print0 |
            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

    # 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 "Preset not found: \`%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: '$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[*]}"
        "$0" "${preset_cmd[@]}"
        (( $? )) && ret=1
    done

    exit $ret
}

install_modules() {
    if (( $# == 0 )); then
        warning "No modules were added to the image. This is probably not what you want."
        return 0
    fi

    printf '%s\0' "$@" | sort -zu |
        xargs -0 cp -t "$BUILDROOT/usr/lib/modules/$KERNELVERSION/kernel"

    # unzip modules prior to recompression
    gzip -dr "$BUILDROOT/usr/lib/modules/$KERNELVERSION/kernel"

    msg "Generating module dependencies"
    install -m644 -t "$BUILDROOT/usr/lib/modules/$KERNELVERSION" \
        "$_d_kmoduledir"/modules.builtin

    # we install all modules into kernel/, making the .order file incorrect
    # for the module tree. munge it, so that we have an accurate index.
    awk -F'/' '{ print "kernel/" $NF }' \
        "$_d_kmoduledir"/modules.order >"$BUILDROOT/lib/modules/$KERNELVERSION/modules.order"

    depmod -b "$BUILDROOT" "$KERNELVERSION"

    # remove all non-binary module.* files (except devname for on-demand module loading)
    rm "$BUILDROOT/usr/lib/modules/$KERNELVERSION"/modules.!(*.bin|devname)
}

. "$_f_functions"

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

_opt_short='A:c:g:H:hk:nLMp:r:S:st:vz:'
_opt_long=('add:' 'addhooks:' 'config:' 'generate:' 'hookhelp:' 'help'
          'kernel:' 'listhooks' 'automods' 'moduleroot:' 'nocolor'
          'preset:' 'skiphooks:' 'save' 'builddir:' '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
            _optkver=$1 ;;
        -s|--save)
            _optsavetree=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 ;;
        -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
            if script=$(find_in_dirs "$1" "${_d_install[@]}"); then
                . "$script"
                if ! declare -f help >/dev/null; then
                    error "No help for hook $1"
                    exit 1
                fi
            else
                error "No hook '$1'"
                exit 1
            fi
            msg "Help for hook '$1':"
            help
            list_hookpoints "$script"
            exit 0 ;;
        -L|--listhooks)
            msg "Available hooks"
            for dir in "${_d_install[@]}"; do
                ( cd "$dir" &>/dev/null && printf '   %s\n' * )
            done | sort -u | column -c$(tput cols)
            exit 0 ;;
        -M|--automods)
            _optshowautomods=1 ;;
        -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
    # prefer terminal safe colored and bold text when tput is supported
    if tput setaf 0 &>/dev/null; then
        _color_none="$(tput sgr0)"
        _color_bold="$(tput bold)"
        _color_blue="$_color_bold$(tput setaf 4)"
        _color_green="$_color_bold$(tput setaf 2)"
        _color_red="$_color_bold$(tput setaf 1)"
        _color_yellow="$_color_bold$(tput setaf 3)"
    else
        _color_none="\e[1;0m"
        _color_bold="\e[1;1m"
        _color_blue="$_color_bold\e[1;34m"
        _color_green="$_color_bold\e[1;32m"
        _color_red="$_color_bold\e[1;31m"
        _color_yellow="$_color_bold\e[1;33m"
    fi
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)
[[ $_optpreset ]] && process_preset "$_optpreset"

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

_d_workdir=$(initialize_buildroot "$KERNELVERSION") || cleanup 1
BUILDROOT=$_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"
    _f_autodetect_hook=$(find_in_dirs 'autodetect' "${_d_install[@]}")
    . "$_f_autodetect_hook"
    build
    printf '%s\n' "${!_autodetect_cache[@]}" | sort
    cleanup 0
fi

if [[ -z $_optgenimg ]]; then
    msg "Starting dry run: %s" "$KERNELVERSION"
else
    # 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
        die "Unable to locate compression method: %s" "$_optcompress"
    fi

    msg "Starting build: %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_hook
write_image_config

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

install_modules "${_modpaths[@]}"

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

if [[ $_optgenimg ]]; then
    build_image "$_optgenimg" "$_optcompress"
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: