5 files changed, 52 insertions, 3 deletions
diff --git a/files/firstboot/21-generate-ssh-keys.sh b/files/firstboot/21-generate-ssh-keys.sh
index 928ee2e..b3345b4 100644
--- a/files/firstboot/21-generate-ssh-keys.sh
+++ b/files/firstboot/21-generate-ssh-keys.sh
@@ -1,8 +1,11 @@
 logger -t "rc.firstboot" "Generating SSH host keys"
 
 if [ -d "/etc/ssh/" ] ; then
+  # Remove ssh host keys
   rm -f /etc/ssh/ssh_host_*
   systemctl stop sshd
+
+  # Regenerate ssh host keys
   ssh-keygen -q -t rsa -N "" -f /etc/ssh/ssh_host_rsa_key
   ssh-keygen -q -t dsa -N "" -f /etc/ssh/ssh_host_dsa_key
   ssh-keygen -q -t ecdsa -N "" -f /etc/ssh/ssh_host_ecdsa_key
@@ -11,8 +14,11 @@ if [ -d "/etc/ssh/" ] ; then
 fi
 
 if [ -d "/etc/dropbear/" ] ; then
+  # Remove ssh host keys
   rm -f /etc/dropbear/dropbear_*
   systemctl stop dropbear
+
+  # Regenerate ssh host keys
   dropbearkey -t rsa -f /etc/dropbear/dropbear_rsa_host_key
   dropbearkey -t dss -f /etc/dropbear/dropbear_dss_host_key
   dropbearkey -t ecdsa -f /etc/dropbear/dropbear_ecdsa_host_key
diff --git a/files/firstboot/22-expandroot.sh b/files/firstboot/22-expandroot.sh
index 03ec50d..235c411 100644
--- a/files/firstboot/22-expandroot.sh
+++ b/files/firstboot/22-expandroot.sh
@@ -1,10 +1,20 @@
-logger -t "rc.firstboot" "Expanding root"
+logger -t "rc.firstboot" "Expanding root partition"
+
+# Detect root partition device
 ROOT_PART=$(mount | sed -n 's|^/dev/\(.*\) on / .*|\1|p')
-PART_NUM=$(echo ${ROOT_PART} | grep -o '[1-9][0-9]*$')
+if [ -z "$ROOT_PART" ] ; then
+  log_warning_msg "unable to detect root partition device"
+  return 1
+fi
+
+# Extract root device name
 case "${ROOT_PART}" in
   mmcblk0*) ROOT_DEV=mmcblk0 ;;
   sda*)     ROOT_DEV=sda ;;
 esac
+
+# Check detected root partition name
+PART_NUM=$(echo ${ROOT_PART} | grep -o '[1-9][0-9]*$')
 if [ "$PART_NUM" = "$ROOT_PART" ] ; then
   logger -t "rc.firstboot" "$ROOT_PART is not an SD card. Don't know how to expand"
   return 0
@@ -16,6 +26,8 @@ if [ "$PART_NUM" -gt 2 ] ; then
   logger -t "rc.firstboot" "Your partition layout is not currently supported by this tool."
   return 0
 fi
+
+# Check if last partition number
 LAST_PART_NUM=$(parted /dev/${ROOT_DEV} -ms unit s p | tail -n 1 | cut -f 1 -d:)
 if [ $LAST_PART_NUM -ne $PART_NUM ]; then
   logger -t "rc.firstboot" "$ROOT_PART is not the last partition. Don't know how to expand"
@@ -53,4 +65,4 @@ EOF2
 # Reload the partition table, resize root filesystem then remove resizing code from this file
 partprobe &&
   resize2fs /dev/${ROOT_PART} &&
-  logger -t "rc.firstboot" "Root partition successfuly resized."
+  logger -t "rc.firstboot" "Root partition successfully resized."
diff --git a/files/firstboot/23-regenerate-initramfs.sh b/files/firstboot/23-regenerate-initramfs.sh
new file mode 100644
index 0000000..78e2922
--- /dev/null
+++ b/files/firstboot/23-regenerate-initramfs.sh
@@ -0,0 +1,31 @@
+logger -t "rc.firstboot" "Regenerating initramfs to remove encrypted root partition auto-expand"
+
+KERNEL_VERSION=$(uname -r)
+KERNEL_ARCH=$(uname -m)
+INITRAMFS="/boot/firmware/initramfs-${KERNEL_VERSION}"
+INITRAMFS_UBOOT="${INITRAMFS}.uboot"
+
+# Extract kernel arch
+case "${KERNEL_ARCH}" in
+  arm*) KERNEL_ARCH=arm ;;
+esac
+
+# Regenerate initramfs
+if [ -r "${INITRAMFS}" ] ; then
+  rm -f /etc/initramfs-tools/scripts/init-premount/expand_encrypted_rootfs
+  rm -f /etc/initramfs-tools/scripts/local-premount/expand-premount
+  rm -f /etc/initramfs-tools/hooks/expand-tools
+  rm -f "${INITRAMFS}"
+  mkinitramfs -o "${INITRAMFS}" "${KERNEL_VERSION}"
+fi
+
+# Convert generated initramfs for U-Boot using mkimage
+if [ -r "${INITRAMFS_UBOOT}" ] ; then
+  rm -f /etc/initramfs-tools/scripts/init-premount/expand_encrypted_rootfs
+  rm -f /etc/initramfs-tools/scripts/local-premount/expand-premount
+  rm -f /etc/initramfs-tools/hooks/expand-tools
+  rm -f "${INITRAMFS_UBOOT}"
+  mkinitramfs -o "${INITRAMFS}" "${KERNEL_VERSION}"
+  mkimage -A "${KERNEL_ARCH}" -T ramdisk -C none -n "initramfs-${KERNEL_VERSION}" -d "${INITRAMFS}" "${INITRAMFS_UBOOT}"
+  rm -f "${INITRAMFS}"
+fi
diff --git a/files/firstboot/23-generate-machineid.sh b/files/firstboot/24-generate-machineid.sh
index dca48b3..dca48b3 100644
--- a/files/firstboot/23-generate-machineid.sh
+++ b/files/firstboot/24-generate-machineid.sh
diff --git a/files/firstboot/24-create-resolv-symlink.sh b/files/firstboot/25-create-resolv-symlink.sh
index 8a3f00b..8a3f00b 100644
--- a/files/firstboot/24-create-resolv-symlink.sh
+++ b/files/firstboot/25-create-resolv-symlink.sh