#!/bin/bash
#
## Wizard Kit: Live Linux Build Tool

set -o errexit
set -o errtrace
set -o nounset
set -o pipefail

# Prep
DATE="$(date +%F)"
DATETIME="$(date +%F_%H%M)"
ROOT_DIR="$(realpath $(dirname "$0"))"
BUILD_DIR="$ROOT_DIR/BUILD_LINUX"
LIVE_DIR="$BUILD_DIR/live"
LOG_DIR="$BUILD_DIR/logs"
OUT_DIR="$ROOT_DIR/OUT_LINUX"
REPO_DIR="$BUILD_DIR/repo"
SKEL_DIR="$LIVE_DIR/airootfs/etc/skel"
TEMP_DIR="$BUILD_DIR/temp"
MIRRORLIST_SOURCE='https://www.archlinux.org/mirrorlist/?country=US&protocol=http&protocol=https&ip_version=4&use_mirror_status=on'
if which nano >/dev/null 2>&1; then
    EDITOR=nano
elif which vim >/dev/null 2>&1; then
    EDITOR=vim
else
    EDITOR=vi
fi
if [ ! -z ${SUDO_USER+x} ]; then
    REAL_USER="$SUDO_USER"
fi

function ask() {
    while :; do
        read -p "$1 [Y/N] " -r answer
        if echo "$answer" | egrep -iq '^(y|yes|sure)$'; then
            return 0
        elif echo "$answer" | egrep -iq '^(n|no|nope)$'; then
            return 1
        fi
    done
}
        
function cleanup() {
    for d in "$TEMP_DIR" "$LIVE_DIR"; do
        if [[ -d "$d" ]]; then
            if ask "Remove: ${d}?"; then
                rm -Rf "$d"
            fi
        fi
    done
    mkdir -p "$BUILD_DIR" 2>/dev/null
    mkdir -p "$LOG_DIR" 2>/dev/null
    mkdir -p "$OUT_DIR" 2>/dev/null
    mkdir -p "$TEMP_DIR" 2>/dev/null
}

function fix_kit_permissions() {
    # GitHub zip archives don't preserve the correct permissions
    for d in .bin .cbin .kit_items .linux_items .pe_items Images; do
        find "$ROOT_DIR/$d" -type d -exec chmod 755 "{}" \;
    done
}

function load_settings() {
    # Check if settings already loaded
    ## Code based on StackOverflow Q&A
    ## Question:    https://stackoverflow.com/q/3601515
    ## Asked by:    https://stackoverflow.com/users/260127/prosseek
    ## Edited by:   https://stackoverflow.com/users/3924118/nbro
    ## Answer:      https://stackoverflow.com/a/13864829
    ## Answer by:   https://stackoverflow.com/users/1633643/lionel
    ## Answer edit: https://stackoverflow.com/users/-1/community
    _main_path="$BUILD_DIR/main.py"
    if [ ! -z ${KIT_NAME_FULL+x} ]; then
        # KIT_NAME_FULL is set
        return 0 # Skip loading settings from main.py
    fi

    if [[ "${1:-}" == "--edit" ]]; then
        # Copy settings
        if [[ ! -e "$BUILD_DIR/main.py" ]] || ask "Overwrite main.py?"; then
            cp -bv "$ROOT_DIR/.bin/Scripts/settings/main.py" "$BUILD_DIR/main.py"
            dos2unix "$BUILD_DIR/main.py"
        fi

        # Edit settings
        read -p "Press Enter to open settings... " -r
        "$EDITOR" "$BUILD_DIR/main.py"
    else
        # Load settings from $LIVE_DIR
        _main_path="$LIVE_DIR/airootfs/usr/local/bin/settings/main.py"
    fi

    # Load settings
    while read line; do
        if echo "$line" | egrep -q "^\w+='"; then
            line="$(echo "$line" | sed -r 's/[\r\n]+//')"
            eval "$line"
        fi
    done < "$_main_path"
}

function copy_live_env() {
    echo "Copying Archlinux files..."
    rsync -aI /usr/share/archiso/configs/releng/ "$LIVE_DIR/"

    # Remove items
    rm "$LIVE_DIR/airootfs/etc/systemd/scripts/choose-mirror"
    rmdir "$LIVE_DIR/airootfs/etc/systemd/scripts" --ignore-fail-on-non-empty
    rm "$LIVE_DIR/airootfs/etc/systemd/system/choose-mirror.service"
    rm "$LIVE_DIR/airootfs/etc/systemd/system/etc-pacman.d-gnupg.mount"
    rm "$LIVE_DIR/airootfs/etc/systemd/system/pacman-init.service"
    rm "$LIVE_DIR/airootfs/etc/udev/rules.d/81-dhcpcd.rules"
    rmdir "$LIVE_DIR/airootfs/etc/udev/rules.d" --ignore-fail-on-non-empty
    rmdir "$LIVE_DIR/airootfs/etc/udev" --ignore-fail-on-non-empty
    rm "$LIVE_DIR/isolinux"/*.cfg
    rm "$LIVE_DIR/syslinux"/*.cfg "$LIVE_DIR/syslinux"/*.png
    
    # Add items
    rsync -aI "$ROOT_DIR/.linux_items/include/" "$LIVE_DIR/"
    if [[ "${1:-}" != "--minimal" ]]; then
      rsync -aI "$ROOT_DIR/.linux_items/include_x/" "$LIVE_DIR/"
    fi
    mkdir -p "$LIVE_DIR/airootfs/usr/local/bin"
    rsync -aI "$ROOT_DIR/.bin/Scripts/" "$LIVE_DIR/airootfs/usr/local/bin/"
    cp -a "$BUILD_DIR/main.py" "$LIVE_DIR/airootfs/usr/local/bin/settings/"
}

function run_elevated() {
    prog="$1"
    shift
    if which sudo >/dev/null 2>&1; then
        if ! sudo "$prog" $*; then
            echo "ERROR: Failed to run '$prog'"
            if ask "Retry?"; then
                sudo "$prog" $*
            fi
        fi
    else
        echo -n "Root "
        if ! su -c "export REAL_USER=$USER && '$prog' $*"; then
            echo "ERROR: Failed to run '$prog'"
            if ask "Retry?"; then
                su -c "export REAL_USER=$USER && '$prog' $*"
            fi
        fi
    fi
}

function update_live_env() {
    hostname="$(echo "${KIT_NAME_SHORT}-linux" | tr "[:upper:]" "[:lower:]")"
    username="tech"
    label="${KIT_NAME_SHORT}_LINUX"

    # Boot config (legacy)
    mkdir -p "$LIVE_DIR/arch"
    cp "$ROOT_DIR/Images/Pxelinux.png" "$LIVE_DIR/arch/pxelinux.png"
    cp "$ROOT_DIR/Images/Syslinux.png" "$LIVE_DIR/arch/syslinux.png"
    sed -i -r "s/_+/$KIT_NAME_FULL/" "$LIVE_DIR/syslinux/wk_head.cfg"
    mkdir -p "$TEMP_DIR" 2>/dev/null
    curl -Lo "$TEMP_DIR/wimboot.zip" "http://git.ipxe.org/releases/wimboot/wimboot-latest.zip"
    7z e "$TEMP_DIR/wimboot.zip" -o"$LIVE_DIR/arch/boot" 'wimboot*/LICENSE.txt' 'wimboot*/README.txt' 'wimboot*/wimboot'

    # Boot config (UEFI)
    mkdir -p "$LIVE_DIR/EFI/boot"
    cp  "/usr/share/refind/refind_x64.efi" "$LIVE_DIR/EFI/boot/bootx64.efi"
    cp "$ROOT_DIR/Images/rEFInd.png" "$LIVE_DIR/EFI/boot/rEFInd.png"
    rsync -aI "/usr/share/refind/drivers_x64/" "$LIVE_DIR/EFI/boot/drivers_x64/"
    rsync -aI "/usr/share/refind/icons/" "$LIVE_DIR/EFI/boot/icons/" --exclude "/usr/share/refind/icons/svg"
    sed -i "s/%ARCHISO_LABEL%/${label}/" "$LIVE_DIR/EFI/boot/refind.conf"

    # Customize_airootfs.sh
    sed -i -r 's/set -e -u/set -o errexit\nset -o errtrace\nset -o nounset\nset -o pipefail/' "$LIVE_DIR/airootfs/root/customize_airootfs.sh"

    # Memtest86
    mkdir -p "$LIVE_DIR/EFI/memtest86/Benchmark"
    mkdir -p "$TEMP_DIR/memtest86"
    curl -Lo "$TEMP_DIR/memtest86/memtest86-usb.zip" "https://www.memtest86.com/downloads/memtest86-usb.zip"
    7z e "$TEMP_DIR/memtest86/memtest86-usb.zip" -o"$TEMP_DIR/memtest86" "memtest86-usb.img"
    7z e "$TEMP_DIR/memtest86/memtest86-usb.img" -o"$TEMP_DIR/memtest86" "MemTest86.img"
    7z x "$TEMP_DIR/memtest86/MemTest86.img" -o"$TEMP_DIR/memtest86"
    rm "$TEMP_DIR/memtest86/EFI/BOOT/BOOTIA32.efi"
    mv "$TEMP_DIR/memtest86/EFI/BOOT/BOOTX64.efi" "$LIVE_DIR/EFI/memtest86/memtestx64.efi"
    mv "$TEMP_DIR/memtest86/EFI/BOOT"/* "$LIVE_DIR/EFI/memtest86"/
    mv "$TEMP_DIR/memtest86/help"/* "$LIVE_DIR/EFI/memtest86"/
    mv "$TEMP_DIR/memtest86/license.rtf" "$LIVE_DIR/EFI/memtest86"/

    # build.sh
    if ! grep -iq 'wizardkit additions' "$LIVE_DIR/build.sh"; then
        sed -i -r 's/^(run_once make_iso)$/# wizardkit additions\n\1/' "$LIVE_DIR/build.sh"
        sed -i '/# wizardkit additions/r .linux_items/build_additions.txt' "$LIVE_DIR/build.sh"
    fi

    # Hostname
    echo "$hostname" > "$LIVE_DIR/airootfs/etc/hostname"
    echo "127.0.1.1  $hostname.localdomain  $hostname" >> "$LIVE_DIR/airootfs/etc/hosts"
    
    # Live packages
    while read -r p; do
        sed -i "/$p/d" "$LIVE_DIR/packages.x86_64"
    done < "$ROOT_DIR/.linux_items/packages/live_remove"
    cat "$ROOT_DIR/.linux_items/packages/live_add" >> "$LIVE_DIR/packages.x86_64"
    if [[ "${1:-}" == "--minimal" ]]; then
        cat "$ROOT_DIR/.linux_items/packages/live_add_min" >> "$LIVE_DIR/packages.x86_64"
    else
        cat "$ROOT_DIR/.linux_items/packages/live_add_x" >> "$LIVE_DIR/packages.x86_64"
    fi
    echo "[custom]" >> "$LIVE_DIR/pacman.conf"
    echo "SigLevel = Optional TrustAll" >> "$LIVE_DIR/pacman.conf"
    echo "Server = file://$REPO_DIR" >> "$LIVE_DIR/pacman.conf"
    echo "" >> "$LIVE_DIR/pacman.conf"

    # Mirrors
    sed -i -r 's/^(.*mirrorlist.*)$/#NOPE#\1/' "$LIVE_DIR/airootfs/root/customize_airootfs.sh"
    echo "curl -o '/etc/pacman.d/mirrorlist' '$MIRRORLIST_SOURCE'" >> "$LIVE_DIR/airootfs/root/customize_airootfs.sh"
    echo "sed -i 's/#Server/Server/g' /etc/pacman.d/mirrorlist" >> "$LIVE_DIR/airootfs/root/customize_airootfs.sh"
    
    # MOTD
    sed -i -r "s/_+/$KIT_NAME_FULL Linux Environment/" "$LIVE_DIR/airootfs/etc/motd"

    # Oh My ZSH
    git clone --depth=1 https://github.com/robbyrussell/oh-my-zsh.git "$SKEL_DIR/.oh-my-zsh"
    rm -Rf "$SKEL_DIR/.oh-my-zsh/.git"
    curl -o "$SKEL_DIR/.oh-my-zsh/themes/lean.zsh-theme" https://raw.githubusercontent.com/miekg/lean/master/prompt_lean_setup
    
    if [[ "${1:-}" != "--minimal" ]]; then
      # Openbox theme
      git clone --depth=1 https://github.com/addy-dclxvi/Openbox-Theme-Collections.git "$TEMP_DIR/ob-themes"
      mkdir -p "$LIVE_DIR/airootfs/usr/share/themes"
      cp -a "$TEMP_DIR/ob-themes/Triste-Orange" "$LIVE_DIR/airootfs/usr/share/themes/"
    fi

    # Services
    sed -i -r 's/^(.*pacman-init.*)$/#NOPE#\1/' "$LIVE_DIR/airootfs/root/customize_airootfs.sh"
    sed -i -r 's/^(.*choose-mirror.*)$/#NOPE#\1/' "$LIVE_DIR/airootfs/root/customize_airootfs.sh"

    # Shutdown stall fix
    echo "sed -i -r 's/^.*(DefaultTimeoutStartSec)=.*$/\1=15s/' /etc/systemd/system.conf" >> "$LIVE_DIR/airootfs/root/customize_airootfs.sh"
    echo "sed -i -r 's/^.*(DefaultTimeoutStopSec)=.*$/\1=15s/' /etc/systemd/system.conf" >> "$LIVE_DIR/airootfs/root/customize_airootfs.sh"
    
    # SSH
    mkdir -p "$SKEL_DIR/.ssh"
    ssh-keygen -b 4096 -C "$username@$hostname" -N "" -f "$SKEL_DIR/.ssh/id_rsa"
    echo 'rm /root/.ssh/id*' >> "$LIVE_DIR/airootfs/root/customize_airootfs.sh"
    echo 'rm /root/.zlogin' >> "$LIVE_DIR/airootfs/root/customize_airootfs.sh"
    sed -i -r '/.*PermitRootLogin.*/d' "$LIVE_DIR/airootfs/root/customize_airootfs.sh"
    echo "sed -i -r '/.*PermitRootLogin.*/d' /etc/ssh/sshd_config" >> "$LIVE_DIR/airootfs/root/customize_airootfs.sh"
    cp "$ROOT_DIR/.linux_items/authorized_keys" "$SKEL_DIR/.ssh/authorized_keys"
    
    # Root user
    echo "echo 'root:$ROOT_PASSWORD' | chpasswd" >> "$LIVE_DIR/airootfs/root/customize_airootfs.sh"

    # Sudo
    echo "echo '%wheel ALL=(ALL) NOPASSWD: ALL' >> /etc/sudoers" >> "$LIVE_DIR/airootfs/root/customize_airootfs.sh"

    # Tech user
    echo "groupadd -r autologin" >> "$LIVE_DIR/airootfs/root/customize_airootfs.sh"
    echo "useradd -m -s /bin/zsh -G autologin,power,storage,wheel -U $username" >> "$LIVE_DIR/airootfs/root/customize_airootfs.sh"
    echo "echo '$username:$TECH_PASSWORD' | chpasswd" >> "$LIVE_DIR/airootfs/root/customize_airootfs.sh"

    # Tech user autologin
    mkdir -p "$LIVE_DIR/airootfs/etc/systemd/system/getty@tty1.service.d"
    echo "[Service]" > "$LIVE_DIR/airootfs/etc/systemd/system/getty@tty1.service.d/autologin.conf"
    echo "ExecStart=" >> "$LIVE_DIR/airootfs/etc/systemd/system/getty@tty1.service.d/autologin.conf"
    echo "ExecStart=-/sbin/agetty --autologin $username --noclear %I 38400 linux" >> "$LIVE_DIR/airootfs/etc/systemd/system/getty@tty1.service.d/autologin.conf"

    # Timezone
    echo "ln -sf '/usr/share/zoneinfo/$LINUX_TIME_ZONE' '/etc/localtime'" >> "$LIVE_DIR/airootfs/root/customize_airootfs.sh"
    echo 'sed -i "s/#FallbackNTP/NTP/" /etc/systemd/timesyncd.conf' >> "$LIVE_DIR/airootfs/root/customize_airootfs.sh"

    # udevil fix
    echo "mkdir /media" >> "$LIVE_DIR/airootfs/root/customize_airootfs.sh"

    if [[ "${1:-}" != "--minimal" ]]; then
      # VNC password
      echo "mkdir '/home/$username/.vnc'" >> "$LIVE_DIR/airootfs/root/customize_airootfs.sh"
      echo "echo '$TECH_PASSWORD' | vncpasswd -f > '/home/$username/.vnc/passwd'" >> "$LIVE_DIR/airootfs/root/customize_airootfs.sh"

      # Wallpaper
      mkdir -p "$LIVE_DIR/airootfs/usr/share/wallpaper"
      cp "$ROOT_DIR/Images/Linux.png" "$LIVE_DIR/airootfs/usr/share/wallpaper/burned.in"
    fi

    # WiFi
    cp "$ROOT_DIR/.linux_items/known_networks" "$LIVE_DIR/airootfs/root/known_networks"
    echo "add-known-networks --user=$username" >> "$LIVE_DIR/airootfs/root/customize_airootfs.sh"
}

function update_repo() {
    if [[ "$EUID" -eq 0 ]]; then
        echo "This section not meant to be run as root."
        echo "Aborted."
        exit 1
    fi
    echo "Updating custom repo..."

    # Archive current files
    if [[ -d "$REPO_DIR" ]]; then
        mkdir -p "$BUILD_DIR/Archive" 2>/dev/null
        archive="$BUILD_DIR/Archive/$(date "+%F_%H%M%S")"
        mv -bv "$REPO_DIR" "$archive"
    fi
    sleep 1s

    # Build custom repo packages
    mkdir -p "$REPO_DIR" 2>/dev/null
    mkdir -p "$TEMP_DIR" 2>/dev/null
    pushd "$TEMP_DIR" >/dev/null
    while read -r p; do
        echo "Building: $p"
        curl -LsfO https://aur.archlinux.org/cgit/aur.git/snapshot/$p.tar.gz
        tar xf $p.tar.gz
        pushd $p >/dev/null
        makepkg -d
        popd >/dev/null
        mv -n $p/*xz "$REPO_DIR"/
    done < "$ROOT_DIR/.linux_items/packages/aur"
    popd >/dev/null

    # Build custom repo database
    pushd "$REPO_DIR" >/dev/null
    repo-add custom.db.tar.gz *xz
    popd >/dev/null
}

function install_deps() {
    echo "Installing dependencies..."
    packages=
    while read -r line; do
        packages="$packages $line"
    done < "$ROOT_DIR/.linux_items/packages/dependencies"
    run_elevated pacman -Syu --needed --noconfirm $packages
}

function build_iso() {
    if [[ "$EUID" -ne 0 ]]; then
        echo "This section is meant to be run as root."
        echo "Aborted."
        exit 1
    fi
    
    # Set permissions
    echo "Setting permissions..."
    chown root:root "$LIVE_DIR" -R
    chmod 700 "$LIVE_DIR/airootfs/etc/skel/.ssh"
    chmod 600 "$LIVE_DIR/airootfs/etc/skel/.ssh/id_rsa"

    # Removing cached (and possibly outdated) custom repo packages
    for package in $(cat "$ROOT_DIR/.linux_items/packages/aur"); do
        for p in /var/cache/pacman/pkg/*${package}*; do
            if [[ -f "${p}" ]]; then
                rm "${p}"
            fi
        done
    done
    
    # Build ISO
    prefix="${KIT_NAME_SHORT}-Linux"
    label="${KIT_NAME_SHORT}_LINUX"
    "$LIVE_DIR/build.sh" -N "$prefix" -V "$DATE" -L "$label" -w "$TEMP_DIR/Linux" -o "$OUT_DIR" -v | tee -a "$LOG_DIR/$DATETIME.log"

    # Cleanup
    echo "Removing temp files..."
    rm "$TEMP_DIR/Linux" -Rf | tee -a "$LOG_DIR/$DATETIME.log"

    echo "Reverting permissions..."
    chown $REAL_USER:$REAL_USER "$LIVE_DIR" -R
    chown $REAL_USER:$REAL_USER "$OUT_DIR" -R
}

function build_full() {
    if [[ "$EUID" -eq 0 ]]; then
        echo "This section not meant to be run as root."
        echo "Aborted."
        exit 1
    fi
    
    # Prep for build
    cleanup
    fix_kit_permissions
    install_deps
    load_settings --edit
    update_repo
    copy_live_env
    update_live_env
    # Rerun script as root to start Archiso build process
    run_elevated "$(realpath "$0")" --build-iso
}

# Check input
case ${1:-} in
    -b|--build-full)
        build_full
        echo Done
        ;;

    -f|--fix-perms)
        fix_kit_permissions
        echo Done
        ;;

    -i|--install-deps)
        install_deps
        echo Done
        ;;

    -m|--prep-minimal-env)
        load_settings --edit
        copy_live_env --minimal
        update_live_env --minimal
        echo Done
        ;;

    -o|--build-iso)
        load_settings
        build_iso
        echo Done
        ;;
    
    -p|--prep-live-env)
        load_settings --edit
        copy_live_env
        update_live_env
        echo Done
        ;;

    -u|--update-repo)
        update_repo
        echo Done
        ;;

    *)
        echo "Usage: $(basename "$0") [OPTIONS]"
        echo ""
        echo "Options:"
        echo "    -b --build-full       Perform all tasks and build iso"
        echo "    -h --help             Show usage"
        echo ""
        echo "Advanced options:"
        echo "    -f --fix-perms        Fix folder permissions"
        echo "    -i --install-deps     Install build dependencies"
        echo "    -m --prep-minimal-env Prep live & airootfs folders (minimal packages)"
        echo "    -o --build-iso        Build ISO (using current setup)"
        echo "    -p --prep-live-env    Prep live & airootfs folders"
        echo "    -u --update-repo      Update custom pacman repo"
        ;;
esac