#!/sbin/sh
##########################################################################################
#
# Magisk Boot Image Patcher
# by topjohnwu
# 
# This zip will patch your boot image with Magisk support
#
##########################################################################################

if [ -z "$BOOTMODE" ]; then
  BOOTMODE=false
fi

TMPDIR=/tmp
($BOOTMODE) && TMPDIR=/dev/tmp

INSTALLER=$TMPDIR/magisk

COREDIR=/magisk/.core

##########################################################################################
# Flashable update-binary preparation
##########################################################################################

OUTFD=$2
ZIP=$3

readlink /proc/$$/fd/$OUTFD 2>/dev/null | grep /tmp >/dev/null
if [ "$?" -eq "0" ]; then
  OUTFD=0

  for FD in `ls /proc/$$/fd`; do
    readlink /proc/$$/fd/$FD 2>/dev/null | grep pipe >/dev/null
    if [ "$?" -eq "0" ]; then
      ps | grep " 3 $FD " | grep -v grep >/dev/null
      if [ "$?" -eq "0" ]; then
        OUTFD=$FD
        break
      fi
    fi
  done
fi

mkdir -p $INSTALLER
cd $INSTALLER
unzip -o "$ZIP"

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

ui_print() {
  if ($BOOTMODE); then
    echo "$1"
  else 
    echo -n -e "ui_print $1\n" >> /proc/self/fd/$OUTFD
    echo -n -e "ui_print\n" >> /proc/self/fd/$OUTFD
  fi
}

getvar() {
  local VARNAME=$1
  local VALUE=$(eval echo \$"$VARNAME");
  for FILE in /data/.magisk /cache/.magisk /system/.magisk; do
    if [ -z "$VALUE" ]; then
      LINE=$(cat $FILE 2>/dev/null | grep "$VARNAME=")
      if [ ! -z "$LINE" ]; then
        VALUE=${LINE#*=}
      fi
    fi
  done
  eval $VARNAME=\$VALUE
}

find_boot_image() {
  if [ -z "$BOOTIMAGE" ]; then
    for PARTITION in kern-a KERN-A android_boot ANDROID_BOOT kernel KERNEL boot BOOT lnx LNX; do
      BOOTIMAGE=$(readlink /dev/block/by-name/$PARTITION || readlink /dev/block/platform/*/by-name/$PARTITION || readlink /dev/block/platform/*/*/by-name/$PARTITION)
      if [ ! -z "$BOOTIMAGE" ]; then break; fi
    done
  fi
  if [ -z "$BOOTIMAGE" ]; then
    FSTAB="/etc/recovery.fstab"
    [ ! -f "$FSTAB" ] && FSTAB="/etc/recovery.fstab.bak"
    BOOTIMAGE=$(grep -E '\b/boot\b' "$FSTAB" | grep -oE '/dev/[a-zA-Z0-9_./-]*')
  fi
}

is_mounted() {
  if [ ! -z "$2" ]; then
    cat /proc/mounts | grep $1 | grep $2, >/dev/null
  else
    cat /proc/mounts | grep $1 >/dev/null
  fi
  return $?
}

mount_image() {
  if [ ! -d "$2" ]; then
    mount -o rw,remount rootfs /
    mkdir -p $2 2>/dev/null
    ($BOOTMODE) && mount -o ro,remount rootfs /
    [ ! -d "$2" ] && return 1
  fi
  if (! is_mounted $2); then
    LOOPDEVICE=
    for LOOP in 0 1 2 3 4 5 6 7; do
      if (! is_mounted $2); then
        LOOPDEVICE=/dev/block/loop$LOOP
        if [ ! -f "$LOOPDEVICE" ]; then
          mknod $LOOPDEVICE b 7 $LOOP
        fi
        losetup $LOOPDEVICE $1
        if [ "$?" -eq "0" ]; then
          mount -t ext4 -o loop $LOOPDEVICE $2
          if (! is_mounted $2); then
            /system/bin/toolbox mount -t ext4 -o loop $LOOPDEVICE $2
          fi
          if (! is_mounted $2); then
            /system/bin/toybox mount -t ext4 -o loop $LOOPDEVICE $2
          fi
        fi
        if (is_mounted $2); then
          ui_print "- Mounting $1 to $2"
          break;
        fi
      fi
    done
  fi
}

grep_prop() {
  REGEX="s/^$1=//p"
  shift
  FILES=$@
  if [ -z "$FILES" ]; then
    FILES='/system/build.prop'
  fi
  cat $FILES 2>/dev/null | sed -n $REGEX | head -n 1
}

unpack_boot() {
  rm -rf $UNPACKDIR $RAMDISK 2>/dev/null
  mkdir -p $UNPACKDIR
  mkdir -p $RAMDISK
  cd $UNPACKDIR
  $BINDIR/bootimgtools --extract $1

  find $TMPDIR/boottmp -type d -exec chmod 755 {} \;
  find $TMPDIR/boottmp -type f -exec chmod 644 {} \;
  chmod 755 $(find $TMPDIR/boottmp -type d)
  chmod 644 $(find $TMPDIR/boottmp -type f)

  cd $RAMDISK
  gunzip -c < $UNPACKDIR/ramdisk.gz | cpio -i
}

repack_boot() {
  cd $RAMDISK
  find . | cpio -o -H newc 2>/dev/null | gzip -9 > $UNPACKDIR/ramdisk.gz
  cd $UNPACKDIR
  $BINDIR/bootimgtools --repack $ORIGBOOT
  if [ -f chromeos ]; then
    echo " " > config
    echo " " > bootloader
    $CHROMEDIR/futility vbutil_kernel --pack new-boot.img.signed --keyblock $CHROMEDIR/kernel.keyblock --signprivate $CHROMEDIR/kernel_data_key.vbprivk --version 1 --vmlinuz new-boot.img --config config --arch arm --bootloader bootloader --flags 0x1
    rm -f new-boot.img
    mv new-boot.img.signed new-boot.img
  fi
  if ($SAMSUNG); then
    SAMSUNG_CHECK=$(cat new-boot.img | grep SEANDROIDENFORCE)
    if [ $? -ne 0 ]; then
      echo -n "SEANDROIDENFORCE" >> new-boot.img
    fi
  fi
  mv new-boot.img $NEWBOOT
  $BINDIR/bootimgtools --hexpatch $NEWBOOT \
  49010054011440B93FA00F71E9000054010840B93FA00F7189000054001840B91FA00F7188010054 \
  A1020054011440B93FA00F7140020054010840B93FA00F71E0010054001840B91FA00F7181010054
}

##########################################################################################
# Detection
##########################################################################################

ui_print "****************************"
ui_print "Magisk v7 Boot Image Patcher"
ui_print "****************************"

if [ ! -d "$INSTALLER/common" ]; then
  ui_print "! Failed: Unable to extract zip file!"
  exit 1
fi

ui_print "- Mounting /system(ro), /cache, /data"
mount -o ro /system 2>/dev/null
mount /cache 2>/dev/null
mount /data 2>/dev/null

if [ ! -f '/system/build.prop' ]; then
  ui_print "! Failed: /system could not be mounted!"
  exit 1
fi

API=$(grep_prop ro.build.version.sdk)
ABI=$(grep_prop ro.product.cpu.abi | cut -c-3)
ABI2=$(grep_prop ro.product.cpu.abi2 | cut -c-3)
ABILONG=$(grep_prop ro.product.cpu.abi)

ARCH=arm
IS64BIT=
if [ "$ABI" = "x86" ]; then ARCH=x86; fi;
if [ "$ABI2" = "x86" ]; then ARCH=x86; fi;
if [ "$ABILONG" = "arm64-v8a" ]; then ARCH=arm64; IS64BIT=1; fi;
if [ "$ABILONG" = "x86_64" ]; then ARCH=x64; IS64BIT=1; fi;


if [ "$API" -lt "21" ]; then
  ui_print "! Magisk is only for Lollipop 5.0+ (SDK 21+)"
  exit 1
fi

ui_print "- Device platform: $ARCH"

BINDIR=$INSTALLER/arm
if [ "$ARCH" = "x86" -o "$ARCH" = "x64" ]; then
  BINDIR=$INSTALLER/x86
fi

find_boot_image
if [ -z "$BOOTIMAGE" ]; then
  ui_print "! Unable to detect boot image"
  exit 1
fi

if [ -z "$NOOVERRIDE" ]; then
  # read override variables
  getvar KEEPVERITY
  getvar KEEPFORCEENCRYPT
  getvar NORESTORE
fi

if [ -z "$KEEPVERITY" ]; then
  # we don't keep dm-verity by default
  KEEPVERITY=false
fi
if [ -z "$KEEPFORCEENCRYPT" ]; then
  # we don't keep forceencrypt by default
  KEEPFORCEENCRYPT=false
fi
if [ -z "$NORESTORE" ]; then
  # we don't keep ramdisk by default
  NORESTORE=false
fi

SAMSUNG=false
SAMSUNG_CHECK=$(cat /system/build.prop | grep "ro.build.fingerprint=" | grep -i "samsung")
if [ $? -eq 0 ]; then
  SAMSUNG=true
fi

##########################################################################################
# Environment
##########################################################################################

ui_print "- Constructing environment"

if (is_mounted /data); then
  rm -rf /data/busybox /data/magisk 2>/dev/null
  mkdir -p /data/busybox
  cp -af $BINDIR /data/magisk
  chmod 755 /data/busybox /data/magisk /data/magisk/*
  chcon 'u:object_r:system_file:s0' /data/busybox /data/magisk /data/magisk/*
  /data/magisk/busybox --install -s /data/busybox
  # Prevent issues
  rm -f /data/busybox/su /data/busybox/sh
else
  rm -rf /cache/data_bin 2>/dev/null
  mkdir -p /cache/data_bin
  cp -af $BINDIR /cache/data_bin
fi

##########################################################################################
# Image
##########################################################################################

# Fix SuperSU.....
($BOOTMODE) && /data/magisk/sepolicy-inject -s fsck --live

if (is_mounted /data); then
  IMG=/data/magisk.img
else
  IMG=/cache/magisk.img
  ui_print "- Data unavalible, use cache workaround"
fi

if [ -f "$IMG" ]; then
  ui_print "- $IMG detected!"
else
  ui_print "- Creating $IMG"
  make_ext4fs -l 64M -a /magisk -S $INSTALLER/common/file_contexts_image $IMG
fi

mount_image $IMG /magisk
if (! is_mounted /magisk); then
  ui_print "! Image mount failed... abort"
  exit 1
fi
MAGISKLOOP=$LOOPDEVICE

##########################################################################################
# Boot image patch
##########################################################################################

ui_print "- Found Boot Image: $BOOTIMAGE"

rm -rf $TMPDIR/boottmp 2>/dev/null
mkdir -p $TMPDIR/boottmp

CHROMEDIR=$INSTALLER/chromeos
ORIGBOOT=$TMPDIR/boottmp/boot.img
NEWBOOT=$TMPDIR/boottmp/new-boot.img
UNPACKDIR=$TMPDIR/boottmp/bootunpack
RAMDISK=$TMPDIR/boottmp/ramdisk

chmod 777 $CHROMEDIR/futility $BINDIR/*

ui_print "- Dumping boot image"
dd if=$BOOTIMAGE of=$ORIGBOOT

ui_print "- Unpacking boot image"
unpack_boot $ORIGBOOT

SUPERSU=false

if (! $NORESTORE); then
  #  Backups
  if [ -d ".backup" ]; then
    ui_print "- Restoring ramdisk with backup"
    cp -af .backup/* .
    rm -rf magisk init.magisk.rc sbin/magic_mask.sh sbin/magisk_wrapper.sh 2>/dev/null
  else
    if [ -f "sbin/launch_daemonsu.sh" ]; then
      SUPERSU=true
      # Save it for helper module
      cp -af sbin/launch_daemonsu.sh $INSTALLER/roothelper/launch_daemonsu.sh
    fi
    # Non-standard boot image restores
    if ($SUPERSU); then
      ui_print "- SuperSU patched boot detected"
      # Restore with SuperSU's backup
      MOUNTSU=false
      (! is_mounted /su) && (is_mounted /data) && mount_image /data/su.img /su && MOUNTSU=true
      if (is_mounted /su); then
        # Use sukernel's built-in functions
        ui_print "- Using sukernel to restore boot image"
        cd $UNPACKDIR
        gunzip -c < $UNPACKDIR/ramdisk.gz > suramdisk
        /su/bin/sukernel --restore suramdisk /data/stock_boot.img
        if [ "$?" -ne "0" ]; then
          # No boot backup found, use ramdisk backup
          ui_print "- Restoring ramdisk with backup"
          /su/bin/sukernel --cpio-restore suramdisk suramdisk
          rm -rf $RAMDISK
          mkdir -p $RAMDISK
          cd $RAMDISK
          cpio -i < $UNPACKDIR/suramdisk
          rm -f $UNPACKDIR/suramdisk
        else
          ui_print "- Restoring boot image with backup"
          cp -af /data/stock_boot.img $ORIGBOOT
          unpack_boot $ORIGBOOT
        fi
        if ($MOUNTSU); then
          ui_print "- Unmounting su.img"
          umount /su
          losetup -d $LOOPDEVICE
        fi
      else
        # Find the boot backup ourselves
        ui_print "! su.img mount failed... find the backup ourselves"
        cp -af /data/stock_boot_*.gz /data/stock_boot.img.gz 2>/dev/null
        gzip -d /data/stock_boot.img.gz 2>/dev/null
        rm -rf /data/stock_boot.img.gz 2>/dev/null
        if [ -f "/data/stock_boot.img" ]; then
          ui_print "- Restoring boot image with backup"
          cp -af /data/stock_boot.img $ORIGBOOT
          unpack_boot $ORIGBOOT
        else
          ui_print "! No backups found"
          ui_print "! Installer will still proceed, but might cause issues"
          ui_print "! If possible, please restore to stock boot then flash Magisk again"
          # Force removing SuperSU parts
          rm -rf su init.supersu.rc sbin/launch_daemonsu.sh 2>/dev/null
        fi
      fi
    else
      # Magisk's own situation
      if [ -d "magisk" ]; then
        cp -af /data/stock_boot_*.gz /data/stock_boot.img.gz 2>/dev/null
        gzip -d /data/stock_boot.img.gz 2>/dev/null
        if [ -f "/data/stock_boot.img" ]; then
          ui_print "- Restoring boot image with backup"
          cp -af /data/stock_boot.img $ORIGBOOT
          unpack_boot $ORIGBOOT
        else
          ui_print "! No backups found"
          ui_print "! Installer will still proceed, but might cause issues"
          ui_print "! If possible, please restore to stock boot then flash Magisk again"
          # Removing other boot image modifications
          rm -rf sbin/su init.xposed.rc sbin/mount_xposed.sh 2>/dev/null
        fi
      fi
    fi
    ui_print "- Creating backups"
    mkdir .backup
    cp -af init.rc *fstab* verity_key sepolicy .backup 2>/dev/null
    if (is_mounted /data); then
      cp -af $ORIGBOOT /data/stock_boot.img
    else
      cp -af $ORIGBOOT /cache/stock_boot.img
    fi
  fi
fi

ui_print "- Installing root helper module"
cp -af $INSTALLER/roothelper /magisk/00roothelper
chmod 755 /magisk/00roothelper /magisk/00roothelper/*

# Patch ramdisk
ui_print "- Patching ramdisk"

if [ $(grep -c "import /init.magisk.rc" init.rc) -eq "0" ]; then
  sed -i "/import \/init\.environ\.rc/iimport /init.magisk.rc" init.rc
fi

sed -i "/selinux.reload_policy/d" init.rc
find . -type f -name "*fstab*" 2>/dev/null | while read FSTAB ; do
  if (! $KEEPVERITY); then
    sed -i "s/,support_scfs//g" $FSTAB
    sed -i 's;,\{0,1\}verify\(=[^,]*\)\{0,1\};;g' $FSTAB
  fi
  if (! $KEEPFORCEENCRYPT); then
    sed -i "s/forceencrypt/encryptable/g" $FSTAB
    sed -i "s/forcefdeorfbe/encryptable/g" $FSTAB
  fi
done
rm verity_key 2>/dev/null

# sepolicy patches
$BINDIR/sepolicy-inject --magisk -P sepolicy

# Add new items
mkdir -p magisk 2>/dev/null
cp -af $INSTALLER/common/init.magisk.rc init.magisk.rc
cp -af $INSTALLER/common/magic_mask.sh sbin/magic_mask.sh

chmod 0755 magisk
chmod 0750 init.magisk.rc sbin/magic_mask.sh

ui_print "- Repacking boot image"
repack_boot

ORIGSIZE=$(ls -l $ORIGBOOT | awk '{print $5}')
NEWSIZE=$(ls -l $NEWBOOT | awk '{print $5}')
if [ "$NEWSIZE" -gt "$ORIGSIZE" ]; then
  ui_print "! Boot partition space insufficient"
  ui_print "! Try to remove ramdisk backups"
  rm -rf $RAMDISK/.backup $NEWBOOT 2>/dev/null
  repack_boot
  NEWSIZE=$(ls -l $NEWBOOT | awk '{print $5}')
  if [ "$NEWSIZE" -gt "$ORIGSIZE" ]; then
    ui_print "! Boot partition size still too small..."
    ui_print "! Unable to install Magisk"
    exit 1
  fi
fi

chmod 644 $NEWBOOT

if [ -L "$BOOTIMAGE" ]; then
  ui_print "- Block symlink detected!"
else
  dd if=/dev/zero of=$BOOTIMAGE bs=4096 2>/dev/null
fi
ui_print "- Flashing new boot image"
dd if=$NEWBOOT of=$BOOTIMAGE bs=4096

if (! $BOOTMODE); then
  ui_print "- Unmounting partitions"
  umount /magisk
  losetup -d $MAGISKLOOP
  umount /system
fi

ui_print "- Done"
exit 0