#!/bin/bash
# -*- mode: sh; indent-tabs-mode: nil; sh-basic-offset: 4; -*-
# vim: et sts=4 sw=4

#  SPDX-License-Identifier: LGPL-2.1+
#
#  Copyright © 2019-2021 Collabora Ltd.
#  Copyright © 2019-2021 Valve Corporation.
#
#  This file is part of steamos-customizations.
#
#  steamos-customizations is free software; you can redistribute it and/or
#  modify it under the terms of the GNU Lesser General Public License as
#  published by the Free Software Foundation; either version 2.1 of the License,
#  or (at your option) any later version.

# NOTE: large portions shared with atomic-update/rauc/post-install.sh
# TODO: deduplicate / shere code between there and here

set -e
set -u

export LC_ALL=C

DEVICE_OUT=      # $1

FORCE=0          # --force
GROW_LAST=1      # --grow-last=
FACTORY_RESET=0  # --factory-reset=

BOOTCONF_OLDPATH=SteamOS/bootconf
BOOTCONF_RELDIR=SteamOS/conf
PARTSETS_RELDIR=SteamOS/partsets
GRUB_BINARY_RELPATH=EFI/steamos/grubx64.efi
GRUB_CONFIG_RELPATH=EFI/steamos/grub.cfg

STEAMOS_N_PARTITIONS=10
STEAMOS_ALL_PARTLABELS='esp efi-A efi-B verity-A verity-B rootfs-A rootfs-B var-A var-B home'

[ -e /usr/lib/steamos/steamos-partitions-lib ] && \
    . /usr/lib/steamos/steamos-partitions-lib  || \
    { echo "Failed to source '/usr/lib/steamos/steamos-partitions-lib'"; exit 1; }

# Helpers

log()  { echo >&2 "$@"; }
fail() { echo >&2 "$@"; exit 1; }
log_step() { echo >&2 "$(tput bold)>>>" "$@" "...$(tput sgr0)"; }

usage() {
    local status=${1-2}

    if [ $status -ne 0 ]; then
        exec >&2
    fi

    echo
    echo "Usage: $(basename $0) [OPTIONS] DESTINATION_DISK"
    echo
    echo "Clone the currently running SteamOS to another disk."
    echo
    echo "Note that it works only if SteamOS is the only OS installed."
    echo
    echo "Example: $(basename $0) /dev/sdX"
    echo
    echo "OPTIONS"
    echo
    echo "  --force"
    echo "        Do things even when it looks unsafe. Use with caution."
    echo
    echo "  --grow-last 0|1       (default: 1)"
    echo "        Use this option if you want to grow the last partition (presumably home)"
    echo "        so that it fills the whole disk."
    echo
    echo "  --factory-reset 0|1   (default: 0)"
    echo "        Install the system as if it was out of factory. It means that persistent"
    echo "        data won't be copied from the current system, ie. home and var partitions"
    echo "        will be formatted and left empty."
    echo
    echo "OPTIONS - NOT IMPLEMENTED YET"
    echo
    echo "  --skip-other     Use this if you don't want to clone the OTHER set of partitions."
    echo "                   (say you're running A, and you don't want to install B)."
    echo "                   This will create the partitions, but won't copy the data."
    echo "                   I guess a wipefs would be welcome though..."
    echo

    exit $status
}

ask() {
    local message="$1 [y/n] "

    while read -r -t 0; do
        read -n 256 -r -s
    done

    while true; do
        read -p "$message" answer
        case "$answer" in
            [Yy]|YES|Yes|yes) return 0;;
            [Nn]|NO|No|no)    return 1;;
            *) echo "Please answer yes or no.";;
        esac
    done
}

get_all_partlabels() {

    # Get the list of all partlabels on a given disk, for example:
    # /dev/sda -> efi esp-A esp-B ...

    local disk=$1

    [ -b "$disk" ] || return

    # trick: use echo to turn '/n' into ' ' WITHOUT adding an ending ' '
    echo $(sfdisk -ql -o name "$disk" | tail +2)
}

validate_partlabels() {

    # Ensure that there's no un-expected partition labels. In other words,
    # there should be only SteamOS installed on the disk, and nothing else.
    # This is because we just copy one disk to another, without trying to
    # be smart. So we just don't handle complicated cases.

    local all_partlabels=$1
    local all_partlabels_sorted=$(echo $all_partlabels | tr ' ' '\n' | sort)
    local expected_partlabels_sorted=$(echo $STEAMOS_ALL_PARTLABELS | tr ' ' '\n' | sort)

    [ "$all_partlabels_sorted" == "$expected_partlabels_sorted" ]
}

replicate_partition_table() {

    # Replicate partition table from one disk to another
    # TODO What about protective MBR? Is it replicated as well?

    local src=$1
    local dst=$2

    # before doing anything foolish, dry run
    sgdisk --pretend --replicate $dst $src

    sgdisk --zap-all $dst
    sgdisk --replicate $dst $src
    sgdisk --randomize-guids $dst

    if [ $GROW_LAST -eq 1 ]; then
        # can do with parted as well: parted $dst resizepart $NPARTS 100%
        echo ", +" | sfdisk -q -N $NPARTS $dst
    fi
}

format_block_device() {

    # Format a block device

    local srcdisk=$1
    local dstdisk=$2
    local partlabel=$3
    local srcdev=
    local dstdev=
    local fs=
    local label=
    local mkfs=

    srcdev=$(get_device_by_partlabel $srcdisk $partlabel)
    if ! [ "$srcdev" ]; then
        log "Could not find partlabel '$partlabel' on disk '$srcdisk'"
        return 1
    fi

    dstdev=$(get_device_by_partlabel $dstdisk $partlabel)
    if ! [ "$dstdev" ]; then
        log "Could not find partlabel '$partlabel' on disk '$dstdisk'"
        return 1
    fi

    fs=$(lsblk -no fstype "$srcdev")

    case "$fs" in
        (ext4)
            label=$(e2label "$srcdev")
            if [ "$label" ]; then
                mkfs="mkfs.ext4 -F -L $label"
            else
                mkfs="mkfs.ext4 -F"
            fi
            echo yes | $mkfs "$dstdev"
            ;;

        (*)
            log "Can't format block device '$dstdev', unsupported fs: '$fs'"
            return 1
            ;;
    esac
}

copy_block_device() {

    # Copy a partition from one disk to another, given its partlabel

    local srcdisk=$1
    local dstdisk=$2
    local partlabel=$3
    local srcdev=
    local dstdev=
    local mntpoint=
    local please_unfreeze=0
    local please_remount=0

    srcdev=$(get_device_by_partlabel $srcdisk $partlabel)
    if ! [ "$srcdev" ]; then
        log "Could not find partlabel '$partlabel' on disk '$srcdisk'"
        return 1
    fi

    dstdev=$(get_device_by_partlabel $dstdisk $partlabel)
    if ! [ "$dstdev" ]; then
        log "Could not find partlabel '$partlabel' on disk '$dstdisk'"
        return 1
    fi

    mntpoint=$(find_mountpoint_for_device $srcdev)
    if [ "$mntpoint" ]; then
        local fstype=

        fstype=$(findmnt --real -no fstype $mntpoint)

        log "Source device '$srcdev' mounted at '$mntpoint' already"

        case "$fstype" in
            (ext*)
                # For ext filesystems, we prefer feezing over unmounting,
                # as we know that the fs might be busy and unmount might fail.
                log "Freezing filesystem at '$mntpoint'"
                fsfreeze -f $mntpoint
                please_unfreeze=1
                ;;
            (vfat)
                # We hope to remount the filesystem later on, ie. we assume
                # that it's present in /etc/fstab
                log "Unmounting filesystem at '$mntpoint'"
                umount $mntpoint
                please_remount=1
                ;;
            (*)
                log "Unexpected filesystem '$fstype'"
                return 1
        esac
    fi

    log "$srcdev > $dstdev"
    cat $srcdev > $dstdev

    if [ $please_unfreeze -eq 1 ]; then
        log "Unfreezing '$mntpoint' ..."
        fsfreeze -u $mntpoint
    fi

    if [ $please_remount -eq 1 ]; then
        log "Remounting '$mntpoint' ..."
        mount $mntpoint || : # failure ok
    fi
}

grow_fs() {

    # Grow a filesystem

    local device=$1
    local fs=

    fs=$(lsblk -no fstype "$device")

    case "$fs" in
        (ext*)
            e2fsck -fp "$device"
            resize2fs "$device"
            ;;
        (*)
            log "Can't grow fs on '$device', unsupported fs: '$fs'"
            return 1
            ;;
    esac
}

cleanup_var() {

    # Cleanup var partition
    #
    # Some data on the var partition don't make sense when we clone the
    # disk to another device, so we remove it. Namely:
    # - /var/boot: might contain a dkms-rebuilt initrd, it's hardware specific

    local disk=$1
    local partset=$2
    local partlabel=var-$partset
    local vardev=
    local mntpoint=

    vardev=$(get_device_by_partlabel $disk $partlabel)
    if ! [ "$vardev" ]; then
        log "Could not find partlabel '$partlabel' on disk '$disk'"
        return 1
    fi

    # TODO trap and exit!

    mntpoint=$(mktemp -d)
    mount $vardev $mntpoint

    rm -fr $mntpoint/boot

    umount $mntpoint
    rmdir $mntpoint
}

setup_esp() {
    local disk=$1
    local id=$2
    local espdev=$(get_device_by_partlabel $disk esp)
    local espmnt=
    local confdir=

    if ! [ "$espdev" ]; then
        log "Could not find partlabel '$partlabel' on disk '$disk'"
        return 1
    fi

    espmnt=$(mktemp -d)
    mount $espdev $espmnt

    # reset the bootconf file:
    confdir=$espmnt/$BOOTCONF_RELDIR
    bootconf=$confdir/$id.conf
    mkdir -p $confdir
    steamos-bootconf create --conf-dir $confdir --image $id --set title $id

    umount $espmnt
    rmdir $espmnt
}

setup_efi() {

    # Setup efi partition, a subtle process...
    #
    # Note that an efi partition might be empty (say only A was booted
    # up to now, and B was never booted yet, so efi-B is be empty),
    # and in this case we leave it empty.

    local disk=$1
    local self=$2
    local partlabel=
    local efidev=
    local mntpoint=
    local bootconf=
    local partsets=
    local grub_binary=
    local grub_rootcfg=
    local grub_varcfg=

    partlabel=efi-$self
    efidev=$(get_device_by_partlabel $disk $partlabel)
    if ! [ "$efidev" ]; then
        log "Could not find partlabel '$partlabel' on disk '$disk'"
        return 1
    fi

    # TODO trap and exit!

    mntpoint=$(mktemp -d)
    mount $efidev $mntpoint

    # replace the partition definitions, only if it exists already

    partsets=$mntpoint/$PARTSETS_RELDIR
    if [ -e $partsets ]; then
        rm -fr $partsets
        steamos-chroot --disk $disk --partset $self -- \
            steamos-partsets /efi/$PARTSETS_RELDIR
    fi

    setup_esp $disk $self

    # erase the bootconf file at the old location:
    # the new location is on the ESP
    bootconf=$mntpoint/$BOOTCONF_OLDPATH
    if [ -e $bootconf ]; then
        rm -f $bootconf
    fi

    # re-install the grub bootloader, only if it exists already

    grub_binary=$mntpoint/$GRUB_BINARY_RELPATH
    if [ -e $grub_binary ]; then
        steamos-chroot --disk $disk --partset $self -- \
            grub-mkimage
    fi

    # update the grub config for the root partition, only if it exists already

    grub_cfg=$mntpoint/$GRUB_CONFIG_RELPATH
    if [ -e $grub_cfg ]; then
        steamos-chroot --disk $disk --partset $self -- \
            update-grub
    fi

    umount $mntpoint
    rmdir $mntpoint
}

# main

while [ $# -gt 0 ]; do
    case "$1" in
        -h|--help)
            usage 0
            ;;
        -f|--force)
            shift
            FORCE=1
            ;;
        --factory-reset)
            shift
            [ "${1:-}" -eq 0 ] || [ "${1:-}" -eq 1 ] || usage 1
            FACTORY_RESET=$1
            shift
            ;;
        --grow-last)
            shift
            [ "${1:-}" -eq 0 ] || [ "${1:-}" -eq 1 ] || usage 1
            GROW_LAST=$1
            shift
            ;;
        *)
            if [ -z "$DEVICE_OUT" ]; then DEVICE_OUT=$1; shift; continue; fi
            usage 1
            ;;
    esac
done

[ "$DEVICE_OUT" ] || usage 1

ROOTDEV=$(find_device_for_mountpoint_deep '/')
DEVICE_IN=$(get_parent_device $ROOTDEV)
[ -b "$DEVICE_IN" ] || fail "'$DEVICE_IN' is not a block device"

DEVICE_OUT=$(realpath $DEVICE_OUT)
[ -b "$DEVICE_OUT" ] || fail "'$DEVICE_OUT' is not a block device"

NPARTS=$(get_partition_count $DEVICE_IN)
[ $NPARTS -eq $STEAMOS_N_PARTITIONS ] || \
    fail "Unexpected number of partitions on '$DEVICE_IN': $NPARTS (expected $STEAMOS_N_PARTITIONS)"

ALL_PARTLABELS=$(get_all_partlabels $DEVICE_IN)
if ! validate_partlabels "$ALL_PARTLABELS"; then
    log "Something wrong with partlabels on '$DEVICE_IN'"
    log "We got     : '$ALL_PARTLABELS'"
    log "We expected: '$STEAMOS_ALL_PARTLABELS'"
    log "Are you sure that ONLY SteamOS is installed on '$DEVICE_IN'?"
    exit 1
fi

if ! verity_is_enabled && ! [ $FORCE -eq 1 ]; then
    echo "Cloning SteamOS disk is available only when dm-verity is enabled."
    echo "Use --force if you really want to proceed."
    exit 1
fi

echo
echo "You're about to clone your SteamOS install to another disk."
echo "Basically, the disk '$DEVICE_IN' will be copied to '$DEVICE_OUT'."
echo "ALL THE DATA ON '$DEVICE_OUT' WILL BE LOST! There's no way back."
if ! ask "Do you want to continue?"; then
    echo "Alright buddy. Come back when you feel ready."
    exit 0
fi
echo

log_step "Replicating partition table from disk '$DEVICE_IN' to '$DEVICE_OUT'"
replicate_partition_table $DEVICE_IN $DEVICE_OUT

partprobe $DEVICE_OUT
udevadm trigger $DEVICE_OUT
udevadm settle

for partlabel in $ALL_PARTLABELS; do
    COPY_PART=1
    case $partlabel in
        home|var-A|var-B)
            [ $FACTORY_RESET -eq 1 ] && COPY_PART=0
            ;;
    esac
    if [ $COPY_PART -eq 1 ]; then
        log_step "Copying partition '$partlabel'"
        copy_block_device $DEVICE_IN $DEVICE_OUT $partlabel
    else
        # TODO For the home partition, we should enable the casefold feature
        log_step "Formatting partition '$partlabel'"
        format_block_device $DEVICE_IN $DEVICE_OUT $partlabel
    fi
done

# Disable reserved blocks on the home parition, as requested by Valve.
tune2fs -m 0 $(get_device_by_partlabel $dstdisk home)

partprobe $DEVICE_OUT
udevadm trigger $DEVICE_OUT
udevadm settle

if [ $GROW_LAST -eq 1 ]; then
    LASTPART=$(get_device_by_partnumber $DEVICE_OUT $NPARTS)
    log_step "Growing filesystem on last partition '$LASTPART'"
    grow_fs $LASTPART
fi

# TODO Talking about cleaning, we must remove ssh host keys and machine-id
# Depending whether this is atomic image or not, these files might be on
# the root partition, or on the etc overlay (on the var partition).

if ! [ $FACTORY_RESET -eq 1 ]; then
    log_step "Cleaning up var A"
    cleanup_var $DEVICE_OUT A
    log_step "Cleaning up var B"
    cleanup_var $DEVICE_OUT B
fi

log_step "Setting up efi A"
setup_efi $DEVICE_OUT A
log_step "Setting up efi B"
setup_efi $DEVICE_OUT B

echo
echo "Installation successful, please power off the board,"
echo "remove the media and turn it on again."
echo
