sdmcryptfs

#!/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 errifrc() {
    [ $1 -ne 0 ] && errexit "$2"
}

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

function pmsg() {
    local msg="$1" tty="/dev/console"
    [ $fssh -eq 1 ] && tty="/dev/tty"
    echo "> $(thisdate) $msg" >$tty
}

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

function dbgprint() {
    [ $fdebug -eq 1 ] && pmsg "D! $1"
}

function dbgshell() {
    local msg="${1:-Debug}" tty="/dev/console"
    [ $fssh -eq 1 ] && tty="/dev/tty"
    echo "dbgshell: $msg"
    bash <$tty >$tty
    echo "dbgexit: $msg"
}

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 printdt() {
    #
    # Computes time from seconds ($1) and returns it as a string of Days:Hours:Minutes:Seconds
    #
    local days hours minutes seconds eseconds="$1"
    local btime="$1" etime="$2"
    days=$((eseconds/86400))
    hours=$(((eseconds - $((days*86400)))/3600))
    minutes=$(((eseconds - $((days*86400)) - $((hours*3600)))/60))
    seconds=$((eseconds - $((days*86400)) - $((hours*3600)) - $((minutes*60))))
    [ $days -ne 0 ] && printf "%02d:%02d:%02d:%02d\n" "$days" "$hours" "$minutes" "$seconds" || printf "%02d:%02d:%02d\n" "$hours" "$minutes" "$seconds"
}
 
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)))
    printdt "$eseconds"
}

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" =~ "nvme" ]] && return 0 || return 1
}

function isgpt() {
    local dev=$1
    [ "$(getdisktype $dev)" == "gpt" ] && return 0 || return 1
}

function getpartname() {
    local dev=$1 partnum=${2:-2}
    local p1=1 p2=2

    ispdevp $dev && p1="p1" && p2="p2"
    [ $partnum -eq 1 ] && echo "${dev}$p1"
    [ $partnum -eq 2 ] && echo "${dev}$p2"
}

function getdisktype() {
    local dev=$1
    echo "$(lsblk -o PTTYPE --noheadings $(getpartname $dev 1))"
}

function expandpartition() {
    # $1: devicename (/dev/sdX)
    # $2: 0 (fill available space) or nnnn (expand by MB)
    # $3: message output (write_burnmsg/logtoboth/or similar)
    # $4: fs type (ext4, btrfs, or lvm)
    local fulldevname="$1" sexp=${2:-0} msgrtn=${3:-logtoboth} rootfstype=${4:-ext4}
    local phydev=${fulldevname##/dev/}
    local partname fsize partstart partend devsize partsize newend pbytes nbytes pend

    partname=$(getpartname $phydev 2)
    devsize=$(blockdev --getsz $fulldevname)
    partsize=$(blockdev --getsz /dev/$partname)
    pbytes=$((partsize*512))
    if [ $sexp -eq 0 ]
    then
	devsize=$((devsize-(33*2)))  # Eliminate sgdisk -p 33 'block' complaint
	nbytes=$(((devsize*512)-partend+pbytes))
	newend=$((devsize-1))                  # 512byte blocks
    else
	nbytes=$(((pbytes+(sexp*1048576)))) #1024*1024
	newend=$(((nbytes+partstart)/512))
    fi
    [ "$plugindebug" == "" ] && plugindebug=0
    if [ $plugindebug -eq 1 ]
    then
	$msgrtn "> Plugin ${pfx}ExpandPartition: devsize: $devsize"
	$msgrtn "> Plugin ${pfx}ExpandPartition: partsize: $partsize"
	$msgrtn "> Plugin ${pfx}ExpandPartition: nbytes: $nbytes"
    fi
    $msgrtn "Expand root partition '$partname' ($rootfstype) on device '$fulldevname' from $(getgbstr $pbytes) to $(getgbstr $nbytes)"
    if isgpt ${fulldevname}
    then
	parted -m /dev/$phydev u s resizepart 2 $((newend-48))  || errexit "? parted resizepart error '$?'" # ?? 32 not enough, 48 works
    else
	parted -m /dev/$phydev u s resizepart 2 $newend || errexit "? parted resizepart error '$?'"
    fi
}

function getpartsize() {
    local rootpart=$1
    local pat r2fs partsize

    if [ "$(lsblk --noheadings --output fstype $rootpart)" == "ext4" ]
    then
	pat='([[:space:]][[:digit:]]+)'
	r2fs=$(resize2fs -fM -p $rootpart 2>&1)
	r2fs=${r2fs##* is}
	[[ $r2fs =~ $pat ]]
	partsize=${BASH_REMATCH[1]}
	echo $partsize
    else
	# Get partition size 4k blocks
	partsize=$(($(blockdev --getsize64 $rootpart)/4096))
	echo "$partsize"
    fi
}

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 "Securely wipe scratch disk '$sdisk'"
    then
	pmsg "> Write random data on '$sdisk'"
	starttime="$(getcbbdate)"
	dd bs=4k count=$partsize if=/dev/random of=$sdisk || errexit "? dd write random data error '$?'"
	pmsg "rootfs Wipe elapsed time: $(datediff "$starttime" "$(getcbbdate)")"
    else
	pmsg "> Zap partitions on scratch disk '$sdisk'"
	sgdisk --zap-all $sdisk >/dev/null 2>&1 || errexit "? sgdisk zap all error '$?'"
	pmsg "> The partition table on '$sdisk' has been zapped"
	pmsg "> Scratch disk '$sdisk' still has unencrypted blocks of rootfs on the disk"
	pmsg "  Clean it up after rootfs encryption 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 2)
    local scpart=$(getpartname $sdisk 1)
    local partend fullblocks remainder rootfsize mb
    local starttime endtime pbytes pgb

    mkdir -p /scratch
    pmsg "** NOTE: save and restore time estimates are pessimistic"
    pmsg "** Exact speed measurements dependent on CPU, memory, and disk speeds of this system"

    if [ $fram -eq 1 ]
    then
	mount -t tmpfs -o size=${rootfsize}M tmpfs /scratch || errexit "? mount error '$?'"
	pmsg "Save rootfs '$rootfs' to RAM disk"
	pmsg "rootfs save should take less than $(printdt $(((partsize/256)/43)))"
	xcompress=0
    else
	checkdev $sdisk || return 1
	pmsg "Create file system on '$sdisk'"
	parted -s $sdisk mklabel gpt || errexit "? parted create gpt label error '$?'"
	partend=$(((partsize*8)+2048))
	sgdisk --new 1:2048:$partend $sdisk >/dev/null || errexit "? sgdisk create partition error '$?'"
	sgdisk --typecode 1:8300 $sdisk || errexit "? sgdisk set typecode error '$?'"
	mke2fs -F $scpart >/dev/null || errexit "? mke2fs error '$?'"
	mount $scpart /scratch || errexit "? mount scratch partition error '$?'"
	pmsg "Save rootfs '$rootfs' to '$sdisk'"
	pmsg "rootfs save should take less than $(printdt $(((partsize/256)/51)))"
    fi
    starttime="$(getcbbdate)"
    if [ $xcompress -eq 0 ]
    then
	# no compression
	dd bs=4k count=$partsize if=$rootfs of=/scratch/uncompressed-rootfs.img || errexit "? dd error '$?'"
    else
	# use compression
	fullblocks=$((partsize/4096))   #1M blocks
	remainder=$((partsize % 4096))  #   remainder bits
	pmsg "Compress and save $fullblocks 16MB blocks"
	dd bs=16M count=$fullblocks if=$rootfs | lz4 -c --fast=1 > /scratch/compressed-rootfs.lz4 || errexit "? dd error '$?'"
	if [ $remainder -ne 0 ]
	then
	    pmsg "Compress and save $remainder 4K blocks"
	    dd bs=4k if=$rootfs skip=$((fullblocks*4096)) count=$remainder | lz4 -c --fast=1 >> /scratch/compressed-rootfs.lz4 || errexit $? "? dd error '$?'"
	fi
    fi
    pmsg "rootfs Save elapsed time: $(datediff "$starttime" "$(getcbbdate)")"
    return 0
}

function find_keyfile() {
    local kfl=""

    IFS=":" read lval kfl < <(grep keyfile: /etc/sdmcryptctl)
    [ "$kfl" == "" ] && echo "" || echo "/etc/$kfl"
}

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

    kfl=$(find_keyfile)
    [ "$kfl" == "" ] && return || kfn=$(basename $kfl)
    if [ "$kfn" != "" ]
    then
	pmsg "Add LUKS key '${kfn%.lek}' from file $kfl"
	pmsg "When prompted for an existing passphrase enter the one you just created"
	cryptsetup luksAddKey $rootfs $kfl || errexit "? cryptsetup add key '$kfn' error: '$?'"
	rm -f $kfl
    fi
}

function cpoptions() {
    local sts=$1

    pmsg "% cryptsetup error ($sts)"
    pmsg "% If you mis-typed 'YES' or want to try again enter Y or y"
    pmsg "% Power off the system if you do not want to retry"
    pmsg "%  (Disks will be OK since none are mounted)"
    pmsg "% The only way to 'reset' this disk is to reburn it"
    pmsg "% and start over"
    if ! askyn "Retry"
    then
	pmsg "OK. Exit sdmcryptfs"
	exit 1
    fi
    return 0
}

function cryptpart() {
    local dev=$1
    local cipher sts
    case "$xcrypto" in
	xchacha*)
	    [ "$xcrypto" == "xchacha" ] && xcha=xchacha12 || xcha=$xcrypto
	    cipher="$xcha,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"
    pmsg ""

    if [ $xnopwd -eq 1 ]
    then
	while [ 0 ]
	do
	    printf "YES\n" | cryptsetup luksFormat --type luks2 $rootfs $(find_keyfile)
	    sts=$?
	    [ $sts -eq 0 ] && break
	    cpoptions $sts || { pmsg "Exit sdmcryptfs" ; exit 1 ; }
	done
    else
	while [ 0 ]
	do
	    cryptsetup luksFormat --type luks2 --cipher $cipher --hash sha256 --iter-time 5000 --key-size 256 --pbkdf argon2i $rootfs
	    sts=$?
	    [ $sts -eq 0 ] && break
	    cpoptions $sts || { pmsg "Exit sdmcryptfs" ; exit 1 ; }
	done
	config_keyfile $dev
    fi
    return 0
}

function cryptopen() {
    local dev=$1
    local rootfs=$(getpartname $dev 2)
    pmsg "Unlock encrypted partition '$rootfs'"
    pmsg "Unlock will take several seconds"
    if [ $xnopwd -eq 1 ]
    then
	while [ 0 ]
	do
	    cryptsetup luksOpen $rootfs $cryptdevice --key-file $(find_keyfile)
	    [ $? -eq 0 ] && break
	done
    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 fstype=$4
    local rootfs=$(getpartname $dev 2)
    local cpartnum cpartstart cpartend cpartsize etc
    local starttime endtime pbytes pgb

    [ $fram -eq 0 ] && { checkdev $sdisk || return 1 ; }
    pmsg "Restore '$rootfs' from '$sdisk'"

    starttime="$(getcbbdate)"
    if [ $xcompress -eq 0 ]
    then
	pmsg "rootfs restore should take less than $(printdt $(((partsize/256)/15)))"
	#echo "dd bs=16M if=/scratch/uncompressed-rootfs.img of=/dev/mapper/$cryptdevice"
	dd bs=16M if=/scratch/uncompressed-rootfs.img of=/dev/mapper/$cryptdevice || errexit "? dd error '$?'"
    else
	pmsg "rootfs restore should take less than $(printdt $(((partsize/256)/11)))"
	#echo "dd bs=16M if=/scratch/compressed-rootfs.lz4 | lz4 -d | bs=16M dd of=/dev/mapper/$cryptdevice"
	dd bs=16M if=/scratch/compressed-rootfs.lz4 | lz4 -d | dd bs=16M of=/dev/mapper/$cryptdevice || errexit "? dd error '$?'"
    fi
    pmsg "rootfs Restore elapsed time: $(datediff "$starttime" "$(getcbbdate)")"
    pmsg "Restore complete"
    umount /scratch >/dev/null 2>&1
    case  "$fstype"
    in
	ext4)
	    pmsg "Expand ext4 rootfs file system..."
	    resize2fs -f /dev/mapper/$cryptdevice || errexit "? resize2fs error '$?'"
	    cpartsize=$(blockdev --getsize64 /dev/mapper/$cryptdevice)
	    pmsg "rootfs partition size: $cpartsize Bytes $(getgbstr $cpartsize)"
	    ;;
	btrfs)
	    pmsg "Expand btrfs rootfs file system..."
	    mkdir -p /mnt
	    #dbgshell "restore_rootfs before btrfs expand"
	    mount -t btrfs /dev/mapper/$cryptdevice /mnt
	    btrfs filesystem resize max /mnt || errexit "? brtfs resize error '$?'"
	    umount /mnt
	    ;;
    esac
    pmsg "cryptsetup status for /dev/mapper/$cryptdevice:"
    cryptsetup status /dev/mapper/$cryptdevice
    return 0
}

function runall() {
    local dev=$1 sdisk=$2 pbytes pgb
    local zpartnum zpartstart zpartend zpartsize etc
    local rootfs=$(getpartname $dev 2)
    fstype=$(lsblk --noheadings --output fstype $rootfs)
    checkdev $dev || return 1
    [ $fram -eq 0 ] && { checkdev $sdisk || return 1 ; }
    [ "$fstype" == "ext4" ] && 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
    if grep -q -s xnoexpand:0 /etc/sdmcryptctl
    then
	# Want to expand the partition
	IFS=":" read zpartnum zpartstart zpartend zpartsize etc < <(parted -ms $dev unit s print | tail -n 1)
	if [ "$zpartnum" == "2" ]
	then
	    # Last partition on the disk so take it all
	    pmsg "Expand rootfs partition with expandpartition"
	    expandpartition $dev 0 pmsg $fstype
	else
	    # Partition following rootfs so can expand up to it's start -1
	    IFS=":" read zpartnum zpartstart zpartend zpartsize etc < <(parted -ms $dev unit s print | grep ^3:)
	    IFS=":" read rpartnum rpartstart rpartend rpartsize etc < <(parted -ms $dev unit s print | grep ^2:)
	    newend=$((${zpartstart%s}-1))
	    #pmsg "DD2 rpartnum:|$rpartnum| rpartstart:|$rpartstart| rpartend:|$rpartend| rpartsize:|$rpartsize|"
	    #pmsg "DD3 zpartnum:|$zpartnum| zpartstart:|$zpartstart| zpartend:|$zpartend| zpartsize:|$zpartsize|"
	    pmsg "Expand rootfs partition from ${rpartend%s} to $((${zpartstart%s}-1)) (increase by $((${zpartstart%s}-${rpartend%s}-1)))"
	    parted -m $dev u s resizepart 2 $newend || errexit "? parted resize partition error '$?'"
	fi
    else
	if [ "$fstype" == "btrfs" ]
	then
	    # Do minimal expansion (so there's room for btrfs + LUKS header) using space allocated in parted or bc it's the rest of the disk
	    IFS=":" read rpartnum rpartstart rpartend rpartsize etc < <(parted -ms $dev unit s print | grep ^2:)
	    newend=$((${rpartend%s}+65535))
	    #pmsg "DD4 rpartnum:|$rpartnum| rpartstart:|$rpartstart| rpartend:|$rpartend| rpartsize:|$rpartsize| newend:|$newend|"
	    pmsg "Do minimal rootfs expansion so btrfs and LUKS header both fit"
	    parted -m $dev u s resizepart 2 $newend || errexit "? parted resize partition error '$?'"
	fi
    fi
    cryptpart $dev || return 1
    cryptopen $dev || return 1
    restore_rootfs $dev $sdisk $partsize $fstype || return 1
}

fauto=0
fram=0
fispi=0
fdebug=1
dev="$1"
sdisk="$2"
datefmt="%Y-%m-%d %H:%M:%S"
[ "$SSH_CLIENT" == "" ] && fssh=0 || fssh=1
if grep -q -s autosysdisk: /etc/sdmcryptctl
then
    fauto=1
    IFS=":" read label dev < <(grep autosysdisk: /etc/sdmcryptctl)
    IFS=":" read label sdisk < <(grep autoscratch: /etc/sdmcryptctl)
fi
[[ "$dev" == "" ]] || [[ "$sdisk" == "" ]] && usage

[ -e /sys/firmware/devicetree/base/model ] && pimodel=$(tr -d '\0' < /sys/firmware/devicetree/base/model) || pimodel=""
IFS=" " read ign pimem units <<< $(grep -i MemTotal /proc/meminfo)
[ "$pimodel" != "" ] && [ "$pimem" != "" ] && pmsg "Start sdmcryptfs on a $pimem $units '$pimodel'"

#
# Determine if we can use a RAM disk
#
if [ "$pimodel" != "" ]
then
    fispi=1
    read junk mb junk < <(grep -i MemTotal /proc/meminfo)
    partsize=$(getpartsize $(getpartname $dev 2))
    rootfsize=$((((partsize/256)*6)/5))    # convert to MB and scale to 1.2X
    mb=$(((mb/1024)-3))            # Memory size in MB minus a bit
    [ $mb -gt $rootfsize ] && fram=1 && xcompress=0 && sdisk=RAM
fi

if [ $fauto -eq 1 ]
then
    pmsg ""
    pmsg "************"
    pmsg "Automatic sdmcryptfs conversion on '$dev' with '$sdisk' as scratch disk"
    pmsg "sdm Disk Encryption process will continue automatically when sdmcryptfs completes"
    if [ $fram -eq 0 ]
    then
	cdelay=15
	pmsg "Conversion will start in $cdelay seconds"
	pmsg "*** Ensure that the scratch disk '$sdisk' is inserted ***"
	pmsg "    If it's not inserted, we'll wait for it"
	pmsg "************"
	pmsg ""
	sleep $cdelay
	until [ -b $sdisk ]
	do
	    pmsg "Waiting for scratch disk '$sdisk'"
	    sleep 1
	done
    else
	pmsg "************"
	pmsg ""
    fi
fi

xcrypto=""
IFS=":" read lval xcrypto < <(grep xcrypto: /etc/sdmcryptctl)
[ "$xcrypto" == "" ] && xcrypto="aes"
[[ "aes|xchacha|xchacha12|xchacha20" =~ $xcrypto ]] || [[ "$xcrypto" =~ aes- ]] || errexit "? Unrecognized crypto '$xcrypto'"

IFS=":" read lval xnopwd < <(grep xnopwd: /etc/sdmcryptctl)
[ "$xnopwd" == "" ] && xnopwd=0

IFS=":" read lval xcompress < <(grep xcompress: /etc/sdmcryptctl)
[ "$xcompress" == "" ] && xcompress=1

partsize=""
IFS=":" read lval cryptdevice < <(grep mappername: /etc/sdmcryptctl)
[ "$cryptdevice" == "" ] && cryptdevice="nomappername"

runall $dev $sdisk || exit 1

if [ $fauto -eq 1 ]
then
    [ $fram -eq 0 ] && wipescratch $partsize
    echo $"
Encryption complete. System boot will continue automatically in 5 seconds
"
    sleep 5
    exit 0
fi

sleepsecs=10
if [ $fssh -eq 0 ]
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