#!/bin/bash
#
# Copyright (C) 2007-2009 Gauvain Pocentek <gpocentek@linutop.com>
#
# This program is free software; you can redistribute it and/or
# modify it under the terms of the GNU General Public License
# as published by the Free Software Foundation; either version 2
# of the License, or (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program; if not, write to the Free Software
# Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA

set -e

PROGRAM=$(basename $0)
DEVICE=
MPOINT=$(mktemp -d /tmp/netinst.XXXXXX)
RESUME=0
UPDATE=0
FORCE=0
FORCENEW=0
NOCHECK=0
DEFAULT_URL="http://images.linutop.com/"

echo "Linutop tools - $PROGRAM"
echo "lusm stands for Linutop USB Stick Maker"
echo "$PROGRAM will perform a new Linutop installation"
echo "by downloading the files from a web space."
echo ""

clean()
{
    umount $MPOINT >/dev/null 2>&1 || true
    rmdir $MPOINT 2>/dev/null || true
    rm -f /tmp/livefs.tgz /tmp/md5sums
}

usage()
{
    echo "$PROGRAM [-d device] [-s size] [-u] [-x] [-y] [-h] url"
    echo "options:"
    echo -e "\t-d device  specify the device on which the image will"
    echo -e "\t           be installed (default to $DEVICE)"
    echo -e "\t-u         perform an update, don't format the target"
    echo -e "\t-s size    specify the size of the first partition"
    echo -e "\t           (this option has no effect with -u)"
    echo -e "\t-c         don't perform the md5 sums check"
    echo -e "\t-r         don't resume the download, perform a new"
    echo -e "\t           installation"
    echo -e "\t-y         don't ask for confirmations"
    echo -e "\t-x         debug on"
    echo -e "\t-h         display this help"

    exit $1
}

die_err()
{
    echo "(EE): $1"
    exit 1
}

check_target_device() {
    if [ "$1" = "$ROOT_DEVICE" ]; then
        die_err "You can't use the device running the current system as target."
    fi
}

get_target_size() {
    SOURCE=$1
    local ret=$(parted --script $SOURCE print 1 | grep ^Size: | sed -r "s/^Size: *([^ ]*) .*/\1/")
    # manage MB and GB
    ret=$(echo $ret | sed -r "s/([0-9]*)MB/\1/")
    ret=$(echo $ret | sed -r "s/([0-9]*)GB/\1"000"/")
    echo $ret
}

guess_size() {
    LL=/tmp/livefs.tgz
    wget -c -O $LL $LIVEFS
    LIVEFS_SIZE=$(gunzip --list --name $LL | grep "livefs.tar$" | awk '{print $2}')
    SQUASHFS_SIZE=$(wget --spider -S $SQUASHFS 2>&1 | grep Content-Length: | cut -d: -f2)
    SIZE=$(( ($LIVEFS_SIZE+$SQUASHFS_SIZE) / (1000*1000)))
    # to be safe, an extra 20M
    SIZE=$(( $SIZE+20 ))
    echo $SIZE
}

get_device_infos() {
    UDI=$(hal-find-by-property --key block.device --string $1)
    DEVICE_MODEL=$(hal-get-property --udi $UDI --key storage.model)
    IS_REMOVABLE=$(hal-get-property --udi $UDI --key storage.removable)
    if [ "$IS_REMOVABLE" = "true" ]; then
        DEVICE_SIZE=$(hal-get-property --udi $UDI --key storage.removable.media_size 2>/dev/null)
    else
        DEVICE_SIZE=$(hal-get-property --udi $UDI --key storage.size 2/dev/null)
    fi
    [ -z "$DEVICE_SIZE" ] && return
    DEVICE_SIZE=$(( $DEVICE_SIZE / 1000000 ))
    if [ $DEVICE_SIZE -ge 1000 ]; then
        DEVICE_SIZE="$(( $DEVICE_SIZE / 1000 ))G"
    else
        DEVICE_SIZE="$DEVICE_SIZE"M
    fi
    echo "$DEVICE_MODEL ($DEVICE_SIZE)"
}

# Partition the target
partition () {
    # $1: device
    # $2: size of the first partition
    # TODO: possibility to add a swap (for large drives)

    DEVICE=$1
    SIZE=$2
    SIZEO=$(($SIZE+1))

    echo "Creating the partitions..."
    partition_list=$(parted --script $DEVICE print | grep "^ [1-9]" | awk '{print $1}')
    for part in $partition_list; do
        if grep -q "^$DEVICE$part " /proc/mounts; then
            umount $DEVICE$part >/dev/null 2>&1
            [ $? -eq 0 ] || die_err "Unable to umount the partitions."
        fi

        # if we've removed an extended partition earlier the primary partitions
        # that it contained have also been removed. So we make sure that a
        # partition still exists before trying to remove it.
        t=$(parted --script $DEVICE print | grep "^ $part")
        [ -z "$t" ] && continue
        parted --script $DEVICE rm $part || die_err "Unable to delete parition $part."
    done
    # erase MBR
    dd if=/dev/zero of=$DEVICE bs=512 count=1 >/dev/null 2>&1

    parted --script $DEVICE mklabel msdos
    sleep 3

    parted --script $DEVICE mkpartfs primary fat32 0 $SIZE
    sleep 5
    # we mount the partition to make sure it's usable
    mount -o rw -t vfat "$DEVICE"1 $MPOINT >/dev/null 2>&1
    if [ $? -ne 0 ]; then
        dir_err "The first partition is not usable."
    else
        umount "$DEVICE"1
    fi

    parted --script $DEVICE mkpartfs primary ext2 $SIZEO 100%
    sleep 5
    mount -o rw -t ext2 "$DEVICE"2 $MPOINT >/dev/null 2>&1
    if [ $? -ne 0 ]; then
        dir_err "The second partition is not usable."
    else
        umount "$DEVICE"2
    fi

    parted --script $DEVICE set 1 boot on
    sleep 5
}

ask_for_device () {
    local ALL_DEVICES=$(ls /dev/[sh]d?)
    local POSSIBLE_DEVICES=
    for d in $ALL_DEVICES; do
        [ $d != "$ROOT_DEVICE" ] && POSSIBLE_DEVICES="$POSSIBLE_DEVICES $d"
    done

    [ -z "$POSSIBLE_DEVICES" ]&& \
        die_err "$PROGRAM could not find a target device. Use the -d option to specify one."

    # We need to keep track of the association number/device
    T_FILE=$(mktemp)
    NB=0
    while ! grep -q "^$NB " $T_FILE; do
        echo -n > $T_FILE
        COUNT=0
        echo "Select a device on which to install the new system (<enter> to cancel):"
        for d in $POSSIBLE_DEVICES; do
            COUNT=$(( $COUNT + 1 ))
            INFOS=$(get_device_infos $d)
            [ -z "$INFOS" ] && continue
            echo "$COUNT $d" >> $T_FILE
            echo " $COUNT: $INFOS"
        done
        echo
        echo -n "Device number to use: "
        read NB
        [ -z "$NB" ] && exit 0
    done

    DEVICE=$(grep "^$NB " $T_FILE | cut -d ' ' -f 2)
    rm $T_FILE
}

trap 'clean' INT TERM EXIT

# make sure we speak english
LANG=C

# Are we root
if [ $USER != "root" ]; then
    echo "WARNING: You need to run this script as root:"
    echo "  sudo $PROGRAM <url>"
    exit 1
fi

ROOT_DEVICE=$(grep "^/dev/.*1 /cdrom " /proc/mounts | awk '{print $1}' | cut -d 1 -f 1)

while getopts "d:s:ucrhxy" options
do
    case "$options" in
        d)
          DEVICE=$OPTARG;;
        s)
          SIZE=$OPTARG;;
        u)
          UPDATE=1;;
        c)
          NOCHECK=1;;
        r)
          FORCENEW=1;;
        y)
          FORCE=1;;
        h)
          usage 0;;
        x)
          set -x;;
        *)
          usage 1;;
    esac
done

# -y needs -d DEVICE
if [ $FORCE -eq 1 -a -z "$DEVICE" ]; then
    die_err "The -y option can't be used without the -d option."
fi

# user can't force the partition size on update
if [ "$UPDATE" = 1 ]; then
    unset SIZE
fi

shift $(($OPTIND-1))
if [ -z "$1" ]; then
    FULL_VERSION=$(grep ^LINUTOP_CODENAME /etc/lsb-release | cut -d'=' -f2)
    T=$(echo $FULL_VERSION | sed -r "s/^.*(l2|l3|eee)/\1/")
    case "$T" in
        l3) EXT=3;;
        eee) EXT=eee;;
        *) EXT=2;;
    esac
    BASE_URL="$DEFAULT_URL"current"$EXT"
else
    case "$1" in
    http://* | ftp://*)
        BASE_URL=$1
        ;;
    *)
        BASE_URL="$DEFAULT_URL$(echo $1 | tr A-Z a-z)"
        ;;
    esac
fi

LIVEFS="$BASE_URL/livefs.tgz"
SQUASHFS="$BASE_URL/filesystem.squashfs"
INITRD="$BASE_URL/livefs/initrd.gz"
VMLINUZ="$BASE_URL/livefs/vmlinuz"
MD5SUMS="$BASE_URL/md5sums"

# If no device has been specified, try to find the ones we could use
[ -z "$DEVICE" ] && ask_for_device

if [ ! -e "$DEVICE" ]; then
    die_err "No device found. Please plug it in and restart the script."
fi

check_target_device "$DEVICE"

DLREADY=
for u in $LIVEFS $SQUASHFS $INITRD $VMLINUZ $MD5SUM; do
    if ! wget --spider $u >/dev/null 2>&1; then
        DLREADY="n"
        break
    fi
done
[ -z $DLREADY ] || die_err "Unable to fetch the files at $BASE_URL."

# make sure that we have the needed packages
unset PKGTOINSTALL
for soft in syslinux parted; do
    which $soft >/dev/null || PKGTOINSTALL="$PKGTOINSTALL $soft"
done
if [ -n "$PKGTOINSTALL" ]; then
    SUCCESS=0
    if $(which apt-get); then
        apt-get install -y --force-yes $PKGTOINSTALL && SUCCESS=1
    elif $(which yum); then
        yum install $PKGTOINSTALL && SUCCESS=1
    fi
    [ "$SUCCESS" -eq 0 ] && die_err "Please install the following packages:$PKGTOINSTALL"
fi

echo
DEVICE_INFOS=$(get_device_infos $DEVICE)
echo "Installing the image located at $BASE_URL on $DEVICE_INFOS"
echo

# Stop thunar to avoid problems with automounting
if [ -n "$SUDO_USER" ] && ps -U $SUDO_USER | grep -q Thunar >/dev/null; then
    sudo -u $SUDO_USER killall Thunar
fi

if grep $DEVICE /proc/mounts >/dev/null ; then
    if [ "$FORCE" != 1 ]; then
        echo
        echo    "The device is already mounted; would you like me to umount it?"
        echo -n "Please type \"yes\" if you really want to continue: "
        read confirm
        case $confirm in
            "yes") ;;
            *) exit 1;;
        esac
    fi
    for mounted_part in $(cat /proc/mounts | grep "^$DEVICE" | cut -d" " -f1); do
        umount $mounted_part  || die_err "Unable to umount $mounted_part"
    done
fi

if [ "$FORCENEW" != 1 ]; then
    # Do we have a device and a stamp on it
    if [ -e "$DEVICE"1 ]; then
        if mount -o rw -t vfat "$DEVICE"1 $MPOINT >/dev/null 2>&1; then
            [ -f $MPOINT/net-stamp ] && RESUME=1
            umount "$DEVICE"1
        fi
    fi
fi

# Format the device only if no resume and no update
if [ "$UPDATE" != 1 ]; then
    if [ $RESUME -ne 1 ]; then
        # Partition USB key
        if [ "$FORCE" != 1 ]; then
            echo    "You are about to erase the contents of this device:"
            echo    "  $DEVICE_INFOS"
            echo
            echo -n "Please type \"yes\" if you really want to continue: "
            read confirm
            case $confirm in
                "yes") ;;
                *) exit 1;;
            esac
        fi

        # guess the needed size for the 1st partition, if not set by the user
        if [ -z "$SIZE" ]; then
            echo "Computing the first partition size..."
            SIZE=$(guess_size)
        fi

        partition "$DEVICE" "$SIZE"
        mount -o rw -t vfat "$DEVICE"1 $MPOINT >/dev/null 2>&1
        touch $MPOINT/net-stamp
        umount "$DEVICE"1
    fi
else
    if [ $RESUME -eq 1 ]; then
        echo "The system on $DEVICE is not fully installed."
        echo "Please finish the initial installation first."
        exit 1
    fi
fi

# mount the 1st partition
mount -o rw -t vfat "$DEVICE"1 $MPOINT >/dev/null 2>&1

if [ "$UPDATE" = 1 -a  "$FORCE" != 1 ]; then
    if [ $RESUME -eq 1 ]; then
        echo "The system on $DEVICE is not fully installed."
        echo "Please finish the initial installation first."
        exit 1
    fi

    # Make sure than we have enough room to store the update
    CUR_SIZE=$(get_target_size $DEVICE)
    if [ $SIZE -gt $CUR_SIZE ]; then
        echo "The first partition of $DEVICE is to small to host the update."
        echo "You can either manually resize it, or perform a fresh install."
        exit 1
    fi

    echo    "You are about to replace the system on the device $DEVICE."
    echo -n "Please type \"yes\" if you really want to continue: "
    read confirm
    case $confirm in
    "yes") ;;
    *) exit 1;;
    esac
fi

# ! update => downloading everything
if [ "$UPDATE" != 1 ]; then
    if [ ! -f /tmp/livefs.tgz ]; then
        echo "Downloading the livefs..."
        wget -c -O /tmp/livefs.tgz $LIVEFS || die_err "Can't reach the livefs archive."
    fi
    echo "Uncompressing the livefs..."
    tar -xz --no-same-owner --no-same-permissions -f /tmp/livefs.tgz -C $MPOINT >/dev/null 2>&1
    echo "Downloading the system. Please wait, it may be long..."
    wget -c -O $MPOINT/casper/filesystem.squashfs $SQUASHFS || die_err "Can't reach the filesystem image."
else
    # TODO: test if we have a valid linutop system
    ISLINUTOP=1
    for f in casper/filesystem.squashfs vmlinuz initrd.gz; do
        if [ ! -f $MPOINT/$f ]; then
            ISLINUTOP=0
            break
        fi
    done
    if [ $ISLINUTOP -eq 0 ]; then
        echo "The target device doesn't contain a valid linutop system."
        exit 1
    fi

    echo "Downloading the system. Please wait, it may be long..."
    wget -O $MPOINT/casper/filesystem.squashfs $SQUASHFS || die_err "Can't reach the filesystem image."
    echo "Downloading the kernel..."
    wget -O $MPOINT/vmlinuz $VMLINUZ || die_err "Can't reach the kernel."
    echo "Downloading the initramfs..."
    wget -O $MPOINT/initrd.gz $INITRD || die_err "Can't reach the initramfs."
fi

# copy the boot menu
cp /usr/lib/syslinux/vesamenu.c32 $MPOINT/

# umount and remount the device to make sure than everything
# is written
umount $MPOINT >/dev/null 2>&1
mount -o rw -t vfat "$DEVICE"1 $MPOINT >/dev/null 2>&1

# Now the rest of the data if we don't do an update
if [ "$UPDATE" != 1 ]; then
    SYSLINUX=$BASE_URL/syslinux.cfg
    # do we have a custom syslinux.cfg?
    CUSTOM_SYSLINUX=1
    wget --spider $SYSLINUX >/dev/null 2>&1 || CUSTOM_SYSLINUX=0
    if [ $CUSTOM_SYSLINUX -eq 1 ]; then
        echo
        echo "Installing a custom syslinux.cfg..."
        wget $SYSLINUX -O $MPOINT/syslinux.cfg
        # avoid running lconfeditor
        touch $MPOINT/lconfeditor-stamp
    fi
fi

if [ "$NOCHECK" != 1 ]; then
    echo "Verifying md5 sums..."

    # get the data file
    rm -f /tmp/md5sums
    wget $MD5SUMS -O /tmp/md5sums

    # if we have a custom syslinux.cfg don't check its md5 sum
    if [ $CUSTOM_SYSLINUX -eq 1 ]; then
        sed -i '/^syslinux.cfg/d' /tmp/md5sums
    fi

    good=1
    if [ "$UPDATE" != 1 ]; then
        while read f sum; do
            new_sum=$(md5sum "$MPOINT"/"$f" | awk '{print $1}')
            if [ "$sum" != "$new_sum" ]; then
                good=0
                break
            fi
        done < /tmp/md5sums
    else
        for f in "casper/filesystem.squashfs" vmlinuz initrd.gz; do
            sum=$(cat /tmp/md5sums | grep ^$f | cut -d " " -f 2)
            new_sum=$(md5sum "$MPOINT"/"$f" | awk '{print $1}')
            if [ "$sum" != "$new_sum" ]; then
                good=0
                break
            fi
        done
    fi
fi

# Clean up
rm -f $MPOINT/net-stamp
umount $MPOINT >/dev/null 2>&1
syslinux "$DEVICE"1

if [ "$UPDATE" != 1 ]; then
    # Do we have data to install on the second partition?
    EXTRA=$BASE_URL/data.tgz
    HAVE_EXTRA=1
    wget --spider $EXTRA >/dev/null 2>&1 || HAVE_EXTRA=0
    if [ $HAVE_EXTRA -eq 1 ]; then
        echo
        echo "Downloading and installing extra data..."
        mount -o rw -t ext2 "$DEVICE"2 $MPOINT >/dev/null 2>&1
        HERE=$(pwd)
        cd $MPOINT
        wget $EXTRA -O - | tar zxvf -
        cd $HERE
        umount $MPOINT >/dev/null 2>&1
    fi
fi

if [ "$NOCHECK" = 1 ]; then
    echo "Success."
    echo "You can safely use your new Linutop installation."
    exit 0
fi

# if we've checked the md5sums
if [ "$good" -eq 1 ]; then
    echo "Success."
    echo "Each file copied has been verified."
    echo "You can safely use your new Linutop installation."
else
    echo "Failure."
    echo "An error occured while copying the files."
    echo "Your new installation is likely to be corrupted."
    echo "Please restart the installation."
    exit 1
fi

exit 0