#!/bin/bash

# Copyright (C) 2014-2015 Red Hat Inc.
# SPDX-License-Identifier:  GPL-2.0+

# Automate Media Creation for Fedora ARM 
# Current version
VERSION=1.99.12

# usage message
usage() {
    echo "
Usage: $(basename ${0}) <options>

   --image=IMAGE    - xz compressed image file name
   --target=TARGET  - target board
   --media=DEVICE   - media device file (/dev/[sdX|mmcblkX])
   --selinux=ON/OFF - Turn SELinux off/on as needed
   --norootpass     - Remove the root password
   -y		    - Assumes yes, will not wait for confirmation
   --version	    - Display version and exit
   --resizefs	    - Resize root filesystem to fill media device
   --addconsole     - Add system console to extlinux.conf

Example: $(basename ${0}) --image=Fedora-Rawhide.xz --target=Bananapi --media=/dev/mmcblk0 --selinux=OFF

For list of supported boards please check SUPPORTED-BOARDS file.
"
}

# Set some global variables for the command directory, target board directory,
# and valid targets.
DIR=$(dirname $0)
if [ -d "/usr/share/arm-image-installer/boards.d" ]; then
        BOARDDIR="/usr/share/arm-image-installer/boards.d"
else
        DIR=$(dirname $0)
        BOARDDIR="${DIR}/boards.d"
fi

# check the args
while [ $# -gt 0 ]; do
        case $1 in
                --debug)
                        set -x
                        ;;
                -h|--help)
                        usage
                        exit 0
                        ;;
                --target*)
                        if echo $1 | grep '=' >/dev/null ; then
                                TARGET=$(echo $1 | sed 's/^--target=//')
                        else
                                TARGET=$2
                                shift
                        fi
                        ;;
                --image*)
                        if echo $1 | grep '=' >/dev/null ; then
                                IMAGE=$(echo $1 | sed 's/^--image=//')
                        else
                                IMAGE=$2
                                shift
                        fi
                        ;;
                --media*)
                        if echo $1 | grep '=' >/dev/null ; then
                                MEDIA=$(echo $1 | sed 's/^--media=//')
                        else
                                MEDIA=$2
                                shift
                        fi
                        ;;
                --selinux*)
                        if echo $1 | grep '=' >/dev/null ; then
                                SELINUX=$(echo $1 | sed 's/^--selinux=//')
                        else
                                SELINUX=$2
                                shift
                        fi
                        ;;
                --norootpass)
                        NOROOTPASS=1
                        ;;
                --resizefs)
                        RESIZEFS=1
                        ;;
                --addconsole)
                        CONSOLE=1
                        ;;
                --version)
                        echo "$(basename ${0})-"$VERSION""
                        exit 0
                        ;;
                -y)
                        NOASK=1
                        ;;
                *)
                        echo "$(basename ${0}): Error - ${1}"
                        usage
                        exit 1
                        ;;
        esac
        shift
done

contains() {
    string="$1"
    substring="$2"
    if test "${string#*$substring}" != "$string"
    then
        return 0    # $substring is in $string
    else
        return 1    # $substring is not in $string
    fi
}

# ensure sudo user
if [ "$(whoami)" != "root" ] ; then
        echo "Error: This script requires 'sudo' privileges in order to write to disk & mount media."
        exit 1
fi

# check to make sure populated
if [ "$MEDIA" = "" ] ; then
        usage
        exit 1
fi

# change cubietruck target to uppercase
if [ "$TARGET" = "cubietruck" ]; then
        TARGET="Cubietruck"
fi

# check for boards
if [ "$TARGET" != "" -a ! -e "${BOARDDIR}/${TARGET}" ] ; then
        echo "Error: You must choose a supported board or none at all." 
        usage
        exit 1
fi

# image exists
if [ ! -f "$IMAGE" ] && [ "$IMAGE" != "" ] ; then
        echo "Error: $IMAGE not found! Please choose an existing image."
        exit 1
fi

# device exists
if [ ! -e "$MEDIA" ] ; then
        echo "Error: $MEDIA not found! Please choose an existing device."
        exit 1
fi

clear
# Last chance to back out
echo " "
echo "====================================================="
# Image if included
if [ "$IMAGE" != "" ] ; then
        echo "= Selected Image:                                 "
        echo "= $IMAGE"
fi
echo "= Selected Media : $MEDIA"
# target hardware platform
if [ "$TARGET" != "" ] ; then 
        echo "= U-Boot Target : $TARGET"
fi
# SE Linux On/Off
if [ "$SELINUX" != "" ] ; then
        echo "= SELINUX = $SELINUX"
fi
# Remove root password
if [ "$NOROOTPASS" != "" ] ; then
        echo "= Root Password will be removed."
fi
# Resize root filesystem to fill media device
if [ "$RESIZEFS" != "" ] ; then
        echo "= Root partition will be resized"
fi
# Console to be added
if [ "$CONSOLE" != "" ] ; then
        echo "= Console for $TARGET will be added."
fi
echo "====================================================="
echo " "
echo "*****************************************************"
echo "*****************************************************"
echo "******** WARNING! ALL DATA WILL BE DESTROYED ********"
echo "*****************************************************"
echo "*****************************************************"
if [ "$NOASK" != 1 ] ; then 
        echo " "
        echo " Type 'YES' to proceed, anything else to exit now "
        echo " "
        # wait for agreement
        read -p "= Proceed? " PROCEED
        if [ "$(echo ${PROCEED} | tr [:lower:] [:upper:])" != "YES" ] ; then
                echo "User exit, no image written."
                exit 0
        fi
fi
# umount before starting
umount $MEDIA* &> /dev/null

# Write the disk image to media
if [ "$IMAGE" != "" ] ; then
        echo "= Writing: "
        echo "= $IMAGE "
        echo "= To: $MEDIA ...."
        xzcat $IMAGE | dd of=$MEDIA bs=4M; sync; sleep 3
        echo "= Writing image complete!"
        # read the new partition table
        partprobe "$MEDIA"
fi
# check to see how many partitions on the image
partprobe "$MEDIA"
case $MEDIA in
        "/dev/mmcblk"*)
                if [ -e "$MEDIA"p4 ]; then
                        BOOTPART="${MEDIA}p2"
                        ROOTPART="${MEDIA}p4"
                        PARTNUM=4
                else
                        BOOTPART="${MEDIA}p1"
                        ROOTPART="${MEDIA}p3"
                        PARTNUM=3
                fi
                ;;
        *)
                if [ -e "$MEDIA"4 ]; then
                        BOOTPART="${MEDIA}2"
                        ROOTPART="${MEDIA}4"
                        PARTNUM=4
                else
                        BOOTPART="${MEDIA}1"
                        ROOTPART="${MEDIA}3"
                        PARTNUM=3
                fi
                ;;
esac

# resize root filesystem before mounting
if [ "$RESIZEFS" != "" ] ; then
        echo "= Resizing $MEDIA ...."
        sync
        echo ", +" | sfdisk -N "$PARTNUM" "$MEDIA"
        while [ $? != '0' ]; do
                sleep 5
                echo ", +" | sfdisk -N "$PARTNUM" "$MEDIA"
        done
        
        fsck.ext4 -fy "$ROOTPART"
        partprobe "$MEDIA"
        resize2fs "$ROOTPART"
fi

# make temp mount points
mkdir /tmp/{boot,root} &> /dev/null
mount "$BOOTPART" /tmp/boot &> /dev/null
mount "$ROOTPART" /tmp/root &> /dev/null

# turn off selinux
if [ "$SELINUX" != "" ] ; then
        if [ "$(echo ${SELINUX} | tr [:lower:] [:upper:])" = "OFF" ] ; then
                echo "= Turning SELinux off ..."
                sed -i 's/append/& enforcing=0/' /tmp/boot/extlinux/extlinux.conf
                # turn on selinux
        elif [ "$(echo ${SELINUX} | tr [:lower:] [:upper:])" = "ON" ] ; then
                echo "= Turning SELinux on ..."
                sed -i 's/ enforcing=0//' /tmp/boot/extlinux/extlinux.conf
        fi
fi
# Remove root password
if [ "$NOROOTPASS" = "1" ] ; then
        echo "= Removing the root password."
        sed -i 's/root:x:/root::/' /tmp/root/etc/passwd
fi
# determine uboot and write to disk 
if [ "$TARGET" = "" ]; then
        echo "= No U-boot will be written."
        TARGET="Mystery Board"
else
        PREFIX=/tmp/root
        . "${BOARDDIR}/${TARGET}"
        # Add console
        if [ "$CONSOLE" = "1" ] && [ "$SYSCON" != "" ] ; then
                echo "= Adding console $SYSCON to extlinux.conf ..."
                sed -i "s|append|& console=$SYSCON|" /tmp/boot/extlinux/extlinux.conf
        fi

fi

umount $ROOTPART $BOOTPART &> /dev/null

echo ""
echo "= Installation Complete! Insert into the "$TARGET" and boot."
exit 0

# vi: tabstop=8 softtabstop=0 expandtab shiftwidth=8 smarttab
