#!/system/bin/sh
#######################################################################################
# Magisk Boot Image Patcher
#######################################################################################
#
# Usage: boot_patch.sh <bootimage>
#
# The following flags can be set in environment variables:
# KEEPVERITY, KEEPFORCEENCRYPT, RECOVERYMODE
#
# This script should be placed in a directory with the following files:
#
# File name          Type      Description
#
# boot_patch.sh      script    A script to patch boot image for Magisk.
#                  (this file) The script will use files in its same
#                              directory to complete the patching process
# util_functions.sh  script    A script which hosts all functions required
#                              for this script to work properly
# magiskinit         binary    The binary to replace /init
# magisk(32/64)      binary    The magisk binaries
# magiskboot         binary    A tool to manipulate boot images
# chromeos           folder    This folder includes the utility and keys to sign
#                  (optional)  chromeos boot images. Only used for Pixel C.
#
#######################################################################################

############
# Functions
############

# Pure bash dirname implementation
getdir() {
  case "$1" in
    */*)
      dir=${1%/*}
      if [ -z $dir ]; then
        echo "/"
      else
        echo $dir
      fi
    ;;
    *) echo "." ;;
  esac
}

#################
# Initialization
#################

[ -e "$BOOTIMAGE" ] || abort "$BOOTIMAGE does not exist!"

# Dump image for MTD/NAND character device boot partitions
if [ -c "$BOOTIMAGE" ]; then
  nanddump -f boot.img "$BOOTIMAGE"
  BOOTNAND="$BOOTIMAGE"
  BOOTIMAGE=boot.img
fi

# Flags
[ -z $KEEPVERITY ] && KEEPVERITY=false
[ -z $KEEPFORCEENCRYPT ] && KEEPFORCEENCRYPT=false
[ -z $RECOVERYMODE ] && RECOVERYMODE=false
export KEEPVERITY
export KEEPFORCEENCRYPT

chmod -R 755 .

#########
# Unpack
#########

CHROMEOS=false

ui_print "- Unpacking boot image"
$MAGISKBIN/magiskboot unpack "$BOOTIMAGE"

case $? in
  0 ) ;;
  1 )
    abort "! Unsupported/Unknown image format"
    ;;
  2 )
    ui_print "- ChromeOS boot image detected"
    CHROMEOS=true
    ;;
  * )
    abort "! Unable to unpack boot image"
    ;;
esac

[ -f recovery_dtbo ] && RECOVERYMODE=true

###################
# Ramdisk Restores
###################

# Test patch status and do restore
ui_print "- Checking ramdisk status"
if [ -e ramdisk.cpio ]; then
  $MAGISKBIN/magiskboot cpio ramdisk.cpio test
  STATUS=$?
else
  # Stock A only system-as-root
  STATUS=0
fi
case $((STATUS & 3)) in
  0 )  # Stock boot
    ui_print "- Stock boot image detected"
    abort "! magisk must be installed"
    ;;
  1 )  # Magisk patched
    ui_print "- Magisk patched boot image detected"
    # Find SHA1 of stock boot image
    [ -z $SHA1 ] && SHA1=$($MAGISKBIN/magiskboot cpio ramdisk.cpio sha1 2>/dev/null)
    cp -af ramdisk.cpio ramdisk.cpio.stock.magisk
    ;;
  2 )  # Unsupported
    ui_print "! Boot image patched by unsupported programs"
    abort "! Please restore back to stock boot image"
    ;;
esac

##################
# Ramdisk Patches
##################

ui_print "- Patching ramdisk"

echo "KEEPVERITY=$KEEPVERITY" > config
echo "KEEPFORCEENCRYPT=$KEEPFORCEENCRYPT" >> config
echo "RECOVERYMODE=$RECOVERYMODE" >> config
[ ! -z $SHA1 ] && echo "SHA1=$SHA1" >> config

$MAGISKBIN/magiskboot cpio ramdisk.cpio "extract overlay.d/ghid.rc ghid.rc" 2>/dev/null
if ! diff $MODPATH/overlay.d/ghid.rc ghid.rc 2>/dev/null; then
  ui_print "- Adding overlay.d/ghid.rc"

  $MAGISKBIN/magiskboot cpio ramdisk.cpio \
  "mkdir 755 overlay.d" \
  "add 755 overlay.d/ghid.rc $MODPATH/overlay.d/ghid.rc"

  rm -f ramdisk.cpio.orig config magisk*.xz

#################
# Binary Patches
#################

  if [ $((STATUS & 4)) -ne 0 ]; then
    ui_print "- Compressing ramdisk"
    $MAGISKBIN/magiskboot cpio ramdisk.cpio compress
  fi

  rm -f ramdisk.cpio.stock.magisk config ghid.rc
#################
# Repack & Flash
#################

  ui_print "- Repacking boot image"
  $MAGISKBIN/magiskboot repack "$BOOTIMAGE" || abort "! Unable to repack boot image!"
# Sign chromeos boot
$CHROMEOS && sign_chromeos

  # Sign chromeos boot
  $CHROMEOS && sign_chromeos
else
  ui_print "- ghid.rc already installed. Skipping boot patching."
  rm -f ramdisk.cpio.stock.magisk config ghid.rc
fi
# Restore the original boot partition path
[ -e "$BOOTNAND" ] && BOOTIMAGE="$BOOTNAND"

# Reset any error code
true