#!/bin/bash
export PATH=/bin:/usr/bin:/sbin:/usr/sbin

partimageinstall () {
    # based on install.sh from xcat 1.3 by Egan Ford

    imgpath=$(dirname $1) # needs to be path to image files
    myimage=$(basename $1 .img) # needs to be image without path or .gz
    disk=$(basename $2) # needs to be disk - i.e. sda
    
    altdisk=$disk
    if [[ "$disk" = "sdb" ]] &&
        dmesg | grep -q 'SCSI device sdb' &&
        (! dmesg | grep -q 'SCSI device sda' || ls /dev/disk/by-id -l | grep -q sda | grep -q usb) &&   # sda doesn't exist or it exists and is usb
        [[ -f $myimage-sda.sfdisk ]] &&
        [[ ! -f $myimage-sdb.sfdisk ]] ; then
            altdisk="sda"
    fi

    DONE=0
    OLD=""

    cd $imgpath
    NUMPART=$(cat $myimage-$altdisk.sfdisk | grep "^/dev/$altdisk" | grep -v Win95 | grep -v Extended | grep -v Empty | grep -v "Id= 0" | grep -v "Id= 5" | grep -v "Id=82" | grep -v swap | awk '{print $1}' | wc -l)
    for i in $(cat $myimage-$altdisk.sfdisk | grep "^/dev/$altdisk" | grep -v Win95 | grep -v Extended | grep -v Empty | grep -v "Id= 0" | grep -v "Id= 5" | grep -v "Id=82" | grep -v swap | awk '{print $1}')
    do
        PARTNAME=$(basename ${i/$altdisk/$disk})
        if (("$NUMPART" > "1"))
        then
            #PARTNAMEIMG=${PARTNAME/#??/part}
            PARTNAMEIMG=$(basename $i | sed 's/^../part/')
            INPUTFILE=$myimage-$PARTNAMEIMG.gz
        else
            INPUTFILE=$myimage.gz
        fi
        if [ "$DONE" = "0" ]
        then
            if [ -r $myimage-$altdisk.sfdisk ]
            then
                sfdisk /dev/$disk <$myimage-$altdisk.sfdisk
            fi
            if [ -r $myimage-$altdisk.mbr ]
            then
                dd if=$myimage-$altdisk.mbr of=/dev/$disk
                sfdisk /dev/$disk <$myimage-$altdisk.sfdisk
            else
                OLD=".6.0"
                partimage$OLD -b -f3 restmbr $INPUTFILE
            fi
            sfdisk /dev/$disk -R
            DONE=1
        fi
        echo running partimage$OLD -b -f3 restore /dev/$PARTNAME $INPUTFILE
        logger -t xcat "running partimage$OLD -b -f3 restore /dev/$PARTNAME $INPUTFILE"
        sleep 5
        if ! partimage$OLD -B x=x -f3 restore /dev/$PARTNAME $INPUTFILE 2>> /tmp/partimageng ; then
            echo partimage exited with a non-zero status, failing
            logger -t xcat "partimage exited with a non-zero status, failing"
            return 1
        fi
    done

    tput clear
    tput sgr0
    tput cnorm

    for i in $(sfdisk -l /dev/$disk | grep "Linux swap" | awk '{print $1}')
    do
        echo "Setting up swap on $i"
        if [ `echo $myimage | grep fc7` ]
        then
            label=SWAP-`echo $(basename $i) | sed -e "s/[sh]d/part/"`
        else
            label=SWAP-$(basename $i)
        fi
        mkswap -L $label $i
    done
}

partimagesave () {
    # based on save.sh from xcat 1.3 by Egan Ford
    imgpath=$(dirname $1) # needs to be path to image files
    myimage=$(basename $1 .img) # needs to be image without path or .gz
    disk=$(basename $2) # needs to be disk - i.e. sda
    
    altdisk=$disk
    if [[ "$disk" = "sdb" ]] &&
        dmesg | grep -q 'SCSI device sdb' &&
        (! dmesg | grep -q 'SCSI device sda' || ls /dev/disk/by-id -l | grep -q sda | grep -q usb) ; then   # sda doesn't exist or it exists and is usb
            altdisk="sda"
    fi

    logger -t xcat "Getting partitons for $disk"
    echo "Partitions for $disk"
    sfdisk -l /dev/$disk
    echo "==="

    NUMPART=$(sfdisk -l /dev/$disk | grep "^/dev/$disk" | grep -v Win95 | grep -v Extended | grep -v Empty | grep -v "Id= 0" | grep -v swap | awk '{print $1}' | wc -l)
    logger -t xcat "number of partitions found: $NUMPART"

    dd if=/dev/$disk of=$imgpath/$myimage-$altdisk.mbr count=1 bs=512
    sfdisk /dev/$disk -d >$imgpath/$myimage-$disk.sfdisk.tmp
    if [ -s $imgpath/$myimage-$disk.sfdisk.tmp ]
    then
        mv -f $imgpath/$myimage-$disk.sfdisk.tmp $imgpath/$myimage-$disk.sfdisk
		  if [[ "$disk" != "$altdisk" ]]; then
			  cat $imgpath/$myimage-$disk.sfdisk | sed "s/$disk/$altdisk/" > $imgpath/$myimage-$altdisk.sfdisk
			  rm -f $imgpath/$myimage-$disk.sfdisk
        fi
    else
        rm -f $imgpath/$myimage-$disk.sfdisk.tmp
    fi
    logger -t xcat "saving individual partitions"
    for i in $(sfdisk -l /dev/$disk | grep "^/dev/$disk" | grep -v Win95 | grep -v Extended | grep -v Empty | grep -v swap | awk '{print $1}')
    do
        PARTNAME=$(basename $i)
        logger -t xcat "working with $PARTNAME"
        if (("$NUMPART" > "1"))
        then
            PARTNAMEIMG=${PARTNAME/#??/part}
            OUTPUTFILE=$imgpath/$myimage-$PARTNAMEIMG.gz
        else
            OUTPUTFILE=$imgpath/$myimage.gz
        fi
        echo running partimage -z1 -f3 -odbc save /dev/$PARTNAME $OUTPUTFILE
        logger -t xcat "running partimage -z1 -f3 -odbc save /dev/$PARTNAME $OUTPUTFILE"
        sleep 5
        if ! partimage -z1 -f3 -odbc save /dev/$PARTNAME $OUTPUTFILE; then
            echo partimage exited with a non-zero status, failing
            logger -t xcat "partimage exited with a non-zero status, failing"
            exit # do we exit or return a value to fail?
            return 1
        fi
    done
    tput clear
    tput sgr0
    tput cnorm

    cd $imgpath
    for i in *.000
    do
        mv $i ${i%%.000}
    done
    return 0
}

exec 2>/tmp/partimageng.log
set -x

depmod -a
modprobe nfs
modprobe scsi_transport_sas
modprobe mptbase
modprobe mptscsih
modprobe mptsas
modprobe mptspi

# check for problem of $MASTER being an unknown host
if ! ping -c 1 $MASTER &> /dev/null; then
    echo "CRITICAL ERROR: could not ping $MASTER, rebooting"
    sleep 3
    /sbin/reboot
    exit
fi

logger -t xcat "mounting $IMAGELIBSERVER:$IMAGELIBINSTALLDIR --> /install"
mount -o nfsvers=3,tcp,nolock,rw $IMAGELIBSERVER:$IMAGELIBINSTALLDIR /install

reboot=''
usepartimageng=0 # this controls what is used to save images

for o in `cat /proc/cmdline` ; do
    case $o in
    image=*)
        image=${o#image=}
        ;;
    device=*)
        device=${o#device=}
        ;;
    action=*)
        action=${o#action=}
        ;;
    blocks=*)
        blocks=${o#blocks=}
        ;;
    reboot)
        reboot=reboot
        ;;
    usepartimageng=*)
        usepartimageng=${o#usepartimageng=}
    esac
done

/install/image/x86/installer_files/listener.py &

sleep 10

devices=`\ls -l /dev/disk/by-id | awk '{print $11}'`
logger -t xcat "disks: $devices"

if [ "x$device" == "x" ] ; then
    device=$(\ls /dev/disk/by-id -l | grep -v total | grep -v usb | head -1 | awk '{print $11}' | sed 's/[\.\/]//g')
    if ! echo $device | egrep -q '(sd[a-f]|hd[a-f])'; then
        device=$(cat /proc/partitions  | egrep -v '(ram|major)' | egrep '\w' | sort -k 3 -n -r | head -1 | awk '{print $4}')
    fi
    if ! echo $device | egrep -q '(sd[a-f]|hd[a-f])'; then
        echo "No device specified and could not guess a valid device"
        logger -t xcat "partimage failed: no device specified and could not guess a valid device"
        #/xcatpost/updateflag.awk $MASTER 3002
        if [ "x$reboot" == "xreboot" ] ; then
            /sbin/reboot
        else
                while : ; do sleep 3600 ; done
        fi
        exit 1
    else
        echo "No device specified, trying with guessed device /dev/$device"
        logger -t xcat "No device specified, trying with guessed device /dev/$device"
    fi
    device="/dev/$device"
fi

if [ "x$action" == "xsave" ] ; then
    touch /tmp/save
    if [ "$usepartimageng" = "0" ]; then
		  cmd="partimage"
        partimagesave $image $device
        rc=$?
        if [ "$rc" = "0" ]; then
            echo "partimage successfully saved image"
            logger -t xcat "partimage successfully saved image"
        else
            echo "partimage failed to saved image"
            logger -t xcat "partimage failed to saved image"
        fi
    else
		  cmd="partimage-ng"
        echo running /sbin/partimage-ng -b $blocks save $device $image
        logger -t xcat "running /sbin/partimage-ng -p -b $blocks save $device $image"
        /sbin/partimage-ng -p -b $blocks save $device $image > /tmp/partimageng
        rc=$?
        if [ "$rc" = "0" ]; then
            echo gzipping $image
            logger -t xcat "gzipping $image"
            touch /tmp/savegzip
            if gzip $image; then
                echo gzipping finished, removing .gz extension from filename
                logger -t xcat "gzipping finished, removing .gz extension from filename"
                mv $image.gz $image
            else
                echo "gzipping failed"
                logger -t xcat "gzipping failed"
            fi
        fi
    fi
    if [ "$rc" = "0" ]; then
        touch $image.capturedone
        logger -t xcat "$cmd exited normally"
    else
        touch $image.capturefailed
        logger -t xcat "$cmd failed with exit code $rc"
    fi
elif [ "x$action" == "xrestore" ] ; then
    touch /tmp/restore
    testfile=$(echo $image | awk -F '.' '{print $1}')
    if \ls -1 $testfile*.mbr > /dev/null; then
        partimageinstall $image $device
        partimageinstall $image $device
        rc=$?
        if [ "$rc" = "0" ]; then
            echo "partimage successfully loaded image"
            logger -t xcat "partimage successfully loaded image"
        else
            echo "partimage failed to load image"
            logger -t xcat "partimage failed to load image"
        fi
    else
        echo "running zcat $image | /sbin/partimage-ng -p -b $blocks restore stdin $device"
        logger -t xcat "running zcat $image | /sbin/partimage-ng -p -b $blocks restore stdin $device"
        zcat $image | /sbin/partimage-ng -p -b $blocks restore stdin $device > /tmp/partimageng
        rc=$?
        if [ "$rc" = "0" ]; then
            echo "partimage-ng successfully loaded image"
            logger -t xcat "partimage-ng successfully loaded image"
            echo "running /sbin/parted $device set 1 boot on"
            logger -t xcat "running /sbin/parted $device set 1 boot on"
            /sbin/parted $device set 1 boot on
        else
            echo "partimage-ng failed to load image"
            logger -t xcat "partimage-ng failed to load image"
        fi
    fi
else
    echo "unknown action supplied"
    logger -t xcat "unknown action supplied"
    exit 1
fi

touch /tmp/done

if [ "$rc" = "0" ]; then
    /xcatpost/updateflag.awk $MASTER 3002
fi

if [ "x$reboot" == "xreboot" ] ; then
    /sbin/reboot
else
	while : ; do sleep 3600 ; done
fi