#!/bin/bash
set -e
##### CHECK PARAMETERS #####
PRODUCTION=0
while [[ "$#" -gt 0 ]]; do
    case $1 in
        -p|--production) PRODUCTION=1 ;;
        *) echo "Unknown parameter passed: $1"; exit 1 ;;
    esac
    shift
done
if [[ "${PRODUCTION}" -ne "0" ]]; then
  echo "====================================="
  echo "========== PRODUCTION MODE =========="
  echo "====================================="
else
  echo "======================================"
  echo "============ TESTING MODE ============"
  echo "======================================"
fi
##### SETUP SYSTEM #####
echo "Stopping Services"
sudo systemctl stop udisks2
sudo -s echo -1 > /sys/module/usbcore/parameters/autosuspend
sudo ufw disable
echo "Installing packages"
sudo apt-get install dislocker cryptsetup libcryptsetup-dev libcryptsetup12 cryptmount cryptmount overlayroot qemu-user-static pip python-is-python3
pip install cryptography
pip install pycrypto
pip install pycryptodome
echo "Creating Directory"
mkdir tmp
cd tmp
echo "Downloading Packages"
wget https://developer.nvidia.com/downloads/embedded/l4t/r35_release_v4.1/release/jetson_linux_r35.4.1_aarch64.tbz2
wget https://developer.nvidia.com/downloads/embedded/l4t/r35_release_v4.1/release/tegra_linux_sample-root-filesystem_r35.4.1_aarch64.tbz2
wget https://developer.nvidia.com/downloads/embedded/l4t/r35_release_v4.1/sources/public_sources.tbz2
echo "Unpacking Packages"
tar xvf jetson_linux_r35.4.1_aarch64.tbz2
sudo tar xvf tegra_linux_sample-root-filesystem_r35.4.1_aarch64.tbz2 -C Linux_for_Tegra/rootfs/
tar xvf public_sources.tbz2
cd Linux_for_Tegra/source/public/
tar xvf nvidia-jetson-optee-source.tbz2
cd ../..
echo "Running prerequeisites script"
sudo tools/l4t_flash_prerequisites.sh
echo "Applying Binaries"
sudo ./apply_binaries.sh
##### GENERATE A PKC KEY PAIR #####
if [ ! -f rsa.pem ]; then
  echo "Generating PKCS key pair..."
  openssl genrsa -out rsa.pem 3072
fi
PKCS_KEY_XML_HASH=$(./bootloader/tegrasign_v3.py --pubkeyhash rsa.pubkey rsa.hash --key rsa.pem | grep "tegra-fuse format" | awk '{print $NF}')
echo "PKCS Key Hash: ${PKCS_KEY_XML_HASH}"
##### PREPARE AN SBK KEY #####
if [ ! -f sbk.key ] || [ ! -f sbk_xml.key ]; then
  echo "Generating SBK key..."
  SBK_0=$(openssl rand -hex 4)
  SBK_1=$(openssl rand -hex 4)
  SBK_2=$(openssl rand -hex 4)
  SBK_3=$(openssl rand -hex 4)
  SBK_4=$(openssl rand -hex 4)
  SBK_5=$(openssl rand -hex 4)
  SBK_6=$(openssl rand -hex 4)
  SBK_7=$(openssl rand -hex 4)
  SBK_KEY=$(echo "0x${SBK_0} 0x${SBK_1} 0x${SBK_2} 0x${SBK_3} 0x${SBK_4} 0x${SBK_5} 0x${SBK_6} 0x${SBK_7}")
  echo "${SBK_KEY}" > sbk.key
  SBK_KEY_XML="0x${SBK_0}${SBK_1}${SBK_2}${SBK_3}${SBK_4}${SBK_5}${SBK_6}${SBK_7}"
  echo "${SBK_KEY_XML}" > sbk_xml.key
else
  SBK_KEY=$(cat sbk.key)
  SBK_KEY_XML=$(cat sbk_xml.key)
fi
echo "SBK Key: ${SBK_KEY_XML}"
##### PREPARE KEK KEYS #####
if [ ! -f kek.key ] || [ ! -f kek_xml.key ]; then
  echo "Generating KEK key..."
  KEK_2_0=$(openssl rand -hex 4)
  KEK_2_1=$(openssl rand -hex 4)
  KEK_2_2=$(openssl rand -hex 4)
  KEK_2_3=$(openssl rand -hex 4)
  KEK_2_4=$(openssl rand -hex 4)
  KEK_2_5=$(openssl rand -hex 4)
  KEK_2_6=$(openssl rand -hex 4)
  KEK_2_7=$(openssl rand -hex 4)
  KEK_2_KEY=$(echo "0x${KEK_2_0} 0x${KEK_2_1} 0x${KEK_2_2} 0x${KEK_2_3} 0x${KEK_2_4} 0x${KEK_2_5} 0x${KEK_2_6} 0x${KEK_2_7}")
  echo "${KEK_2_KEY}" > kek.key
  KEK_2_KEY_XML="0x${KEK_2_0}${KEK_2_1}${KEK_2_2}${KEK_2_3}${KEK_2_4}${KEK_2_5}${KEK_2_6}${KEK_2_7}"
  echo "${KEK_2_KEY_XML}" > kek_xml.key
  KEK_2_KEY_OPTEE="${KEK_2_0}${KEK_2_1}${KEK_2_2}${KEK_2_3}${KEK_2_4}${KEK_2_5}${KEK_2_6}${KEK_2_7}"
  echo "${KEK_2_KEY_OPTEE}" > kek_optee.key
else
  KEK_2_KEY=$(cat kek.key)
  KEK_2_KEY_XML=$(cat kek_xml.key)
  KEK_2_KEY_OPTEE=$(cat kek_optee.key)
fi
echo "KEK1 Key: ${KEK_2_KEY_XML}"
##### GENERATE FUSE XML #####
echo "<genericfuse MagicId=\"0x45535546\" version=\"1.0.0\">" > fuse.xml
echo "  <fuse name=\"PublicKeyHash\" size=\"64\" value=\"${PKCS_KEY_XML_HASH}\"/>" >> fuse.xml
echo "  <fuse name=\"SecureBootKey\" size=\"32\" value=\"${SBK_KEY_XML}\"/>" >> fuse.xml
echo "  <fuse name=\"OemK2\" size=\"32\" value=\"${KEK_2_KEY_XML}\"/>" >> fuse.xml
echo "  <fuse name=\"BootSecurityInfo\" size=\"4\" value=\"0x209\"/>" >> fuse.xml
if [[ "${PRODUCTION}" -ne "0" ]]; then
  echo "  <fuse name=\"SecurityMode\" size=\"4\" value=\"0x1\"/>" >> fuse.xml
fi
echo "</genericfuse>" >> fuse.xml
echo "==============================================================================="
cat fuse.xml
echo "==============================================================================="
##### GENERATE OPTEE IMAGE #####
echo "Generating OpTee image"
echo "bad66eb4484983684b992fe54a648bb8" > fv_ekb_t234
echo "010203040506070809a0b0c0d0e0f001" > sym_t234.key
echo "f0e0d0c0b0a001020304050607080900" > sym2_t234.key
python3 ./source/public/optee/samples/hwkey-agent/host/tool/gen_ekb/gen_ekb.py -chip t234 -oem_k2_key kek_optee.key -fv fv_ekb_t234 -in_sym_key sym_t234.key -in_sym_key2 sym2_t234.key -out bootloader/eks_t234.img
##### FUSE INSTRUCTIONS #####
echo "THIS WILL NOW FUSE THE ORIN NANO. THIS IS IRREVERSIBLE."
read -p "Press key to continue"
sed -i 's/getiterator/iter/g' bootloader/tegraflash_impl_t234.py
sudo ./odmfuse.sh -i 0x23 -k rsa.pem -S sbk.key -X fuse.xml jetson-orin-nano-devkit
##### QSPI SETUP #####
echo "Creating QSPI Image"
echo "Modify NUM_SECTORS in ./tools/kernel_flash/flash_l4t_t234_nvme_rootfs_enc.xml"
echo "For a 500GB drive, this should equal to `(500(size in GiB) * 1000 * 1000 * 1000) / 500(sector size)` = 1000000000."
echo "Put Jetson Orin Nano Devkit into recovery mode and plug it in"
read -p "Press key to continue"
sudo ./tools/kernel_flash/l4t_initrd_flash.sh --network usb0 -u ./rsa.pem -v ./sbk.key --no-flash --showlogs -p "-c bootloader/t186ref/cfg/flash_t234_qspi.xml" jetson-orin-nano-devkit internal
sudo cp bootloader/eks_t234_sigheader_encrypt.img.signed ./tools/kernel_flash/images/internal/
##### ROOTFS SETUP #####
echo "Creating RootFs Image"
echo "Put Jetson Orin Nano Devkit into recovery mode and plug it in"
echo "Please make sure you have modified the script for a good rootfs size of `-S 400Gib` is default and good for a 500GiB drive"
read -p "Press key to continue"
# Modify 400GiB rootfs size. There needs to be enough room inside NUM_SECTORS replacement, and the remaining partitions #
sudo ROOTFS_ENC=1 ./tools/kernel_flash/l4t_initrd_flash.sh --showlogs -u ./rsa.pem -v ./sbk.key --no-flash --external-device nvme0n1p1 -i ./sym2_t234.key -c ./tools/kernel_flash/flash_l4t_t234_nvme_rootfs_enc.xml -S 400GiB --external-only --append --network usb0 jetson-orin-nano-devkit external
##### FLASH #####
echo "Flash Image"
echo "Put Jetson Orin Nano Devkit into recovery mode and plug it in"
read -p "Press key to continue"
sudo ./tools/kernel_flash/l4t_initrd_flash.sh --showlogs -u rsa.pem -v sbk.key --network usb0 --flash-only