#! /bin/bash

#Declare Variables
ORIGINAL_IMAGE=DFIOriginal.IMG
IMAGE=DFIImage.IMG
HEADER=DFIHeader
offset=
OPTIONS="umask=777,shortname=mixed,quiet,utf8"
INPUT=/home/Splashtop/image
FIX=/home/Splashtop/extracted-image
NEWIMAGE=UpdatedDFI.DFI
RESCAPTURE=
VERSION="1.4.10.9"
SAVE_IMAGE=UpdatedIMG.IMG
ORIGINAL_PATH=$(pwd)

#Display help
if [[ "x$1" = "x--help" ]]; then
	echo -e "USAGE:"
	echo -e "	\E[34m$0\E[30m \E[31mFILE\E[30m"
	echo ""
	echo -e "\E[31mFILE\E[30m - The DFI File which should be extracted"
	echo ""
	echo -e	"This script is designed to simplify the process of extracting a DFI file"
	echo -e "and creating an updated version once all modifications are complete."
	echo -e "It provides interactive prompts and can additionally provide an IMG"
	echo -e "file which can be used for creating a bootable USB drive with embeded"
	echo -e "SplashTop."
	echo ""
	echo -e "This script requires the following packages to be installed on your system:"
	echo -e "\E[36m    - squashfs-tools \E[31m[\E[30mversion 3.4\E[31m]\E[30m - \E[34mhttp://squashfs.sourceforge.net\E[30m"
	echo ""
	echo -e "This script must be run as root:"
	echo -e "	sudo bash \E[34m$0\E[31m FILE\E[30m"
	echo ""			
	echo "The Following Applications must be in the same directory as the script:"
	echo -e "    - \E[36mDVM-Header.exe\E[30m"
	echo -e "    - \E[36mDVM-Version.exe\E[30m"
	echo ""		
	echo -e "For more information please visit: \E[34mwww.sierrasoftworks.com/expressgate\E[30m"
	echo ""
	echo -e "Copyright Benjamin Pannell 2010"	
	exit 1
fi

#If no arguments are present then display help as well
if [[ "x$#" = "x0" ]]; then
	echo -e "USAGE:"
	echo -e "	\E[34m$0\E[30m \E[31mFILE\E[30m"
	echo ""
	echo -e "\E[31mFILE\E[30m - The DFI File which should be extracted"
	echo -e "\E[31mFILE\E[30m must be specified"
	echo ""
	echo -e	"This script is designed to simplify the process of extracting a DFI file"
	echo -e "and creating an updated version once all modifications are complete."
	echo -e "It provides interactive prompts and can additionally provide an IMG"
	echo -e "file which can be used for creating a bootable USB drive with embeded"
	echo -e "SplashTop."
	echo ""
	echo -e "This script requires the following packages to be installed on your system:"
	echo -e "    - \E[36msquashfs-tools \E[31m[\E[30mversion 3.4\E[31m]\E[30m - \E[34mhttp://squashfs.sourceforge.net\E[30m"
	echo ""
	echo -e "This script must be run as root:"
	echo -e "	sudo bash \E[34m$0\E[31m FILE\E[30m"
	echo ""			
	echo "The Following Applications must be in the same directory as the script:"
	echo -e "    - \E[36mDVM-Header.exe\E[30m"
	echo -e "    - \E[36mDVM-Version.exe\E[30m"
	echo ""
	echo -e "For more information please visit: \E[34mwww.sierrasoftworks.com/expressgate\E[30m"
	echo ""
	echo -e "Copyright Benjamin Pannell 2010"	
	exit 1
fi

if [ -f "$ORIGINAL_PATH/DVM-Version.exe" ]; then
	echo -e "\E[32mDVM-Version.exe Detected\E[30m"
else
	echo -e "USAGE:"
	echo -e "	\E[34m$0\E[30m \E[31mFILE\E[30m"
	echo ""
	echo -e "\E[31mFILE\E[30m - The DFI File which should be extracted"
	echo ""		
	echo "The Following Applications must be in the same directory as the script:"
	echo -e "    - \E[36mDVM-Header.exe\E[30m"
	echo -e "    - \E[31mDVM-Version.exe\E[30m"
	echo ""
	exit 1
fi

if [ -f "$ORIGINAL_PATH/DVM-Header.exe" ]; then
	echo -e "\E[32mDVM-Header.exe Detected\E[30m"
else
	echo -e "USAGE:"
	echo -e "	\E[34m$0\E[30m \E[31mFILE\E[30m"
	echo ""
	echo -e "\E[31mFILE\E[30m - The DFI File which should be extracted"
	echo ""		
	echo "The Following Applications must be in the same directory as the script:"
	echo -e "    - \E[31mDVM-Header.exe\E[30m"
	echo -e "    - \E[36mDVM-Version.exe\E[30m"
	echo ""
	exit 1
fi

# Make sure only root can run the script
if [[ $EUID -ne 0 ]]; then
	echo -e "USAGE:"
	echo -e "	\E[34m$0\E[30m \E[31mFILE\E[30m"
	echo ""
	echo -e "\E[31mFILE\E[30m - The DFI File which should be extracted"
	echo ""
	echo -e "\E[31mThis script must be run as root:\E[30m"
	echo -e "	sudo bash \E[34m$0\E[31m FILE\E[30m"
	exit 1
fi


#Prompt user to extract original files for modification
echo -e -n "\E[31mWould you like to extract the original SQX files for modification? [y/n]:\E[30m"
read -n 1 EXTRACT_ORIGINALS
echo ""
if [[ "x$EXTRACT_ORIGINALS" = "xy" ]]; then
echo -e "\E[31mYou have chosen to extract the original SQX files. This may cause your image to work incorrectly\E[30m"
echo -e "\E[30mWhen prompted, it is recommended that you choose to patch the version file to prevent SplashTop from checking the modified SQX files\E[30m"
echo -e "\E[30mIt is also recommended that any SQX files which you do not modify be removed from the \E[36m$FIX\E[30m directory before updating the image.\E[30m"
echo -e "\E[30mThis may result in better system stability and a higher chance of success.\E[30m"
echo -e "\E[30mFor more information please visit: \E[34mwww.sierrasoftworks.com/expressgate\E[30m"
fi

#Prompt user to keep updated IMG file for flashing to USB
echo -e -n "\E[31mWould you like to keep a copy of the new IMG file? [y/n]:\E[30m"
read -n 1 KEEP_IMAGE
echo ""
if [[ "x$KEEP_IMAGE" = "xy" ]]; then
echo -e "\E[30mYou have chosen to keep a copy of the IMG file. You can use this file to create a SplashTop USB stick.\E[30m"
fi

#Extract Image from DFI File
echo -e "\E[32mProcessing File [\E[36m$1\E[32m]:\E[30m"
echo ""
echo -e -n "\E[32mCopying DFI File [\E[36m$ORIGINAL_IMAGE.gz\E[32m]...\E[31m"
cp $1 $ORIGINAL_IMAGE.gz
echo -e "\E[32mDone\E[30m"
echo -e -n "\E[32mExtracting DFI File [\E[36m$ORIGINAL_IMAGE\E[32m]...\E[31m"
gzip -d $ORIGINAL_IMAGE.gz
echo -e "\E[32mDone\E[30m"
echo -e "\E[32mRemoving Header [\E[36m$HEADER\E[32m]...\E[30m"
RESCAPTURE=$(dd if=$ORIGINAL_IMAGE of=$HEADER bs=1 count=32)
echo -e "\E[32mDone\E[30m"
echo -e "\E[32mExtracting IMG File From Archive [\E[36m$IMAGE\E[32m]...\E[30m"
RESCAPTURE=$(dd if=$ORIGINAL_IMAGE of=$IMAGE bs=48 skip=1)
echo -e "\E[32mDone\E[30m"
echo -e -n "\E[32mRemoving Extracted Archive [\E[33m$ORIGINAL_IMAGE\E[32m]...\E[30m"
rm $ORIGINAL_IMAGE
echo -e "\E[32mDone\E[30m"
echo ""

offset=$(($(sfdisk -d $IMAGE|grep start=|head -n1|sed 's/.*start=\s*//;s/,.*//')*512))

#Create Working Directories
echo -e "\E[32mCreating Working Directories\E[30m"
echo -e "	[\E[36m$INPUT\E[36m]\E[30m"
echo -e "	[\E[36m$FIX\E[36m]\E[30m"
rm -rf $FIX
mkdir -p $INPUT $FIX

#Mount Image to Working Directory
echo -e "\E[32mChecking to make sure that working directories are not in use\E[30m"
umount $INPUT
echo -e "\E[32mMounting Image [\E[30m$IMAGE\E[32m]\E[31m"
MOUNTRES=$( mount -w -o loop,offset=$offset,$OPIONS $IMAGE $INPUT )
MOUNTRES2=$(echo $MOUNTRES | grep "filesystem")
if [ "x$MOUNTRES2" != "x" ]; then
	echo -e "\E[31mError Mounting Image\E[30m"
	exit 1
fi
if [[ "x$MOUNTRES2" = "x"]]; then
	echo -e "\E[32mImage Mounted\E[30m"
fi

cd $INPUT

#Decompress SQX files
if [[ "x$EXTRACT_ORIGINALS" = "xy" ]]; then
echo -e "\E[32mDecompressing SQX Files\E[30m"
for x in *.sqx; 
do 
	echo -e "\E[34m--------$x--------\E[30m";
	unsquashfs -dest "$FIX/"${x/.sqx/} $x;
	echo -e "\E[34m--------Completed--------\E[30m";
	echo -e ""
done
fi

#Reset Colour
echo -e "\E[30m"

#Display editing Instructions to user

echo -e "\E[32mThe DFI File has been successfully extracted!\E[30m"
echo -e "You can edit the contents of the image by doing the following:"
echo ""
echo "a) Change files in the directory you want to edit."
echo "b) Create new packages by creating a new folder with the correct"
echo "   package name and contents."
echo "c) Type \"exit\" when done to apply changes."
echo

#Grant all users permissions on files
chmod -R ugo+rwx $FIX
chmod -R ugo+rwx $INPUT

#Navigate to Working Directory
cd $FIX
#Allow user access
echo -e "\E[34mCurrent Directory is: \E[30m$FIX"
bash

#Remove all permissions on Working Directories
chmod go-rwx $FIX

#Re-Compress all SQX Files
cd $FIX
echo -e "\E[32mRecompressing SQX Files\E[31m"
for x in $(ls -p --format=single-column | grep -E "^.*/$"); do
	echo -e "\E[34m${x/\/}.sqx\E[30m";
	mksquashfs $FIX/$x "$FIX/${x/\/}.sqx" -noappend;
done



#Copy all new SQX files to expanded directory
cd $FIX
echo -e "\E[32mCopying new SQX Files to image\E[30m"
for x in *.sqx;do
	echo -e -n "\E[36m    - $x...\E[30m";
	cp -f -u $FIX/$x $INPUT/$x;
	echo -e "\E[32mDone\E[30m";
done

cd $INPUT

#Patch version file only if original SQX files have been modified
#if [[ "x$EXTRACT_ORIGINALS" = "xy" ]]; then

echo ""
echo -e "\E[31mFile Integrity Checking"
echo -e "\E[30mIt is recommended that you patch the version file, this will enable file intergity checking on all files that you have added as well as allowing you to modify original files without causing SplashTop to fail."
echo ""
echo -e -n "\E[31mWould you like to patch the version file? [y/n]: "
read -n 1 PATCH_VERSION
echo ""
echo ""


if [[ "x$PATCH_VERSION" = "xy" ]]; then

	echo -e "\E[32mPatching Version file to prevent MD5 check...\E[30m"
	cp $INPUT/version $INPUT/versionBackup
	mono $ORIGINAL_PATH/DVM-Version.exe $INPUT $INPUT/version $VERSION
	echo -e "\E[32mDone\E[30m"

fi

#Unmount IMG
echo -e -n "\E[32mUnmounting Image [\E[30m$INPUT\E[32m]...\E[31m"
RESCAPTURE=$(umount $INPUT)
echo -e "\E[32mDone\E[30m"

#Remove working Directory
echo -e -n "\E[32mRemoving Working Directory [\E[30m$FIX\E[32m]...\E[31m"
rm -rf $FIX
echo -e "\E[32mDone\E[30m"

cd $ORIGINAL_PATH

#If the user wanted to keep a copy then grant them full access
if [[ "x$KEEP_IMAGE" = "xy" ]]; then
	echo -e -n "\E[32mSaving Copy of Image...\E[30m"
	cp $IMAGE $SAVE_IMAGE
	chmod ugo+rwx $SAVE_IMAGE
	echo -e "\E[32mDone\E[30m"	
fi


#Create Updated Header File
echo -e "\E[32mCreating Updated Header File\E[30m"
mono $ORIGINAL_PATH/DVM-Header.exe $IMAGE $HEADER $VERSION

#Recreate DFI File
echo -e -n "\E[32mInserting Header into Image...\E[31m"
cat $HEADER $IMAGE > Updated.IMG
rm new-md5sum
echo -e "\E[32mDone\E[30m"
echo -e -n "\E[32mRe-Packaging Image into DFI File...\E[31m"
gzip -9 Updated.IMG
mv Updated.IMG.gz $NEWIMAGE
echo -e "\E[32mDone\E[30m"

#Granting user permissions on new image
echo -e -n "\E[32mGranting All Users Permissions on new DFI...\E[30m"
chmod ugo+rwx $NEWIMAGE
echo -e "\E[32mDone\E[30m"

#Removing old image parts
echo -n -e "\E[32mRemoving Old Image Parts...\E[31m"
cd $ORIGINAL_PATH
rm $IMAGE

rm $HEADER
echo -e "\E[32mDone\E[30m"

#Removing old image directory
echo -e -n "\E[32mUnmounting Image [\E[36m$INPUT\E[32m]...\E[31m"
RESCAPTURE=$(umount $INPUT)
echo -e "\E[32mDone\E[30m"
echo -e -n "\E[32mRemoving Image Directory [\E[36m$INPUT\E[32m]...\E[31m"
rm -rf $INPUT
echo -e "\E[32mDone\E[30m"

#Done
echo ""
echo -e "\E[34m################\E[30m"
echo -e "\E[34m#   Complete   #\E[30m"
echo -e "\E[34m################\E[30m"