#!/bin/sh
# vim: set ts=4 et:

set -eu

DEVICE=/dev/xvdf
TARGET=/mnt/target
SETUP=/tmp/setup-ami.d

[ "$VERSION" = 'edge' ] && V= || V=v
MAIN_REPO="https://dl-cdn.alpinelinux.org/alpine/$V$VERSION/main/$ARCH"

die() {
    printf '\033[1;31mERROR:\033[0m %s\n' "$@" >&2  # bold red
    exit 1
}

einfo() {
    printf '\n\033[1;36m> %s\033[0m\n' "$@" >&2  # bold cyan
}

rc_add() {
    runlevel="$1"; shift  # runlevel name
    services="$*"         # names of services

    for svc in $services; do
        mkdir -p "$TARGET/etc/runlevels/$runlevel"
        ln -s "/etc/init.d/$svc" "$TARGET/etc/runlevels/$runlevel/$svc"
        echo " * service $svc added to runlevel $runlevel"
    done
}

validate_block_device() {
    lsblk -P --fs "$DEVICE" >/dev/null 2>&1 || \
        die "'$DEVICE' is not a valid block device"

    if lsblk -P --fs "$DEVICE" | grep -vq 'FSTYPE=""'; then
        die "Block device '$DEVICE' is not blank"
    fi
}

main_repo_pkgs() {
    wget -T 10 -q -O - "$MAIN_REPO/" | grep '^<a href=' | cut -d\" -f2
}

fetch_apk_tools() {
    store="$(mktemp -d)"
    tarball="$(main_repo_pkgs | grep ^apk-tools-static- | sort -V | tail -n 1)"

    wget -T 10 -q -O "$store/$tarball" "$MAIN_REPO/$tarball"
    tar -C "$store" --warning=no-unknown-keyword -xf "$store/$tarball"

    find "$store" -name apk.static
}

# mostly from Alpine's /sbin/setup-disk
setup_partitions() {
    start=2M # Needed to align EBS partitions
    line=

    # create new partitions
    (
        for line in "$@"; do
            case "$line" in
            0M*) ;;
            *) echo "$start,$line"; start= ;;
            esac
        done
    ) | sfdisk --quiet --label gpt "$DEVICE"

    # we assume that the build host will create the new devices within 5s
    tries=5
    while [ ! -e "${DEVICE}1" ]; do
        [ $tries -eq 0 ] && break
        sleep 1
        tries=$(( tries - 1 ))
    done
    [ -e "${DEVICE}1" ] || die "Expected new device ${DEVICE}1 not created"
}

make_filesystem() {
    root_dev="$DEVICE"

    if [ "$BOOTLOADER" = grub-efi ] || [ "$BOOTLOADER" = EFI_STUB ]; then
        # create a small EFI partition (remainder for root)
        if [ "$BOOTLOADER" = EFI_STUB ]; then
            setup_partitions '11M,U,*' ',L'   # kernel + initfs
        else
            setup_partitions '512K,U,*' ',L'  # currently 278K used
        fi
        root_dev="${DEVICE}2"
        mkfs.vfat -n EFI "${DEVICE}1"
    fi

    mkfs.ext4 -O ^64bit -L / "$root_dev"
    mount "$root_dev" "$TARGET"

    if [ "$BOOTLOADER" = grub-efi ] || [ "$BOOTLOADER" = EFI_STUB ]; then
        # mount small EFI partition
        mkdir -p "$TARGET/boot/efi"
        mount -t vfat "${DEVICE}1" "$TARGET/boot/efi"
    fi
}

setup_repositories() {
    mkdir -p "$TARGET/etc/apk/keys"
    echo "$REPOS" > "$TARGET/etc/apk/repositories"
}

fetch_keys() {
    tmp="$(mktemp -d)"
    tarball="$(main_repo_pkgs | grep ^alpine-keys- | sort -V | tail -n 1)"

    wget -T 10 -q -O "$tmp/$tarball" "$MAIN_REPO/$tarball"
    tar -C "$TARGET" --warning=no-unknown-keyword -xvf "$tmp/$tarball" etc/apk/keys
    rm -rf "$tmp"
}

install_base() {
    $apk add --root "$TARGET" --no-cache --initdb alpine-base
    # verify release matches
    if [ "$VERSION" != edge ]; then
        ALPINE_RELEASE=$(cat "$TARGET/etc/alpine-release")
        [ "$RELEASE" = "$ALPINE_RELEASE" ] || \
            die "Release '$RELEASE' does not match /etc/alpine-release: $ALPINE_RELEASE"
    fi
}

setup_chroot() {
    mount -t proc none "$TARGET/proc"
    mount --bind /dev "$TARGET/dev"
    mount --bind /sys "$TARGET/sys"

    # Needed for bootstrap, will be removed in the cleanup stage.
    install -Dm644 /etc/resolv.conf "$TARGET/etc/resolv.conf"
}

install_core_packages() {
    chroot "$TARGET" apk --no-cache add $PKGS

    # EFI_STUB requires no bootloader
    [ "$BOOTLOADER" = EFI_STUB ] || \
        chroot "$TARGET" apk --no-cache add --no-scripts "$BOOTLOADER"

    # Disable starting getty for physical ttys because they're all inaccessible
    # anyhow. With this configuration boot messages will still display in the
    # EC2 console.
    sed -Ei '/^tty[0-9]/s/^/#/' "$TARGET/etc/inittab"

    # Enable the getty for the serial terminal. This will show the login prompt
    # in the get-console-output API that's accessible by the CLI and the web
    # console.
    sed -Ei '/^#ttyS0:/s/^#//' "$TARGET/etc/inittab"

    # Make it a little more obvious who is logged in by adding username to the
    # prompt
    sed -i "s/^export PS1='/&\\\\u@/" "$TARGET/etc/profile"

}

setup_mdev() {
    install -o root -g root -Dm755 -t "$TARGET/lib/mdev" \
        "$SETUP/nvme-ebs-links" \
        "$SETUP/eth-eni-hotplug"

    # insert nvme ebs mdev configs just above "# fallback" comment
    sed -n -i \
        -e "/# fallback/r $SETUP/etc/mdev-ec2.conf" \
        -e 1x -e '2,${x;p}' -e '${x;p}' \
        "$TARGET/etc/mdev.conf"
}

create_initfs() {
    sed -Ei "s/^features=\"([^\"]+)\"/features=\"\1 $INITFS_FEATURES\"/" \
        "$TARGET/etc/mkinitfs/mkinitfs.conf"

    chroot "$TARGET" /sbin/mkinitfs $(basename $(find "$TARGET/lib/modules/"* -maxdepth 0))
}

install_bootloader() {
    case "$BOOTLOADER" in
        syslinux)   install_extlinux ;;
        grub-efi)   install_grub_efi ;;
        EFI_STUB)   install_EFI_STUB ;;
        *)          die "unknown bootloader '$BOOTLOADER'" ;;
    esac
}

install_extlinux() {
    # Must use disk labels instead of UUID or devices paths so that this works
    # across instance familes. UUID works for many instances but breaks on the
    # NVME ones because EBS volumes are hidden behind NVME devices.
    #
    # Enable ext4 because the root device is formatted ext4
    #
    # Shorten timeout (1/10s) as EC2 has no way to interact with instance console
    #
    # ttyS0 is the target for EC2s "Get System Log" feature whereas tty0 is the
    # target for EC2s "Get Instance Screenshot" feature. Enabling the serial
    # port early in extlinux gives the most complete output in the system log.
    sed -Ei -e "s|^[# ]*(root)=.*|\1=LABEL=/|" \
        -e "s|^[# ]*(default_kernel_opts)=.*|\1=\"$KERNEL_OPTS\"|" \
        -e "s|^[# ]*(serial_port)=.*|\1=ttyS0|" \
        -e "s|^[# ]*(modules)=.*|\1=$KERNEL_MODS|" \
        -e "s|^[# ]*(default)=.*|\1=virt|" \
        -e "s|^[# ]*(timeout)=.*|\1=1|" \
        "$TARGET/etc/update-extlinux.conf"

    chroot "$TARGET" /sbin/extlinux --install /boot
    chroot "$TARGET" /sbin/update-extlinux --warn-only
}

install_grub_efi() {
    [ -d "/sys/firmware/efi" ] || die "/sys/firmware/efi does not exist"

    case "$ARCH" in
        x86_64)     grub_target=x86_64-efi ; fwa=x64 ;;
        aarch64)    grub_target=arm64-efi ; fwa=aa64 ;;
        *)          die "ARCH=$ARCH is currently unsupported" ;;
    esac

    # disable nvram so grub doesn't call efibootmgr
    chroot "$TARGET" /usr/sbin/grub-install --target="$grub_target" --efi-directory=/boot/efi \
        --bootloader-id=alpine --boot-directory=/boot --no-nvram

    # fallback mode
    install -D "$TARGET/boot/efi/EFI/alpine/grub$fwa.efi" "$TARGET/boot/efi/EFI/boot/boot$fwa.efi"

    # install default grub config
    envsubst < "$SETUP/etc/grub.template" > "$SETUP/etc/grub"
    install -o root -g root -Dm644 -t "$TARGET/etc/default" \
        "$SETUP/etc/grub"

    # generate/install new config
    chroot "$TARGET" grub-mkconfig -o /boot/grub/grub.cfg
}

install_EFI_STUB() {
    [ -d "/sys/firmware/efi" ] || die "/sys/firmware/efi does not exist"

    case "$ARCH" in
        x86_64)     fwa=x64 ;;
        aarch64)    fwa=aa64 ;;
        *)          die "ARCH=$ARCH is currently unsupported" ;;
    esac

    # TODO: kernel modules/options?
    # TODO: will also need initfs in here too
    # TODO: make it work

    # install kernel as UEFI fallback
    install -o root -g root -Dm644 "$TARGET/boot/vmlinuz-virt" \
        "$TARGET/boot/efi/EFI/boot/boot$fwa.efi"
    # replace original with a symlink
    rm "$TARGET/boot/vmlinuz-virt"
    ln -s "efi/EFI/boot/boot$fwa.efi" "$TARGET/boot/vmlinuz-virt"
}

setup_fstab() {
    install -o root -g root -Dm644 -t "$TARGET/etc" \
        "$SETUP/etc/fstab"

    # if we're using an EFI bootloader, add extra line for EFI partition
    if [ "$BOOTLOADER" = grub-efi ] || [ "$BOOTLOADER" = EFI_STUB ]; then
        cat "$SETUP/etc/fstab.grub-efi" >> "$TARGET/etc/fstab"
    fi
}

setup_networking() {
    # configure standard interfaces
    IFACE_CFG="$TARGET/etc/network/interfaces"
    install -o root -g root -Dm755 -d "$SETUP/etc/interfaces.d" "$IFACE_CFG.d"
    install -o root -g root -Dm644 -t "$IFACE_CFG.d" \
        "$SETUP/etc/interfaces.d/"*
    cat "$IFACE_CFG.d/lo" "$IFACE_CFG.d/eth0" > "$IFACE_CFG"

    install -o root -g root -Dm755 -t "$TARGET/etc/network" \
        "$SETUP/assemble-interfaces"

    # install ucdhcp hooks for EC2 ENI IPv6 and secondary IPv4
    install -o root -g root -Dm755 -t "$TARGET/usr/share/udhcpc" \
        "$SETUP/eth-eni-hook"
    for i in post-bound post-renew; do
        mkdir -p "$TARGET/etc/udhcpc/$i"
        ln -sf /usr/share/udhcpc/eth-eni-hook \
            "$TARGET/etc/udhcpc/$i"
    done

    # install ENI interface setup init script
    install -o root -g root -Dm755 -t "$TARGET/etc/init.d" \
        "$SETUP/eth-eni-setup"
}

enable_services() {
    for lvl_svcs in $SVCS; do
        rc_add $(echo "$lvl_svcs" | tr '=,' ' ')
    done
}

create_alpine_user() {
    # Allow members of the wheel group to sudo without a password. By default
    # this will only be the alpine user. This allows us to ship an AMI that is
    # accessible via SSH using the user's configured SSH keys (thanks to
    # tiny-ec2-bootstrap) but does not allow remote root access which is the
    # best-practice.
    sed -i '/%wheel .* NOPASSWD: .*/s/^# //' "$TARGET/etc/sudoers"

    # explicitly lock the root account
    chroot "$TARGET" /usr/bin/passwd -l root

    # There is no real standard ec2 username across AMIs, Amazon uses ec2-user
    # for their Amazon Linux AMIs but Ubuntu uses ubuntu, Fedora uses fedora,
    # etc... (see: https://alestic.com/2014/01/ec2-ssh-username/). So our user
    # and group, by default, are alpine because this is Alpine Linux.
    user="${EC2_USER:-alpine}"
    chroot "$TARGET" /usr/sbin/addgroup "$user"
    chroot "$TARGET" /usr/sbin/adduser -h "/home/$user" -s /bin/sh -G "$user" -D "$user"
    chroot "$TARGET" /usr/sbin/addgroup "$user" wheel
    chroot "$TARGET" /usr/bin/passwd -u "$user"

    # Let tiny-ec2-bootstrap know what the EC2 user of the AMI is
    echo "EC2_USER=$user" > "$TARGET/etc/conf.d/tiny-ec2-bootstrap"
}

configure_ntp() {
    # EC2 provides an instance-local NTP service syncronized with GPS and
    # atomic clocks in-region. Prefer this over external NTP hosts when running
    # in EC2.
    #
    # See: https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/set-time.html
    sed -e 's/^pool /server /' \
        -e 's/pool.ntp.org/169.254.169.123/g' \
        -i "$TARGET/etc/chrony/chrony.conf"
}

setup_script() {
    if [ -f "$SETUP/setup_script" ]; then
        einfo "Executing additional setup script"
        (
            cd "$SETUP"
            chmod u+x ./setup_script
            TARGET="$TARGET" ./setup_script
        )
    else
        einfo "No additional setup script"
    fi
}

cleanup() {
    # Sweep cruft out of the image that doesn't need to ship or will be
    # re-generated when the image boots
    rm -f \
        "$TARGET/var/cache/apk/"* \
        "$TARGET/etc/resolv.conf" \
        "$TARGET/root/.ash_history" \
        "$TARGET/etc/"*-

    # unmount extra EFI mount
    if [ "$BOOTLOADER" = grub-efi ] || [ "$BOOTLOADER" = EFI_STUB ]; then
      umount "$TARGET/boot/efi"
    fi

    umount \
        "$TARGET/dev" \
        "$TARGET/proc" \
        "$TARGET/sys"

    umount "$TARGET"
}

main() {
    validate_block_device

    [ -d "$TARGET" ] || mkdir "$TARGET"

    einfo "Fetching static APK tools"
    apk="$(fetch_apk_tools)"

    einfo "Creating root filesystem"
    make_filesystem

    einfo "Configuring Alpine repositories"
    setup_repositories

    einfo "Fetching Alpine signing keys"
    fetch_keys

    einfo "Installing base system"
    install_base

    setup_chroot

    einfo "Installing core packages"
    install_core_packages

    einfo "Configuring and enabling '$BOOTLOADER' boot loader"
    create_initfs
    install_bootloader

    einfo "Configuring system"
    setup_mdev
    setup_fstab
    setup_networking
    enable_services
    create_alpine_user
    configure_ntp

    setup_script

    einfo "All done, cleaning up"
    cleanup
}

main "$@"