#!/bin/bash

# The purpose of the script is to reinstall the operating system (debrick) on 
# a harddrive that has been extracted from the housing of a WD MyBook Live.

# Edited by Fox_exe

# Color codes
NORM="\033[0m"
INFO="\033[0;32mInfo:$NORM"
ERROR="\033[0;31mError:$NORM"
WARNING="\033[1;33mWarning:$NORM"
INPUT="\033[1;32m => $NORM"

# Vars
LogFile=recovery.log
RootfsDir=/mnt/rootfs
RootfsDev=/dev/md0
HDD=/dev/sda
if [ -n $1 ]; then HDD=$1; fi
ImgFile=hdd.img
if [ -n $2 ]; then ImgFile=$2; fi
Format="False"
if [ $3 = "format" ]; then Format="True"; fi

HDDRoot1=${HDD}1
HDDRoot2=${HDD}2
HDDSwap=${HDD}3
HDDData=${HDD}4

# Usage info
if  [ $# = 1 -a "$1" = "--help" ]; then
	echo -e "$INFO: Usage: sudo ./debricker.sh /dev/sda [hdd.img|latest] [format|none]"
	exit 1
fi

# Check that requirements are fullfilled
if [ "$(id -u)" != "0" ]; then
    echo -e "$ERROR This script must be run as root."
    exit 1
fi

if [ ! -b $HDD ]; then
	echo -e "$ERROR $HDD not exist or path not correct."
	exit 1
fi

check_install () {
	if ! which $1 > /dev/null; then
		echo -e "$WARNING $1 not found. Trying to install..."
		apt-get update >> $LogFile 2>&1
		apt-get install -f $1 >> $LogFile 2>&1
		if ! which $1 > /dev/null; then
			echo -e "$ERROR Failed to install $1. Please, install it manually."
			exit 1
		fi
	fi
}

check_install mdadm
check_install parted
check_install build-essential
# Curl checked only if no firmware specified.

# Making sure the mountpoint is available
if [ -e ${RootfsDir} ]; then
    if mountpoint -q ${RootfsDir}; then
        echo -e "$WARNING ${RootfsDir} Already mounted."
        umount $RootfsDir >> $LogFile 2>&1
        if mountpoint -q ${RootfsDir}; then
			echo -e "$ERROR Cant unmount $RootfsDir, please, unmount it manually."
			exit 1;
		fi
    fi
fi
mkdir -p $RootfsDir

# Making sure that there is no raid unit running
if [ -e $RootfsDev ]; then
    echo -e "$WARNING $RootfsDev already exists! Ill try to stop it..."
    mdadm --stop $RootfsDev >> $LogFile 2>&1
    mdadm --wait $RootfsDev >> $LogFile 2>&1
	rm -r $RootfsDev > $LogFile 2>&1
	sync
	if [ -e $RootfsDev ]; then
		echo -e "$ERROR Something wrong, $RootfsDev still exist. Please, stop it manually."
		exit 1
	fi
fi

echo -e "$INFO $HDD properties:"
parted --script $HDD "print"
echo -en "$INPUT This is correct disk? [y/n]"
read userAnswer
if [ "$userAnswer" != "y" ] || [ "$userAnswer" != "Y" ] || [ "$userAnswer" != "yes" ]; then
    echo -e "$ERROR Wrong disk. Exiting..."
    exit 1
fi

if [ ! -e $ImgFile ]; then
	echo -e "$WARNING $ImgFile not exist or path not correct, so i try to download latest firmware..."
    check_install curl
    # Get latest firmware from WD website:
    wdc_homepage="http://websupport.wdc.com/firmware/list.asp"
    wdc_latestfirmware=$(curl "${wdc_homepage}?type=AP1NC&fw=01.03.03" 2> /dev/null | awk ' {
        if ( match($0, "upgrade file" ) != 0 ) {
            split($0, http, "\"");
            print http[2];
            exit 1;
        }
    }')
    latestversion_simple=$(echo $wdc_latestfirmware | cut -d'-' -f 2)
    latestversion_pattern=$(echo $latestversion_simple | sed 's/../&*/g;s/:$//')
    if [ "$latestversion_simple" == "" -o "$latestversion_pattern" == "" ]; then
        echo -e "$ERROR Could not fetch the latest version!"
        exit 1;
    fi
    echo -e "$INFO Checking:	${latestversion_simple}"
    echo -e "$INFO Searching:	./*/*${latestversion_pattern}.img"
    ImgFile=$(find ./ -type f -name "*${latestversion_pattern}.img" -print; 2>/dev/null)

    #get the latest firmware either from subdirs or internet
	case `echo $ImgFile | wc -w` in
		0)
			echo -e "$INFO Searching:    ./*/*${latestversion_pattern}.deb"
			mkdir -p firmware
			image_deb=$(find ./ -type f -name "*${latestversion_pattern}.deb" -print; 2>/dev/null)
			ImgFile="firmware/rootfs_$latestversion_simple.img"
			if ! [ `echo $image_deb | wc -w` == 1 ]; then
				image_deb="firmware/rootfs_$latestversion_simple.deb"
				echo -e "$INFO Downloading: $image_deb"
				echo -en "$INPUT Download this firmware and continue? [y/n]"
				read userAnswer
				if [ "$userAnswer" != "y" ] || [ "$userAnswer" != "Y" ] || [ "$userAnswer" != "yes" ]; then
					echo -e "$ERROR Ok, Quit."
					exit 1
				fi
				curl $wdc_latestfirmware > $image_deb
				if [ $? != 0 ]; then
					echo -e "$ERROR Cant download firmware file. Exiting..."
					exit 1;
				fi
			fi
			echo -e "$INFO Extracting $image_deb..."

			ar p $image_deb data.tar.lzma | unlzma | tar -x -C ./firmware
			if [ $? != 0 ]; then
				echo -e "$ERROR Cant unpack firmware file. Exiting..."
				exit 1;
			fi
			mv firmware/CacheVolume/upgrade/rootfs.img firmware/rootfs_${latestversion_simple}.img
			rm -rf firmware/CacheVolume
		;;
		1)
			echo -e "$INFO: File to install: $ImgFile"
			;;
		*)
			echo -e "$ERROR Something went wrong ($ImgFile), exiting..."
			exit 1
	esac
fi

echo ""
echo -e "$WARNING Prepare done. Check if all correct:"
echo "
WD Disk:	$HDD
Firmware:	$ImgFile
Format all?	$Format"
echo ""

echo -en "$INPUT This is correct? [y/n]"
read userAnswer
if [ "$userAnswer" != "y" ] || [ "$userAnswer" != "Y" ] || [ "$userAnswer" != "yes" ]; then
    exit 1;
fi

# Rewrite partition table
if [ $Format = "True" ]; then
    backgroundPattern="${backgroundPattern:-0}"

    dd if=/dev/zero of=$HDD bs=1M count=32
    badblocks -swf -b 1048576 -t ${backgroundPattern} ${HDD} 16 0
    sync
    sleep 2

    parted $HDD --align optimal <<EOL
mklabel gpt
mkpart primary 528M  2576M
mkpart primary 2576M 4624M
mkpart primary 16M 528M
mkpart primary 4624M -1M
set 1 raid on
set 2 raid on
quit
EOL

    sync
    sleep 1
	echo -e "$INFO Formatting $HDD done."
    mkfs.ext4 -F -b 65536 -m 0 $HDDData >> $LogFile 2>&1
    if [ $? != 0 ]; then
		echo -e "$ERROR Formatting $HDDData fail. Check $LogFile"
		exit 1
    fi
    sync
fi

# Write the image to the raid disk
echo -e "$INFO Creating raid..."
mdadm --create --force $RootfsDev --verbose --metadata=0.9 --raid-devices=2 --level=raid1 --run $HDDRoot1 $HDDRoot2 >> $LogFile 2>&1
mdadm --wait $RootfsDev >> $LogFile 2>&1
sync
mkfs.ext3 -c -b 4096 $RootfsDev >> $LogFile 2>&1
sync

echo -e "$INFO Raid creation: Done."
echo -e "$INFO Installing firmware to disk... "
dd if=$ImgFile of=$RootfsDev >> $LogFile 2>&1
echo -e "$INFO Installation done. Mounting FS..."
mount $RootfsDev $RootfsDir >> $LogFile 2>&1
if [ $? != 0 ]; then
	echo -e "$ERROR Mounting $RootfsDev fail. Check $LogFile"
	exit 1
fi

echo -e "$INFO Mount done. Installation boot script..."
cp $RootfsDir/usr/local/share/bootmd0.scr $RootfsDir/boot/boot.scr
echo -e "$INFO Boot script installed. Enabling SSH..."
echo "enabled" > $RootfsDir/etc/nas/service_startup/ssh
echo -e "$INFO Enabling SSH done. Cleanup..."
umount $RootfsDir >> $LogFile 2>&1
rm -r $RootfsDir >> $LogFile 2>&1
mdadm --stop $RootfsDev >> $LogFile 2>&1
mdadm --wait $RootfsDev >> $LogFile 2>&1
echo -e "$INFO Cleanup done."

echo -e "$INFO Fixing raid metadata..."
mdadm --assemble --update=byteorder $RootfsDev $HDDRoot1 $HDDRoot2
mdadm --wait $RootfsDev >> $LogFile 2>&1
mdadm --stop $RootfsDev >> $LogFile 2>&1
mdadm --wait $RootfsDev >> $LogFile 2>&1
echo -e "$INFO Partition data fixed."

echo ""
echo -e "\033[1;33m ================================ $NORM"
echo -e "$INFO All done."
echo ""