From 00132512b5776eec639af2f5989a06ca5bde4f91 Mon Sep 17 00:00:00 2001
From: ading2210 <ading2019@gmail.com>
Date: Fri, 26 Apr 2024 13:11:20 -0700
Subject: [PATCH] experimental arm support in build.sh

---
 build.sh      | 27 +++++++++++----------------
 shim_utils.sh | 39 ++++++++++++++++++++++++++++++++++++---
 2 files changed, 47 insertions(+), 19 deletions(-)

diff --git a/build.sh b/build.sh
index d74def6..053544f 100755
--- a/build.sh
+++ b/build.sh
@@ -8,10 +8,13 @@
 
 print_help() {
   echo "Usage: ./build.sh output_path shim_path rootfs_dir"
+  echo "Valid named arguments (specify with 'key=value'):"
+  echo "  quiet - Don't use progress indicators which may clog up log files."
+  echo "  arch  - Set this to 'arm64' to specify that the shim is for an ARM chromebook."
 }
 
 assert_root
-assert_deps "cpio binwalk pcregrep realpath cgpt mkfs.ext4 mkfs.ext2 fdisk rsync"
+assert_deps "cpio binwalk pcregrep realpath cgpt mkfs.ext4 mkfs.ext2 fdisk rsync lz4"
 assert_args "$3"
 parse_args "$@"
 
@@ -19,27 +22,18 @@ output_path=$(realpath -m "${1}")
 shim_path=$(realpath -m "${2}")
 rootfs_dir=$(realpath -m "${3}")
 
-echo "created loop device for shim"
-shim_loop=$(create_loop "${shim_path}")
-kernel_loop="${shim_loop}p2" #KERN-A should always be p2
-
-echo "copying shim kernel to new file in /tmp"
-kernel_dir=/tmp/shim_kernel
-mkdir $kernel_dir -p
-dd if=$kernel_loop of=$kernel_dir/kernel.bin bs=1M status=none
-
-echo "extracting data from kernel"
+echo "reading the shim image"
 initramfs_dir=/tmp/shim_initramfs
-rm -rf $initramfs_dir
-extract_initramfs $kernel_dir/kernel.bin $kernel_dir $initramfs_dir
-losetup -d $shim_loop
+kernel_img=/tmp/kernel.img
+rm -rf $initramfs_dir $kernel_img
+extract_initramfs_full $shim_path $initramfs_dir $kernel_img "${args['arch']}"
 
 echo "patching initramfs"
 patch_initramfs $initramfs_dir
 
 echo "creating disk image"
 rootfs_size=$(du -sm $rootfs_dir | cut -f 1)
-rootfs_part_size=$(($rootfs_size * 12 / 10))
+rootfs_part_size=$(($rootfs_size * 12 / 10 + 5))
 #create a 20mb bootloader partition
 #rootfs partition is 20% larger than its contents
 create_image $output_path 20 $rootfs_part_size
@@ -48,10 +42,11 @@ echo "creating loop device for the image"
 image_loop=$(create_loop ${output_path})
 
 echo "creating partitions on the disk image"
-create_partitions $image_loop "${kernel_dir}/kernel.bin"
+create_partitions $image_loop $kernel_img
 
 echo "copying data into the image"
 populate_partitions $image_loop $initramfs_dir $rootfs_dir "${args['quiet']}"
+rm -rf $initramfs_dir $kernel_img
 
 echo "cleaning up loop devices"
 losetup -d $image_loop
diff --git a/shim_utils.sh b/shim_utils.sh
index 5581ced..541e060 100755
--- a/shim_utils.sh
+++ b/shim_utils.sh
@@ -8,14 +8,14 @@ extract_initramfs() {
   local working_dir="$2"
   local output_dir="$3"
 
-  #first stage
+  #extract the compressed kernel image from the partition data
   local kernel_file="$(basename $kernel_bin)"
   local binwalk_out=$(binwalk --extract $kernel_bin --directory=$working_dir --run-as=root)
   local stage1_file=$(echo $binwalk_out | pcregrep -o1 "\d+\s+0x([0-9A-F]+)\s+gzip compressed data")
   local stage1_dir="$working_dir/_$kernel_file.extracted"
   local stage1_path="$stage1_dir/$stage1_file"
   
-  #second stage
+  #extract the initramfs cpio archive from the kernel image
   binwalk --extract $stage1_path --directory=$stage1_dir --run-as=root > /dev/null
   local stage2_dir="$stage1_dir/_$stage1_file.extracted/"
   local cpio_file=$(file $stage2_dir/* | pcregrep -o1 "([0-9A-F]+):\s+ASCII cpio archive")
@@ -25,6 +25,29 @@ extract_initramfs() {
   cat $cpio_path | cpio -D $output_dir -imd --quiet
 }
 
+extract_initramfs_arm() {
+  local kernel_bin="$1"
+  local working_dir="$2"
+  local output_dir="$3"
+
+  #extract the kernel lz4 archive from the partition
+  local binwalk_out="$(binwalk $kernel_bin --run-as=root)"
+  local lz4_offset="$(echo "$binwalk_out" | pcregrep -o1 "(\d+).+?LZ4 compressed data")"
+  local lz4_file="$working_dir/kernel.lz4"
+  local kernel_img="$working_dir/kernel_decompressed.bin"
+  dd if=$kernel_bin of=$lz4_file iflag=skip_bytes,count_bytes skip=$lz4_offset
+  lz4 -d $lz4_file $kernel_img -q || true
+
+  #extract the initramfs cpio archive from the kernel image
+  local extracted_dir="$working_dir/_kernel_decompressed.bin.extracted"
+  binwalk --extract $kernel_img --directory=$working_dir --run-as=root > /dev/null
+  local cpio_file=$(file $extracted_dir/* | pcregrep -o1 "([0-9A-F]+):\s+ASCII cpio archive")
+  local cpio_path="$extracted_dir/$cpio_file"
+
+  rm -rf $output_dir
+  cat $cpio_path | cpio -D $output_dir -imd --quiet
+}
+
 copy_kernel() {
   local shim_path="$1"
   local kernel_dir="$2"
@@ -40,6 +63,8 @@ copy_kernel() {
 extract_initramfs_full() {
   local shim_path="$1"
   local rootfs_dir="$2"
+  local kernel_bin="$3"
+  local arch="$4"
   local kernel_dir=/tmp/shim_kernel
 
   echo "copying the shim kernel"
@@ -48,6 +73,14 @@ extract_initramfs_full() {
   copy_kernel $shim_path $kernel_dir
 
   echo "extracting initramfs from kernel (this may take a while)"
-  extract_initramfs $kernel_dir/kernel.bin $kernel_dir $rootfs_dir
+  if [ "$arch" = "arm64" ]; then
+    extract_initramfs_arm $kernel_dir/kernel.bin $kernel_dir $rootfs_dir
+  else
+    extract_initramfs $kernel_dir/kernel.bin $kernel_dir $rootfs_dir
+  fi
+
+  if [ "$kernel_bin" ]; then 
+    cp $kernel_dir/kernel.bin $kernel_bin
+  fi
   rm -rf $kernel_dir
 }
\ No newline at end of file