#!/bin/sh

set -x # script debugging; show commands
set -e # exit on error

# This is to be used by a buildroot minimal rootfs on the target board
# as an init script to:
#   - provision flash filesystem
#   - mount flash filesystem
#
# the flash storage (ie create disk partition, filesystems, with/without encryption, and populate them)

# Some basic vars
DEVICE=mmcblk2 # storage device to be used for the encrypted file system
DEV=/dev/$DEVICE
P1_OFFSETMB=64
PART=p1
NAME=rootfs # name for the encrypted file system (only used here)
KEY=/fs.key
NIC=eth0
INIT=/sbin/init # init of final OS
SERVERIP=172.24.21.238
# Boot firmware support
BOOT_FIRMWARE=http://$SERVERIP/tftpboot/venice/signed-flash.bin
FIT_IMAGE=http://$SERVERIP/tftpboot/venice/fit.itb
# TPM support
KEY_HANDLE=0x81234567 #TPM2 Owner Persistent Handle Range: 0x81000000 - 0x81800000

# Mount things needed by this script
mount -n -t devtmpfs devtmpfs /dev
mount -n -t proc proc /proc
mount -n -t sysfs sysfs /sys
mount -n -t tmpfs tmpfs /run

# Check if "tpm" is present in bootargs
for arg in $(cat /proc/cmdline); do
        [ "$arg" == "tpm" ] && {
                insmod /virt-dma.ko
                insmod /imx-sdma.ko
                insmod /spi-imx.ko
                insmod /tpm_tis_spi.ko
                [ -c /dev/tpm0 ] && {
                        TPM=/dev/tpm0
                        echo "Using TPM: $TPM"
                }
        }
done

user_data() {
        echo "Applying user changes..."

        # in any board specific changes to rootfs
        case "$(cat /proc/device-tree/board)" in
                GW74*)
                        mkdir -p /mnt/etc/udev/rules.d
                        echo 'SUBSYSTEM=="net", ACTION=="add", DEVPATH=="/devices/platform/soc@0/30800000.bus/30bf0000.ethernet/net/e
th*", NAME="en0"' > /mnt/etc/udev/rules.d/70-persistent-net.rules
                        echo 'SUBSYSTEM=="net", ACTION=="add", DEVPATH=="/devices/platform/soc@0/30800000.bus/30be0000.ethernet/net/e
th*", NAME="en1"' >> /mnt/etc/udev/rules.d/70-persistent-net.rules
                        sed -i 's/eth0/en0/' /mnt/etc/network/interfaces
                        ;;
        esac
}

# example1: rootfs via a list of tarballs
#  - gateworks ubuntu tarbarll and linux kernel tarball
ROOTFS_TARS="\
        https://dev.gateworks.com/ubuntu/jammy/jammy-venice.tar.xz \
        https://dev.gateworks.com/venice/kernel/linux-venice-6.6.8.tar.xz \
"
rootfs1() {
        # Create an ext4 file system on the opened LUKS device
        echo "Creating filesystem..."
        mkfs.ext4 -q -F -L rootfs /dev/mapper/$NAME

        # Mount the LUKS device to /mnt
        mount /dev/mapper/$NAME /mnt

        for url in $ROOTFS_TARS; do
                echo "Downloading $url..."
                wget -q --show-progress --no-check-certificate -O data $url
                echo "Extracting..."
                pv data | tar -C /mnt -mxJ --keep-directory-symlink
        done

        user_data

        # Unmount the LUKS device
        echo "Unmounting LUKS device..."
        umount /dev/mapper/$NAME
}

# example2: rootfs via compressed disk image
rootfs2() {
        FILESYSTEM=https://dev.gateworks.com/venice/images/jammy-venice.img.gz
        # a compressed disk image contains a partition table as well as
        # partitions. Here we assume an MBR partition with the OS on P1.
        # we will extract the MBR, determine the start/end of P1 and copy
        # it to the LUKS device.
        echo "Downloading $FILESYSTEM..."
        wget -q --show-progress --no-check-certificate -O data $FILESYSTEM
        # Extract the MBR
        zcat data | dd of=mbr count=1
        start=$(sfdisk -l mbr | tail -1 | tr -s ' ' | cut -d ' ' -f3)
        end=$(sfdisk -l mbr | tail -1 | tr -s ' ' | cut -d ' ' -f4)
        # copy the filesystem in P1 to the LUKS device
        echo "Copying..."
        pv data | zcat | dd of=/dev/mapper/$NAME skip=$start count=$((end-start))

        # apply user data
        mount /dev/mapper/$NAME /mnt
        user_data
        umount /dev/mapper/$NAME
}

# example3: rootfs via existing filesystem
rootfs3() {
        FILESYSTEM=https://dev.gateworks.com/buildroot/venice/minimal/rootfs.ext2.xz
        #FILESYSTEM=https://$SERVERIP/tftpboot/venice/rootfs.ext2.xz
        echo "Downloading $FILESYSTEM..."
        wget -q --show-progress --no-check-certificate -O data $FILESYSTEM
        # copy the filesystem in P1 to the LUKS device
        echo "Copying..."
        pv data | xzcat | dd of=/dev/mapper/$NAME bs=16M

        # apply user data
        mount /dev/mapper/$NAME /mnt
        user_data
        umount /dev/mapper/$NAME
}

# provision the flash with secure boot firmware and encrypted rootfs
provision() {
        rootfs=${1:-rootfs1}

        echo "Provision $DEV via $rootfs..."

        # board specific customizations
        case "$(cat /proc/device-tree/board)" in
                GW74*) NIC=eth1;;
        esac

        # soc specific customizations
        case "$(cat /proc/device-tree/soc@0/compatible | tr \0 \n)" in
                *imx8mm*)
                        boot0_offset_kb=33
                        user_offset_kb=33
                        ;;
                *imx8mn*|*imx8mp*)
                        boot0_offset_kb=0
                        user_offset_kb=32
                        ;;
        esac

        # Mount the /tmp directory as a tmpfs with 75% of the available memory
        mount -t tmpfs /tmp -o size=75%

        # Change the current working directory to /tmp
        cd /tmp

        # bring up network
        echo "Bringing up network $NIC..."
        udhcpc -i $NIC

        [ "$BOOT_FIRMWARE" ] && {
                echo "Writing locked down U-Boot to memory..."
                wget -q --show-progress --no-check-certificate -O data $BOOT_FIRMWARE
                # disable boot0 readonly
                echo 0 > /sys/class/block/${DEVICE}boot0/force_ro
                # write boot firmware to proper offset
                dd if=data of=${DEV}boot0 bs=1K seek=$boot0_offset_kb oflag=sync
        }

        [ "$FIT_IMAGE" ] && {
                echo "Flashing FIT Image to memory..."
                wget -q --show-progress --no-check-certificate -O data $FIT_IMAGE
                # sanity check its not too big for P1_OFFSETMB
                #  - we don't want to automatically choose P1_OFFSETMB as you may want to support
                #    firmware updates and you should allow it to grow and thus choose your slack space
                FIT_SIZEMB=$(($(stat -c '%s' data) / 1024 / 1024))
                [ $FIT_SIZEMB -lt $(($P1_OFFSETMB + 1)) ] || { echo "Error: $FIT_IMAGE is too large; adjust P1_OFFSETMB"; exit 1; }
                # Write fit Image to memory
                dd if=data of=$DEV bs=1M seek=1 oflag=sync
        }

        echo "Creating partition table..."
        # partition disk:
        #   - you could use MBR or GPT based on your needs
        #   - ensure that you do not start your partition data until after your boot firmware
        # Create MBR with a single partition taking up the entire device
        printf "$((P1_OFFSETMB*1024*1024/512)),,L,*" | sfdisk -u S $DEV
        # force re-read of partition table now
        partprobe $DEV
        
        [ "$TPM" ] && {
                # Get inital PCR measurment values
                echo "Retrieving initial PCR measurement values..."
                # PCR0: boot-firmware
                dd if=${DEV}boot0 of=data bs=1M count=4
                tpm2 pcrextend 0:sha1=$(sha1sum data | cut -d" " -f1)
                # PCR8: partition table and FIT image (kernel, dtb, ramdisk)
                dd if=$DEV of=data bs=1M count=$P1_OFFSETMB
                tpm2 pcrextend 8:sha1=$(sha1sum data | cut -d" " -f1)
                # show PCR's
                tpm2 pcrread

                echo "Creating Key from TPM2.0..."
                hexdump -C $KEY # if debugging
                # delete any existing key
                tpm2 evictcontrol -C o -c $KEY_HANDLE
                # create a policy that depends on PCR0 and PCR8
                tpm2 createpolicy --policy-pcr -l sha1:0,8 -L policy.digest
                tpm2 createprimary -g sha256 -G rsa -c primary.context
                # create an object containing the key
                tpm2 create -g sha256 -u obj.pub -r obj.priv -C primary.context -L policy.digest -a "noda|adminwithpolicy|fixedparent
|fixedtpm" -i $KEY
                # load it
                tpm2 load -C primary.context -u obj.pub -r obj.priv -c load.context
                # save it to TPM handle
                tpm2 evictcontrol -C o -c load.context $KEY_HANDLE
                rm -f policy.digest primary.context obj.pub obj.priv load.context
        }

        # Format the partition as a LUKS encrypted file system using the encryption key
        echo "Formating LUKS device..."
        cryptsetup --batch-mode luksFormat ${DEV}${PART} --key-file=$KEY

        # Open the LUKS encrypted partition and map it to the specified name ($NAME)
        echo "Opening LUKS device..."
        cryptsetup luksOpen ${DEV}${PART} $NAME --key-file=$KEY

        # call rootfs option
        $rootfs 

        # Close the LUKS device and remove the mapping
        echo "Closing LUKS device..."
        cryptsetup luksClose $NAME

        echo "Provisioning complete"
        sleep 2

        # Wait forever
        #while [ 1 ]; do sleep 1; done

        # power cycle
        echo 2 > /sys/bus/i2c/devices/0-0020/powerdown
}

# boot to encrypted fs
boot() {
        echo "Using $DEV..."

        # Wait for device to exist
        echo "Waiting for $DEV..."
        while [ ! -b "$DEV" ]; do
                sleep 1
                echo -n .
        done

        # Open the LUKS encrypted partition and map it to the specified name ($NAME)
        [ "$TPM" ] && {
                echo "Fetching Key from TPM..."    
                KEY=fs.keyphrase_decrypted
                # PCR0: boot-firmware
                dd if=${DEV}boot0 of=data bs=1M count=4
                tpm2 pcrextend 0:sha1=$(sha1sum data | cut -d" " -f1)
                # PCR8: partition table and FIT image (kernel, dtb, ramdisk)
                dd if=$DEV of=data bs=1M count=$P1_OFFSETMB
                tpm2 pcrextend 8:sha1=$(sha1sum data | cut -d" " -f1)
                # show PCR's
                tpm2 pcrread
                # unseal our key
                tpm2 unseal -c $KEY_HANDLE -p pcr:sha1:0,8 -o $KEY
                hexdump -C $KEY # if debugging
        }

        DEV=${DEV}${PART}
        echo "Opening $DEV..."
        cryptsetup luksOpen $DEV $NAME --key-file=$KEY

        [ "$TPM" ] && {
                # reseal the key by extending the PCR's
                rm -f $KEY
                tpm2 pcrextend 0:sha1=0000000000000000000000000000000000000000
                tpm2 pcrextend 8:sha1=0000000000000000000000000000000000000000
        }

        # Mount the LUKS device to /mnt
        echo "Mounting $DEV..."
        mount /dev/mapper/$NAME /mnt

        # Switch to the new root and execute init
        echo "Switching to new root and running $INIT $@..."
        cd /mnt
        [ -c dev/console ] || mknod -m 600 dev/console c 5 1
        exec switch_root . "$INIT" "$@"

        # This will only be run if the above line failed
        echo "Failed to switch_root"
}

# boot straight to ramdisk (no disk encryption)
bypass() {
        echo "Booting straight to ramdisk..."
        umount /proc
        umount /sys
        umount /run
        if (exec 0</dev/console) 2>/dev/null; then
                exec 0</dev/console
                exec 1>/dev/console
                exec 2>/dev/console
        fi
        exec /sbin/init "$@"
}

for x in $(cat /proc/cmdline); do
        case "$x" in
                bypass) bypass;;
                provision=*) provision ${x//provision=};;
        esac
done

boot