#!/bin/bash

#########################
# Shared constants
export GREEN='\033[0;32m'
export RED='\033[0;31m'
export WHITE='\033[0;37m'
export NORMAL='\033[0m'
#########################


############################################
# Default values for configuration variables

export default_imgfile=arch.img
export default_architecture=x86_64
export default_disk_size=8G
export default_memory=2048
export default_hostname=qemu
export default_root_password=root
export default_username=user
export default_user_password=password
export default_timezone=UTC
export default_locale=en_US.UTF-8
export default_img_download_page='http://mirror.cj2.nl/archlinux/iso/latest/'
export isofile=archlinux-latest.iso
############################################


################
# install script

# Shared variables
imgfile=
architecture=
disk_size=
memory=
hostname=
root_password=
username=
user_password=
timezone=
locale="$default_locale"
img_download_page="$default_img_download_page"

function install_usage() {
    echo "Install an Arch Linux system on a QEMU disk image"
    echo
    echo "Usage: $(basename "$0") install [-o <output-disk-image>] [-a <architecture>] [-s <disk-size>]"
    echo -e "\t[-m <memory>] [-h <hostname>] [-P <root-password>] [-u <non-root-username>]"
    echo -e "\t[-p <non-root-user-password>] [-z <timezone>] [-l <locale>] [-M <arch-mirror-url>]"
    echo
    echo -e "-o\t<output-disk-image>\t\tPath of the output disk image (default: ./arch.img)"
    echo -e "-a\t<architecture>\t\t\tTarget architecture (default: x86_64)"
    echo -e "-s\t<disk-size>\t\t\tDisk size (default: 8G)"
    echo -e "-m\t<memory>\t\t\tRAM size in KB (default: 2048)"
    echo -e "-h\t<hostname>\t\t\tVM hostname (default: qemu)"
    echo -e "-P\t<root-password>\t\t\tRoot password. If not specified it will be prompted"
    echo -e "-u\t<non-root-username>\t\tUsername for the main non-root user"
    echo -e "-p\t<non-root-user-password>\tPassword for the non-root user. If not specified it will be prompted"
    echo -e "-z\t<timezone>\t\t\tSystem timezone (default: UTC)"
    echo -e "-l\t<locale>\t\t\tSystem locale (default: en_US.UTF-8)"
    echo -e "-M\t<arch-mirror-url>\t\tArch Linux download mirror URL (default: http://mirror.cj2.nl/archlinux/iso/latest/)"
    echo -e "\t\t\t\t\tConsult https://archlinux.org/download/ for a full list of the available download mirrors."
    echo
    echo "If you want to install an extra list of packages besides the default ones, then"
    echo "specify them in a file named PKGLIST in the same directory as the disk image file."
    echo
    echo "If you want to run a custom post-installation script after the core system has been"
    echo "installed, then create a custom script named post-install.sh in the same directory as"
    echo "the disk image file".
    exit 1
}

function download_latest_arch_iso() {
    latest_iso=$(
        curl -s "$img_download_page" |
        grep -e '<a href="archlinux-.*\.iso">' |
        head -1 |
        sed -r -e 's/^.*<a href="(archlinux-.*\.iso)">.*/\1/'
    )

    if [ ! -f "$latest_iso" ]; then
        echo -e "${GREEN}Downloading the latest Arch Linux ISO image${NORMAL}"
        rm -f archlinux*.iso
        curl -o "${latest_iso}" "${img_download_page}/${latest_iso}"
        ln -sf "$latest_iso" "$isofile"
    else
        echo -e "${GREEN}Latest Arch Linux image already downloaded${NORMAL}"
    fi
}

function create_disk_image() {
    # Create a backing COW image
    # qemu-img create -o backing_file="$imgfile",backing_fmt=raw -f qcow2 img1.cow

    if [ ! -f "$imgfile" ]; then
        echo -e "${GREEN}Creating base disk image${NORMAL}"
        qemu-img create -f raw "$imgfile" $disk_size
    else
        echo -e "${RED}The base disk image file $imgfile already exists.${NORMAL}"
        echo -e "${RED}Remove it and re-run this script if you intend to run a new fresh installation.${NORMAL}"
    fi
}

function _get_ssh_key_name() {
    keyfiles=$(cat <<EOF
id_rsa
id_dsa
id_ecdsa
id_ed25519
id_xmss
EOF
)

    while IFS= read -r keyfile; do
        k="$HOME/.ssh/$keyfile"
        if [ -f "$k" ]; then
            echo "$k"
            return
        fi
    done <<< "$keyfiles"

    echo -e "${RED}Can't find a valid SSH key under $HOME/.ssh${NORMAL}"
    exit 1
}

function install_os() {
    ssh_keyfile="$(_get_ssh_key_name)"
    ssh_pubkey="$(cat "$ssh_keyfile.pub")"
    ssh_privkey="$(cat "$ssh_keyfile")"
    imgdir="$(cd "$(dirname "$imgfile")" && pwd)"
    logfile="$imgdir/install.log"
    pkgfile="$imgdir/PKGLIST"
    post_install_script="$imgdir/post-install.sh"
    post_install=
    packages=$(cat <<EOF
sudo
curl
wget
dhcpcd
net-tools
netctl
openssh
EOF
)

    if [ -f "$pkgfile" ]; then
        packages="$packages
$(cat "$pkgfile")"
    fi

    [ -f "$post_install_script" ] && post_install="$(cat "$post_install_script")"
    echo -e "${GREEN}Installing base operating system${NORMAL}"
    echo -e "\tISO image: $isofile"
    echo -e "\tDisk file: $imgfile"
    echo -e "\tLog file: $logfile"

    echo "--- Log started at $(date)" > "$logfile"
    expect <<EOF | tee -a "$logfile"
        set prompt "*@archiso*~*#* "
        set chroot_prompt "*root@archiso* "
        set timeout -1
        spawn qemu-system-$architecture \
            -cdrom "$isofile" \
            -boot d \
            -cpu host \
            -enable-kvm \
            -m $memory \
            -smp 2 \
            -nographic \
            -drive file=$imgfile,format=raw

        match_max 100000

        # Pass the console boot options
        # and wait for the system to boot
        expect "*Arch Linux install*"
        send -- "\t"
        expect "*archisobasedir*"
        send -- " console=ttyS0,38400\r"
        expect "archiso login: "
        send -- "root\r"
        expect \$prompt

        # Partition the disk
        send -- "fdisk /dev/sda\r"
        expect "Command (m for help): "
        send -- "n\r"
        expect "Select (default p): "
        send -- "p\r"
        expect "Partition number (1-4, default 1): "
        send -- "\r"
        expect "First sector*: "
        send -- "\r"
        expect "Last sector*: "
        send -- "\r"
        expect "Command (m for help): "
        send -- "a\r"
        expect "Command (m for help): "
        send -- "w\r"
        expect \$prompt

        # Create and mount the filesystem
        send -- "mkfs.ext4 /dev/sda1\r"
        expect \$prompt
        send -- "mount /dev/sda1 /mnt\r"
        expect \$prompt

        # Install the system
        send -- "pacstrap /mnt base linux linux-firmware archlinux-keyring\r"
        expect \$prompt

        # Generate the fstab file
        send -- "genfstab -U /mnt >> /mnt/etc/fstab\r"
        expect \$prompt

        # chroot to the newly installed system
        send -- "arch-chroot /mnt\r"
        expect \$chroot_prompt

        # Set the timezone and sync the clock
        send -- "ln -sf /usr/share/zoneinfo/$timezone /etc/localtime\r"
        expect \$chroot_prompt
        send -- "hwclock --systohc\r"
        expect \$chroot_prompt

        # Configure localization
        send -- "echo $locale UTF-8 >> /etc/locale.gen\r"
        expect \$chroot_prompt
        send -- "locale-gen\r"
        expect \$chroot_prompt
        send -- "echo LANG=en_US.UTF-8 > /etc/locale.conf\r"
        expect \$chroot_prompt

        # Set the hostname
        send -- "echo $hostname > /etc/hostname\r"
        expect \$chroot_prompt

        # Generate /etc/hosts
        send -- "echo -e '127.0.0.1  localhost\\n::1  localhost' >> /etc/hosts\r"
        expect \$chroot_prompt

        # Generate the initpcio
        send -- "mkinitcpio -P\r"
        expect \$chroot_prompt

        # Update the keyring
        # This may currently currently take a long time
        # see https://www.reddit.com/r/archlinux/comments/rbjbcr/pacman_keyring_update_taking_too_long/
        send -- "pacman-key --init\r"
        expect \$chroot_prompt
        send -- "pacman-key --populate archlinux\r"
        expect \$chroot_prompt
        send -- "pacman-key --refresh-keys\r"
        expect \$chroot_prompt

        # As a workaround, you can temporarily disable signature check on pacman.conf
        #send -- "cp /etc/pacman.conf /etc/pacman.conf.orig\r"
        #expect \$chroot_prompt
        #send -- "sed -i /etc/pacman.conf -r -e 's/^(SigLevel\\\\s*=\\\\s*).*$/\\\\1 Never/g'\r"
        #expect \$chroot_prompt

        # Install extra packages
        send -- {echo "$packages" | pacman -S --noconfirm -}
        send -- "\r"
        expect \$chroot_prompt

        # Install syslinux
        send -- "pacman -S --noconfirm syslinux\r"
        expect \$chroot_prompt
        send -- "syslinux-install_update -i -a -m\r"
        #expect \$chroot_prompt
        #send -- "dd bs=440 count=1 conv=notrunc if=/usr/lib/syslinux/bios/mbr.bin of=/dev/sda\r"
        expect \$chroot_prompt
        send -- "sed -i /boot/syslinux/syslinux.cfg -e '1i SERIAL 0 115200'\r"
        expect \$chroot_prompt
        send -- "sed -i /boot/syslinux/syslinux.cfg -e 's|APPEND root=/dev/sda3|APPEND console=tty0 console=ttyS0,115200 root=/dev/sda1|g'\r"
        expect \$chroot_prompt

        # Restore the original pacman.conf
        send -- "mv /etc/pacman.conf.orig /etc/pacman.conf\r"
        expect \$chroot_prompt

        # Set the root password
        send -- "passwd\r"
        expect "New password: "
        send -- "$root_password\r"
        expect "Retype new password: "
        send -- "$root_password\r"
        expect \$chroot_prompt

        # Create a non-admin user
        send -- "useradd -d /home/$username -m $username\r"
        expect \$chroot_prompt
        send -- "passwd $username\r"
        expect "New password: "
        send -- "$user_password\r"
        expect "Retype new password: "
        send -- "$user_password\r"
        expect \$chroot_prompt

        # Enable the dhcpcd service
        send -- "ln -sf /usr/lib/systemd/system/dhcpcd.service /etc/systemd/system/multi-user.target.wants/dhcpcd.service\r"
        expect \$chroot_prompt

        # Enable and configure the SSH daemon
        send -- "ln -sf /usr/lib/systemd/system/sshd.service /etc/systemd/system/multi-user.target.wants/sshd.service\r"
        expect \$chroot_prompt
        send -- "cat <<_EOF_ >> /etc/ssh/sshd_config
RSAAuthentication yes
PubkeyAuthentication yes
_EOF_
"

        expect \$chroot_prompt

        # Copy the user SSH key
        send -- "mkdir -p /home/$username/.ssh\r"
        expect \$chroot_prompt
        send -- {echo "$ssh_pubkey" >> "/home/$username/.ssh/authorized_keys"}
        send -- "\r"
        expect \$chroot_prompt
        send -- {echo "$ssh_privkey" > "/home/$username/.ssh/$(basename $ssh_keyfile)"}
        send -- "\r"
        expect \$chroot_prompt
        send -- {chmod 0600 "/home/$username/.ssh/$(basename $ssh_keyfile)"}
        send -- "\r"
        expect \$chroot_prompt
        send -- {echo "$ssh_pubkey" > "/home/$username/.ssh/$(basename $ssh_keyfile).pub"}
        send -- "\r"
        expect \$chroot_prompt
        send -- {chown -R $username "/home/$username/.ssh"}
        send -- "\r"
        expect \$chroot_prompt

        # Run any post-install scripts
        send -- {$post_install}
        send -- "\r"
        expect \$chroot_prompt

        # Clear the pacman cache
        send -- "rm -rf /var/cache/pacman/pkg/*\r"
        expect \$chroot_prompt

        # Exit and shutdown
        send -- "exit\r"
        expect \$prompt
        send -- "umount /mnt\r"
        expect \$prompt
        send -- "shutdown -h now\r"
        expect eof

        system {echo -e "\n${GREEN}Arch Linux system installed under ${imgfile}${NORMAL}"}
EOF

    echo "--- Log closed at $(date)" >> "$logfile"
}

function install() {
    optstring=':o:a:s:m:h:P:u:p:z:l:M:'
    [[ "$1" == '--help' ]] && install_usage

    while getopts ${optstring} arg; do
        case ${arg} in
            o) imgfile="${OPTARG}";;
            a) architecture="${OPTARG}";;
            s) disk_size="${OPTARG}";;
            m) memory="${OPTARG}";;
            h) hostname="${OPTARG}";;
            P) root_password="${OPTARG}";;
            u) username="${OPTARG}";;
            p) user_password="${OPTARG}";;
            z) timezone="${OPTARG}";;
            l) locale="${OPTARG}";;
            M) img_download_page="${OPTARG}";;
            ?)
                echo "Invalid option: -${OPTARG}" >&2
                install_usage;;
        esac
    done

    [ -z "$imgfile" ] && read -p "Output disk image file [$default_imgfile]: " imgfile
    [ -z "$architecture" ] && read -p "Architecture [$default_architecture]: " architecture
    [ -z "$disk_size" ] && read -p "Disk size [$default_disk_size]: " disk_size
    [ -z "$memory" ] && read -p "Memory in KB [$default_memory]: " memory
    [ -z "$hostname" ] && read -p "Hostname [$default_hostname]: " hostname
    [ -z "$root_password" ] && read -sp "Root password [$default_root_password]: " root_password && echo
    [ -z "$username" ] && read -p "Non-admin username [$default_username]: " username
    [ -z "$user_password" ] && read -sp "Non-admin user password [$default_user_password]: " user_password && echo
    [ -z "$timezone" ] && read -p "Timezone [$default_timezone]: " timezone
    [ -z "$locale" ] && read -p "Locale [$default_locale]: " locale

    for var in imgfile \
        architecture \
        disk_size \
        memory \
        hostname \
        root_password \
        username \
        user_password \
        timezone \
        locale
    do
        default_var=default_$var
        [ -z "${!var}" ] && declare ${var}=${!default_var}
        export $var
    done

    download_latest_arch_iso
    create_disk_image
    install_os
}
################


##################
# Main script

function usage() {
    actions=(install)
    cat <<EOF
Usage: $(basename "$0") <action>

Run $(basename "$0") <action> --help for more details.
Available actions:

$actions
EOF

    exit 1
}

function main() {
    action=$1
    shift

    case "$action" in
        install) install $*;;
        *) usage;;
    esac
}

main $*

##################