#!/system/xbin/bash
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
#boot.img tool
#original author: xiaolu
#heavily modified by: Modding.MyMind

#set -x # for debugging

trap "clean" 2 3 4
workdir=$(pwd)
toolpath=$(readlink -f $0)
tooldir=$(dirname $toolpath)
mkbootimg=$tooldir/mkbootimg
mkbootfs=$tooldir/mkbootfs
busybox=$(command -v busybox)
od=$tooldir/od
gzip=$tooldir/gzip
lz4=$tooldir/lz4
lzma=$tooldir/lzma
xz=$tooldir/xz # Also used for lzma compression
grep=$tooldir/grep
cpio=$tooldir/cpio
magic=$tooldir/magic
file=$tooldir/file
old_bootimg=true
C_OUT="\033[0;1m"
C_ERR="\033[31;1m"
C_CLEAR="\033[0;0m"
	
pout() {
    $busybox printf "${C_OUT}${*}${C_CLEAR}\n"
}
perr() {
    $busybox printf "${C_ERR}${*}${C_CLEAR}\n"
}
clean()
{
    $busybox rm -rf /tmp/mkboot.*
    #pout "..."
    exit
}

# Check for busybox
if [ -z $busybox ]; then
	clear; pout "Busybox is NOT installed!\nThis may cause issues with the script!\nUse at your own risk or install busybox!"; sleep 3
fi

usage()
{
    pout "<Unpack and repack boot.img tool>"
    pout "----------------------------------------------------------------------"
    pout "Not enough parameters or parameter error!"
    pout "unpack boot.img & decompress ramdisk：\n    $($busybox basename $0) [img] [output dir]"
    pout "    $($busybox basename $0) boot.img boot20130905"
    pout "Use the unpacked directory repack boot.img(img_info)：\n    $($busybox basename $0) [unpacked dir] [newbootfile]"
    pout "    $($busybox basename $0) boot20130905 newboot.img"
    clean
}

print_info()
{
    pout "  kernel         : $kernel"
    pout "  ramdisk        : $ramdisk"
    pout "  page size      : $page_size"
    pout "  kernel size    : $kernel_size"
    pout "  ramdisk size   : $ramdisk_size"
    [ ! -z $second_size ] && [ $second_size -gt 0 ] && \
        pout "  second_size    : $second_size"
    [ $dtb_size -gt 0 ] && pout "  dtb size       : $dtb_size"
    pout "  base           : $base_addr"
    pout "  kernel offset  : $kernel_offset"
    pout "  ramdisk offset : $ramdisk_offset"
    [ -z $second_offset ] || pout "  second_offset  : $second_offset"
    pout "  tags offset    : $tags_offset"
    [ $dtb_size -gt 0 ] && pout "  dtb img        : $dt"
    pout "  cmd line       : $cmd_line"
}

mkboot_img()
{
    error=0
    if [ $dtb_size -gt 0 ]; then
        dtb="--dt ${dt}"
    fi
    $mkbootimg --kernel $kernel --ramdisk $ramdisk \
        --base $base_addr --ramdisk_offset $ramdisk_offset \
        --tags_offset $tags_offset --cmdline "$cmd_line" \
        --pagesize $page_size $dtb -o $new_img 2>/dev/null || error=1
    [ $error -eq 1 ] && return $error
    ramdisk_size=$($busybox stat -c "%s" $ramdisk)
    boot_size=$($busybox stat -c "%s" $new_img)
    pout "\nKernel size: $kernel_size, new ramdisk size: $ramdisk_size, $($busybox basename $new_img): $boot_size."
    pout "\n$($busybox basename $new_img) has been created.\n"
}

#decide action
[ $# -lt 2 ] || [ $# -gt 3 ] && usage
if [ $# -eq 2 ] && [ -d $1 ]; then
    mkboot_from_dir=1
elif [ $# -eq 2 ] && [ -s $1 ]; then
    split_boot_to_dir=1
else
    usage
fi

#mkboot_from_dir, img_info
if [ ! -z $mkboot_from_dir ]; then
    pout "\nmkbootimg from $1/img_info.\n"
    unpacked_dir=$1
    new_img=$2
    cd $unpacked_dir
    if [ ! -s img_info ]; then
        pout "Missing img_info file! Can't rebuild $2."
        clean
    fi 
    eval $(cat img_info)
    if [ -z $kernel ] || [ -z $ramdisk ] || [ -z $base_addr ]; then
        pout "Lacking parameters in img_info."
        clean
    fi
    [ -z $dtb_size ] && dtb_size=0
    if [ -d $ramdisk ]; then
	compression_type=$($file -m $magic ./ramdisk.* | $busybox cut -d: -f2 | $busybox cut -d" " -f2)
	
	compression_ext=$compression_type
	case $compression_type in
		gzip) compression_warning=$compression_type; compression_repack=$gzip;;
		xz) compression_warning=$compression_type; compression_repack="$xz -1 --check=crc32";;
		lzma) compression_warning=$compression_type; compression_repack="$xz --format=lzma";;
		lz4) compression_warning=$compression_type; compression_repack="$lz4 -l";;
	esac;
	if [ -z $compression_warning ]; then
		pout "\n****** HAZARD ******* HAZARD ******* HAZARD ******"
		pout "\nRamdisk is $compression_type format. Can't repack ramdisk."
		pout "This tool currently does not support $compression_type."
		pout "\n****** HAZARD ******* HAZARD ******* HAZARD ******\n"
		exit
	fi	
	if [ $compression_type != "gzip" ]; then
		if [ $compression_type != "lzma" ]; then
			if [ $compression_type != "lz4" ]; then
				if [ $compression_type != "xz" ]; then
					pout "\nRamdisk is unknown format. Can't repack ramdisk."
					exit 0
				else
					# XZ
					$mkbootfs $ramdisk | $compression_repack > new_ramdisk.$compression_ext
					ramdisk=new_ramdisk.$compression_ext
					ramdisk_size=$($busybox stat -c "%s" $ramdisk)
				fi
			else
				# LZ4
				$mkbootfs $ramdisk | $compression_repack > new_ramdisk.$compression_ext
				ramdisk=new_ramdisk.$compression_ext
				ramdisk_size=$($busybox stat -c "%s" $ramdisk)
			fi
		else
			# LZMA using xz binary
			$mkbootfs $ramdisk | $compression_repack > new_ramdisk.$compression_ext
			ramdisk=new_ramdisk.$compression_ext
			ramdisk_size=$($busybox stat -c "%s" $ramdisk)
		fi
	else
		# GZIP
		$mkbootfs $ramdisk | $compression_repack > new_ramdisk.$compression_ext
		ramdisk=new_ramdisk.$compression_ext
		ramdisk_size=$($busybox stat -c "%s" $ramdisk)
	fi
    fi
    #cd $unpacked_dir
    print_info
    $busybox rm -f $new_img
    mkboot_img $new_img || perr "Make $new_img Error! pls check img_info file."
    #pout "Add SEANDROIDENFORCE tag."
    #printf SEANDROIDENFORCE >> $new_img
    $busybox rm -f new_ramdisk.gz
    clean
fi

#split boot.img to dir.
if [ -e $2 ]; then
    read -p "$2 exists, delete?(N/y)" reply
    case $reply in
        y | Y)
        $busybox rm -rf $2
        ;;
    *)
        exit
        ;;
    esac
fi
tempdir=$2
$busybox  mkdir -p "$tempdir"
pout "\nUnpack & decompress $1 to $2\n"

#get boot.img info
cp -f $1 $tempdir/
cd $tempdir
bootimg=$($busybox basename $1)
offset=$($grep -abo ANDROID! $bootimg | $busybox cut -f 1 -d :)
[ -z $offset ] && clean
if [ $offset -gt 0 ]; then
    $busybox dd if=$bootimg of=bootimg bs=$offset skip=1 2>/dev/null
    bootimg=bootimg
fi

kernel_addr=0x$($od -A n -X -j 12 -N 4 $bootimg | $busybox sed 's/ //g' | $busybox sed 's/^0*//g')
ramdisk_addr=0x$($od -A n -X -j 20 -N 4 $bootimg | $busybox sed 's/ //g' | $busybox sed 's/^0*//g')
second_addr=0x$($od -A n -X -j 28 -N 4 $bootimg | $busybox sed 's/ //g' | $busybox sed 's/^0*//g')
tags_addr=0x$($od -A n -X -j 32 -N 4 $bootimg | $busybox sed 's/ //g' | $busybox sed 's/^0*//g')

kernel_size=$($od -A n -D -j 8 -N 4 $bootimg | $busybox sed 's/ //g')
base_addr=0x$($od -A n -x -j 14 -N 2 $bootimg | $busybox sed 's/ //g')0000
ramdisk_size=$($od -A n -D -j 16 -N 4 $bootimg | $busybox sed 's/ //g')
second_size=$($od -A n -D -j 24 -N 4 $bootimg | $busybox sed 's/ //g')
page_size=$($od -A n -D -j 36 -N 4 $bootimg | $busybox sed 's/ //g')
dtb_size=$($od -A n -D -j 40 -N 4 $bootimg | $busybox sed 's/ //g')
cmd_line=$($od -A n --strings -j 64 -N 512 $bootimg)

kernel_offset=$((kernel_addr-base_addr))
ramdisk_offset=$((ramdisk_addr-base_addr))
second_offset=$((second_addr-base_addr))
tags_offset=$((tags_addr-base_addr))

#########################################################
# BELOW SECTION HANDLES NON STANDARD IMAGES

non_standard_base=$((tags_addr-0x00000100)) # Need this to get the TRUE Base Address
non_standard_base_addr=$($busybox printf "0x%08x" $non_standard_base)
base_addr_verify=$($busybox printf "0x%08x" $base_addr)
base_addr=$($busybox printf "0x%08x" $base_addr)

if [ $base_addr == $non_standard_base ]; then
	base_addr=$($busybox printf "0x%08x" $base_addr)
else
	base_addr=$($busybox printf "0x%08x" $non_standard_base_addr)
fi

# Accurate offsets used on images with non standard mkbootimg.c
kernel_offset_ns=$((kernel_addr-non_standard_base))
ramdisk_offset_ns=$((ramdisk_addr-non_standard_base))
second_offset_ns=$((second_addr-non_standard_base))
tags_offset_ns=$((tags_addr-non_standard_base))

# Do the Offsets match with the Base Address for standard mkbootimg.c?
kernel_offset_s=$((kernel_addr-0x00008000))
ramdisk_offset_s=$((ramdisk_addr-0x01000000))
second_offset_s=$((second_addr-0x00f00000))
tags_offset_s=$((tags_addr-0x00000100))

if [ $non_standard_base -eq $kernel_offset_s ]; then
	kernel_offset=$($busybox printf "0x%08x" $kernel_offset)
else
	kernel_offset_ns=$($busybox printf "0x%08x" $kernel_offset_ns)
	kernel_offset=$($busybox printf "0x%08x" $kernel_offset_ns)
	kernel_offset_warning=$kernel_offset_ns
fi

if [ $non_standard_base -eq $ramdisk_offset_s ]; then
	ramdisk_offset=$($busybox printf "0x%08x" $ramdisk_offset)
else
	ramdisk_offset_ns=$($busybox printf "0x%08x" $ramdisk_offset_ns)
	ramdisk_offset=$($busybox printf "0x%08x" $ramdisk_offset_ns)
	ramdisk_offset_warning=$ramdisk_offset_ns
fi

if [ $non_standard_base -eq $second_offset_s ]; then
	second_offset=$($busybox printf "0x%08x" $second_offset)
else
	second_offset_ns=$($busybox printf "0x%08x" $second_offset_ns)
	second_offset=$($busybox printf "0x%08x" $second_offset_ns)
	second_offset_warning=$second_offset_ns
fi

if [ $base_addr_verify != $non_standard_base_addr ]; then
	base_warning=$non_standard_base_addr
fi

tags_offset=$($busybox printf "0x%08x" $tags_offset_ns)

# Below are the known offsets for non standard mkbootimg.c
if [[ ! -z $kernel_offset_warning ]] || [[ ! -z $ramdisk_offset_warning ]] || [[ ! -z $second_offset_warning ]]; then
	pout "****** WARNING ******* WARNING ******* WARNING ******\n"
	pout "This image is built using NON-standard mkbootimg!\n"
fi
if [ ! -z $base_warning ]; then
	pout "BASE is $base_warning"
fi
if [ ! -z $kernel_offset_warning ]; then
	pout "KERNEL_OFFSET is $kernel_offset_ns"
fi
if [ ! -z $ramdisk_offset_warning ]; then
	pout "RAMDISK_OFFSET is $ramdisk_offset_ns"
fi
if [ ! -z $second_offset_warning ]; then
	pout "SECOND_OFFSET is $second_offset_ns"
fi
if [[ ! -z $kernel_offset_warning ]] || [[ ! -z $ramdisk_offset_warning ]] || [[ ! -z $second_offset_warning ]]; then
	pout "\nYou can modify mkbootimg.c with the above value(s)"
	pout "\n****** WARNING ******* WARNING ******* WARNING ******\n"
fi

# ABOVE SECTION HANDLES NON STANDARD IMAGES
#########################################################

k_count=$(((kernel_size+page_size-1)/page_size))
r_count=$(((ramdisk_size+page_size-1)/page_size))
s_count=$(((second_size+page_size-1)/page_size))
d_count=$(((dtb_size+page_size-1)/page_size))
k_offset=1
r_offset=$((k_offset+k_count))
s_offset=$((r_offset+r_count))
d_offset=$((s_offset+s_count))

#zImage
$busybox dd if=$bootimg of=zImage_tmp bs=$page_size skip=$k_offset count=$k_count 2>/dev/null
$busybox dd if=zImage_tmp of=zImage bs=$kernel_size count=1 2>/dev/null
#ramdisk.gz
$busybox dd if=$bootimg of=ramdisk_tmp bs=$page_size skip=$r_offset count=$r_count 2>/dev/null
$busybox dd if=ramdisk_tmp of=ramdisk.gz bs=$ramdisk_size count=1 2>/dev/null
#dtb
if [ $dtb_size -gt 0 ]; then
    $busybox dd if=$bootimg of=dt.img_tmp bs=$page_size skip=$d_offset count=$d_count 2>/dev/null
    $busybox dd if=dt.img_tmp of=dt.img bs=$dtb_size count=1 2>/dev/null
    dt="$tempdir/dt.img"
    dt=$($busybox basename $dt)
    dt_name="dt=$dt\n"
    dt_size="dtb_size=$dtb_size\n"
fi
$busybox rm -f *_tmp $($busybox basename $1) $bootimg

kernel=zImage
ramdisk=ramdisk
[ ! -s $kernel ] && clean
#print boot.img info
print_info

#write info to img_info,decompression ramdisk.gz
$busybox printf "kernel=zImage\nramdisk=ramdisk\n${dt_name}page_size=$page_size\n\
kernel_size=$kernel_size\nramdisk_size=$ramdisk_size\n${dt_size}base_addr=$base_addr\nkernel_offset=$kernel_offset\n\
ramdisk_offset=$ramdisk_offset\ntags_offset=$tags_offset\ncmd_line=\"$cmd_line\"\n" > img_info
$busybox mkdir ramdisk
cd ramdisk

compression_type=$($file -m $magic ../ramdisk.gz | $busybox cut -d: -f2 | $busybox cut -d" " -f2)
compression_warning=$compression_type

case $compression_type in
	gzip) compression_type=$gzip; compression_ext=gz;;
	xz) compression_type=$xz; compression_ext=xz;;
	lzma) compression_type=$lzma; compression_ext=lzma;;
	lz4) compression_ext=lz4; decomp_ramdisk="$lz4 -d"; extra="< ../ramdisk.cpio";;
esac;

if [[ $compression_ext != "lz4" ]]; then
	decomp_ramdisk="$compression_type -d -c"
fi
decomp_ramdisk2="$cpio -i -d -m --no-absolute-filenames"

if [ -z $compression_ext ]; then
	pout "\n****** HAZARD ******* HAZARD ******* HAZARD ******"
	pout "\nRamdisk is $compression_warning format. Can't unpack ramdisk."
	pout "This tool currently does not support $compression_warning."
	pout "\n****** HAZARD ******* HAZARD ******* HAZARD ******\n"
	exit
fi

$busybox mv ../ramdisk.gz ../ramdisk.$compression_ext # This is simply to remind the user if they view the folder afterwards
pout "\nRamdisk is $compression_warning format."
$decomp_ramdisk "../ramdisk.$compression_ext" | $decomp_ramdisk2 $extra

#Unpack Finish to exit.
pout "Unpack completed.\n"
exit