#!/bin/bash

#utilties for reading shim disk images

#extract the initramfs from a kernel image
extract_initramfs() {
  local kernel_bin="$1"
  local working_dir="$2"
  local output_dir="$3"

  #first stage
  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
  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")
  local cpio_path="$stage2_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"

  local shim_loop=$(create_loop "${shim_path}")
  local kernel_loop="${shim_loop}p2" #KERN-A should always be p2

  dd if=$kernel_loop of=$kernel_dir/kernel.bin bs=1M status=progress
  losetup -d $shim_loop
}

#copy the kernel image then extract the initramfs
extract_initramfs_full() {
  local shim_path="$1"
  local rootfs_dir="$2"
  local kernel_dir=/tmp/shim_kernel

  echo "copying the shim kernel"
  rm -rf $kernel_dir
  mkdir $kernel_dir -p
  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
  rm -rf $kernel_dir
}