#!/bin/bash
#
# Runs in initramfs to create an encrypted rootfs
# Requires a second drive larger than the rootfs
#

function usage() {
    echo $"Usage: sdmcryptfs rootdev scratchdev

rootdev    Name of device with rootfs (possibly /dev/mmcblk0 or /dev/sda)
scratchdev Name of scratch device (must be larger than rootfs on rootdev)
"
exit 0
}

function errexit() {
    echo -e "$1"
    exit 1
}

function thisdate() {
    echo  "$(date +"$datefmt")"
}

function pmsg() {
    local msg="$1"
    echo "> $(thisdate) $msg"
}

function askyn() {
    local ans
    echo -n "$1" '[y/N]? ' ; read $2 ans
    case "$ans" in
        y*|Y*) return 0 ;;
        *) return 1 ;;
    esac
}

function getcbbdate() {
    #
    # Returns date in busybox canonical format
    #
    local cdtfmt="+%Y-%m-%d %H:%M:%S"  # Canonical date format for manipulation
   echo $(date "$cdtfmt")
}

function datediff() {
    #
    # Computes difference between two dates
    # returns it as a string of Days:Hours:Minutes:Seconds
    #
    local days hours minutes seconds eseconds
    local btime="$1" etime="$2"
    eseconds=$(($(/bin/date -d "$etime" +%s) - $(/bin/date -d "$btime" +%s)))
    days=$((eseconds/86400))
    hours=$(((eseconds - $((days*86400)))/3600))
    minutes=$(((eseconds - $((days*86400)) - $((hours*3600)))/60))
    seconds=$((eseconds - $((days*86400)) - $((hours*3600)) - $((minutes*60))))
    #echo "$days:$hours:$minutes:$seconds"
    [ $days -ne 0 ] && printf "%02d:%02d:%02d:%02d\n" "$days" "$hours" "$minutes" "$seconds" || printf "%02d:%02d:%02d\n" "$hours" "$minutes" "$seconds"
}

function getgbstr() {
    #
    # $1: # of bytes in partition
    #
    # Returns the string "(nn.nnGB, mm.mmGiB)"

    local nbytes=$1
    local gb=1000000000 gib=1073741824 gb2=500000000 gi2=536870912
    local ngbytes ngibytes 
    ngbytes=$(printf %.1f "$(( ((10 * nbytes)+gb2) / gb ))e-1")
    ngibytes=$(printf %.1f "$(( ((10 * nbytes)+gi2) / gib))e-1")
    echo "(${ngbytes}GB, ${ngibytes}GiB)"
    return
}

function ispdevp() {
    local dev="$1"
    [[ "$dev" =~ "mmcblk" ]] || [[ "$dev" =~ "nvme0n1" ]] && return 0 || return 1
}

function getpartname() {
    local dev=$1
    local p1=1 p2=2

    ispdevp $dev && p1="p1" && p2="p2"
    echo "${dev}${p2}"
}

function getpartsize() {
    local rootpart=$1
    local pat r2fs partsize
    pat='([[:space:]][[:digit:]]+)'
    r2fs=$(resize2fs -fM -p $rootpart 2>&1)
    r2fs=${r2fs##* is}
    [[ $r2fs =~ $pat ]]
    partsize=${BASH_REMATCH[1]}
    echo $partsize
}

function diskindrive() {
    local dev=$1
    # parted print will fail if disk partitions are borked
    # so try to mklabel the disk in that case
    # NOTE: Destructive; only use on scratch disk
    parted -ms $dev print >/dev/null 2>&1 && return 0
    parted $dev mklabel msdos >/dev/null 2>&1 < <(echo C) && return 0
    return 1
}

function checkdev() {
    local dev=$1
    [[ -b $dev ]] && diskindrive $dev && return 0
    echo "? Unrecognized device or no disk in drive '$dev'"
    return 1
}

function wipescratch() {
    local starttime partsize=$1

    if askyn "Wipe scratch disk '$sdisk'"
    then
	pmsg "> Write random data on '$sdisk'"
	starttime="$(getcbbdate)"
	dd bs=4k count=$partsize if=/dev/random of=$sdisk
	pmsg "rootfs Wipe elapsed time: $(datediff "$starttime" "$(getcbbdate)")"
    else
	pmsg "> Scratch disk '$sdisk' has an unencrypted copy of rootfs"
	pmsg "  Clean it up after rootfs encryption cleanup completes if that concerns you"
	pmsg "  One way: sudo dd bs=4k count=$partsize if=/dev/random of=$sdisk"
    fi
}

function save_rootfs() {
    local dev=$1 sdisk=$2 partsize=$3
    local rootfs=$(getpartname $dev)
    local starttime endtime pbytes pgb
    checkdev $sdisk || return 1
    pmsg "Save rootfs '$rootfs' to '$sdisk'"
    pbytes=$((partsize*4096))
    pgb=$(((pbytes / 1073741824)+1))
    pmsg "rootfs save should take less than $pgb minutes"
    starttime="$(getcbbdate)"
    dd bs=4k count=$partsize if=$rootfs of=$sdisk
    pmsg "rootfs Save elapsed time: $(datediff "$starttime" "$(getcbbdate)")"
    return 0
}

if [ ! -e /mnt ]; then
    mkdir -p /mnt
    sleep 1
fi

function find_keyfile() {
    local kfl=""

    if [ -f /etc/sdmkeyfile ]
    then
	read kfl < /etc/sdmkeyfile
	[[ "$kfl" != "" ]] && [[ -f /etc/$kfl ]] && kfl="/etc/$(basename $kfl)"
    fi
    echo "$kfl"
}

function config_keyfile() {
    local kfl="" kfn=""

    [ -f /etc/sdmkeyfile ] || return
    kfl=$(find_keyfile)
    [ "$kfl" == "" ] || kfn=$(basename $kfl)
    if [ "$kfn" != "" ]
    then
	pmsg "Add LUKS key '$kfn' from file $kfl"
	pmsg "When prompted for an existing passphrase enter the one you just created"
	cryptsetup luksAddKey $rootfs $kfl
	[ $? -ne 0 ] && echo "? cryptsetup returned an error adding key '$kfn': $?"
	rm -f $kfl
    fi
}

function cryptpart() {
    local dev=$1
    local cipher
    case "$xcrypto" in
	xchacha)
	    cipher="xchacha20,aes-adiantum-plain64"
	    ;;
	aes)
	    cipher="aes-xts-plain64"
	    ;;
	aes-*)
	    cipher="$xcrypto"
	    ;;
    esac

    pmsg "Enable LUKS2 encryption on '$rootfs' with cipher '$cipher'"
    pmsg "Enabling encryption could take a while"
    pmsg "OK to ignore superblock signature warning"

    if [ $xnopwd -eq 1 ]
    then
	cryptsetup luksFormat --type luks2 $rootfs $(find_keyfile)
    else
	while [ 0 ]
	do
	    cryptsetup luksFormat --type luks2 --cipher $cipher --hash sha256 --iter-time 5000 --key-size 256 --pbkdf argon2i $rootfs
	    [ $? -eq 0 ] && break
	done
	config_keyfile $dev
    fi
    return 0
}

function cryptopen() {
    local dev=$1
    local rootfs=$(getpartname $dev)
    pmsg "Unlock encrypted partition '$rootfs'"
    pmsg "Unlock will take several seconds"
    if [ $xnopwd -eq 1 ]
    then
	cryptsetup luksOpen $rootfs $cryptdevice --key-file $(find_keyfile)
    else
	while [ 0 ]
	do
	    cryptsetup luksOpen $rootfs $cryptdevice
	    [ $? -eq 0 ] && break
	done
    fi
    return 0
}

function restore_rootfs() {
    local dev=$1 sdisk=$2 partsize=$3
    local rootfs=$(getpartname $dev)
    local cpartnum cpartstart cpartend cpartsize etc
    local starttime endtime pbytes pgb
    checkdev $sdisk || return 1
    pmsg "Restore '$rootfs' from '$sdisk'"
    pbytes=$((partsize*4096))
    pgb=$(((pbytes / 1073741824)+1))
    pmsg "rootfs restore should take less than $pgb minutes"
    starttime="$(getcbbdate)"
    dd bs=4k count=$partsize if=$sdisk of=/dev/mapper/$cryptdevice
    pmsg "rootfs Restore elapsed time: $(datediff "$starttime" "$(getcbbdate)")"
    pmsg "Restore complete; Expand rootfs..."
    resize2fs -f /dev/mapper/$cryptdevice
    while read line
    do
        if [[ "$line" =~ "ext4" ]]
        then
            IFS=":;" read cpartnum cpartstart cpartend cpartsize fstype etc <<< $line
            cpartsize=${cpartsize%B}
	    break
        fi
    done < <(parted -ms /dev/mapper/$cryptdevice unit B print)
    pmsg "rootfs partition size: $cpartsize Bytes $(getgbstr $cpartsize)"
    pmsg "cryptsetup status for /dev/mapper/$cryptdevice:"
    cryptsetup status /dev/mapper/$cryptdevice
    return 0
}

function runall() {
    local dev=$1 sdisk=$2 pbytes pgb
    local rootfs=$(getpartname $dev)
    checkdev $dev || return 1
    checkdev $sdisk || return 1
    pmsg "Shrink partition '$rootfs' and get its size"
    partsize=$(getpartsize $rootfs)
    pbytes=$((partsize*4096))
    pmsg "Device '$dev' rootfs size: $partsize 4K blocks $(getgbstr $pbytes)"
    save_rootfs $dev $sdisk $partsize || return 1
    cryptpart $dev || return 1
    cryptopen $dev || return 1
    restore_rootfs $dev $sdisk $partsize || return 1
}

dev="$1"
sdisk="$2"
[[ "$dev" == "" ]] || [[ "$sdisk" == "" ]] && usage

xcrypto=""
[ -f /etc/sdmcrypto ] && read xcrypto < /etc/sdmcrypto
[ "$xcrypto" == "" ] && xcrypto="aes"
[[ "aes|xchacha" =~ $xcrypto ]] || [[ "$xcrypto" =~ aes- ]] || errexit "? Supported cryptos: 'aes' and 'xchacha'"

[ -f /etc/sdmnopwd ] && xnopwd=1 || xnopwd=0

partsize=""
cryptdevice="nomappername"
datefmt="%Y-%m-%d %H:%M:%S"
[ -f /etc/mappername ] && read cryptdevice < /etc/mappername || echo "% Where is /etc/mappername?; Using cryptdevice 'nomappername'"
runall $dev $sdisk 
wipescratch $partsize

sleepsecs=10
if [ "$SSH_CLIENT" == "" ]
then
    echo $"
Enter the 'exit' command to resume the system boot process
"
else
    echo $"
Encryption complete. System boot will continue in $sleepsecs seconds
"
    sleep $sleepsecs
    kill -9 -1
fi
exit
