#!/bin/bash
# 
# Reads from a multi-partition flash image file and populates the target
# partition with the requested partition contents.

# Flash system image, uncompressed, e.g. oem-foo-bar.img or filesystem
# image, e.g. oem-foo-bar.fs
IMAGE=$1

# target partition, e.g. /dev/sda1
DEST_PARTITION=$2

# which partition from the sytem image to transfer
COMMAND=$3

FAT_LABEL="IHVCONFIG"  # FIXME: verify this

if [ ! -b ${DEST_PARTITION} ] ; then
    echo "Destination partition device node doesn't exist: ${DEST_PARTITION}"
    exit 1
fi

read_image_partition_offset() {
    IMG=$1
    PNUM=$2

    if [[ $PNUM == 0 ]] ; then
        # Partition 0 means that this is a filesystem image
        OFFSET=0
    else
        BYTES=$(fdisk -l ${IMG} | sed -ne 's/^Units = .* \([0-9]*\) bytes$/\1/p')
        START=$(fdisk -l ${IMG} | sed -ne 's,^'${IMG}'p\?'${PNUM}' *\*\? *\([0-9][0-9]*\) .*$,\1,p')
        if [[ $PNUM == 1 ]] ; then
            OFFSET=512
        else
            OFFSET=$(( ${BYTES} * ${START} - ${BYTES} ))
        fi
    fi
}

mount_image_partition() {
    IMG=$1
    PNUM=$2
    MPNT=$3
    # kpartx is not available on the recovery filesystem,
    # so loop mount with calculated offsets instead.
    read_image_partition_offset ${IMG} ${PNUM}

    IMAGE_LOOP_DEVICE=$(losetup -f)
    losetup ${IMAGE_LOOP_DEVICE} ${IMG} || return 1
    if ! mount -o offset=${OFFSET} $IMAGE_LOOP_DEVICE $MPNT ; then
        losetup -d ${IMAGE_LOOP_DEVICE}
        return 1
    fi
    return 0
}

mount_squashfs() {
    IMG=$1
    MPNT=$2
    SQUASHFS_LOOP_DEVICE=$(losetup -f)
    losetup ${SQUASHFS_LOOP_DEVICE} ${IMG} || return 1
    if ! mount -t squashfs $SQUASHFS_LOOP_DEVICE $MPNT ; then
        losetup -d ${SQUASHFS_LOOP_DEVICE}
        return 1
    fi
    return 0
}

write_system() {
    SYS_PNUM=$1
    RW_MODE=$2

    # Filesystem is transferred using tar, rather than dd, because we do not
    # have full control over the partition table (unlike on flash). We have
    # to preserve any existing OEM partitions, so we cannot guarantee that
    # the partitions we create have _exactly_ the right size due to differing
    # drive geometry.

    MNTPOINT_SRC="$(mktemp -d /tmp/mnt-XXXXXX)"
    if ! mount_image_partition ${IMAGE} ${SYS_PNUM} ${MNTPOINT_SRC} ; then
        echo "Failure mounting partition ${SYS_PNUM} within image: ${IMAGE}"
        exit 1
    fi

    if ! mkfs.ext3 ${DEST_PARTITION} >/dev/null 2>&1 ; then
        umount ${MNTPOINT_SRC}
        rmdir ${MNTPOINT_SRC}
        losetup -d ${IMAGE_LOOP_DEVICE}
        echo "Failure creating filesystem on destination partition: ${DEST_PARTITION}"
        exit 1
    fi
    MNTPOINT_DST="$(mktemp -d /tmp/mnt-XXXXXX)"
    if ! mount -t ext3 ${DEST_PARTITION} ${MNTPOINT_DST} ; then
        rmdir ${MNTPOINT_DST}
        umount ${MNTPOINT_SRC}
        rmdir ${MNTPOINT_SRC}
        losetup -d ${IMAGE_LOOP_DEVICE}
        echo "Failure mounting destination partition: ${DEST_PARTITION}"
        exit 1
    fi

    if [ "X${RW_MODE}" = "Xread-only" ] || [ ! -r ${MNTPOINT_SRC}/rootfs ] ; then
        # Transfer filesystem contents to target partition:
        if ! ( cd ${MNTPOINT_SRC} && tar cf - . ) | ( cd ${MNTPOINT_DST} && tar xf - . ) ; then
            umount ${MNTPOINT_DST}
            rmdir ${MNTPOINT_DST}
            umount ${MNTPOINT_SRC}
            rmdir ${MNTPOINT_SRC}
            losetup -d ${IMAGE_LOOP_DEVICE}
            echo "Error transferring system partition contents to: ${DEST_PARTITION}"
            exit 1
        fi
    else
        MNTPOINT_INNER_SRC="$(mktemp -d /tmp/mnt-XXXXXX)"
        if ! mount_squashfs ${MNTPOINT_SRC}/rootfs ${MNTPOINT_INNER_SRC} ; then
            umount ${MNTPOINT_DST}
            rmdir ${MNTPOINT_DST}
            umount ${MNTPOINT_SRC}
            rmdir ${MNTPOINT_SRC}
            losetup -d ${IMAGE_LOOP_DEVICE}
            echo "Error mounting compressed system image: ${MNTPOINT_SRC}/rootfs"
            exit 1
        fi

        # Transfer filesystem contents to target partition:
        if ! ( cd ${MNTPOINT_INNER_SRC} && tar cf - . ) | ( cd ${MNTPOINT_DST} && tar xf - . ) ; then
            umount ${MNTPOINT_INNER_SRC}
            rmdir ${MNTPOINT_INNER_SRC}
            losetup -d ${SQUASHFS_LOOP_DEVICE}
            umount ${MNTPOINT_DST}
            rmdir ${MNTPOINT_DST}
            umount ${MNTPOINT_SRC}
            rmdir ${MNTPOINT_SRC}
            losetup -d ${IMAGE_LOOP_DEVICE}
            echo "Error transferring system partition contents to: ${DEST_PARTITION}"
            exit 1
        fi

        umount ${MNTPOINT_INNER_SRC}
        rmdir ${MNTPOINT_INNER_SRC}
        losetup -d ${SQUASHFS_LOOP_DEVICE}
    fi

    # The work of the source is done
    umount ${MNTPOINT_SRC}
    rmdir ${MNTPOINT_SRC}
    losetup -d ${IMAGE_LOOP_DEVICE}

    # Reinstall the bootloader, using extlinux from the installed system image:
    INSTALLED=""
    if [ -e "${MNTPOINT_DST}/rootfs" ] ; then
        MNTPOINT_ROOTFS="$(mktemp -d /tmp/mnt-XXXXXX)"
        mount -o loop -t squashfs ${MNTPOINT_DST}/rootfs ${MNTPOINT_ROOTFS}
        mount -o bind ${MNTPOINT_DST} ${MNTPOINT_ROOTFS}/mnt
        mount -t tmpfs none ${MNTPOINT_ROOTFS}/proc
        mount -o bind /dev ${MNTPOINT_ROOTFS}/dev
        echo "${DEST_PARTITION} / ext3 rw,data=ordered 0 0" >${MNTPOINT_ROOTFS}/proc/mounts
        if chroot ${MNTPOINT_ROOTFS} /sbin/extlinux --install /mnt/boot ; then
            INSTALLED="ok"
        else
            echo "Error installing bootloader on: ${DEST_PARTITION}"
        fi
        umount ${MNTPOINT_ROOTFS}/dev
        umount ${MNTPOINT_ROOTFS}/proc
        umount ${MNTPOINT_ROOTFS}/mnt
        umount ${MNTPOINT_ROOTFS}
        rmdir ${MNTPOINT_ROOTFS}
    elif [ -e "${MNTPOINT_DST}/sbin/extlinux" ] ; then
        mount -t tmpfs none ${MNTPOINT_DST}/proc
        mount -o bind /dev ${MNTPOINT_DST}/dev
        echo "${DEST_PARTITION} / ext3 rw,data=ordered 0 0" >${MNTPOINT_DST}/proc/mounts
        if chroot ${MNTPOINT_DST} /sbin/extlinux --install /boot ; then
            INSTALLED="ok"
        else
            echo "Error installing bootloader on: ${DEST_PARTITION}"
        fi
        umount ${MNTPOINT_DST}/dev
        umount ${MNTPOINT_DST}/proc
    fi

    umount ${MNTPOINT_DST}
    rmdir ${MNTPOINT_DST}


    # generate a unique system label:
    e2label ${DEST_PARTITION} xs+$(/bin/uuidgen | /usr/bin/cut -c 1-13)

    # Aside:
    # This is how we would transfer the flash extfs label, *but*
    # we don't want to do that, as the labels may collide if the machine
    # has a flash installation, or multiple installs on different disks.
    #read_image_partition_extfs_label() {
    #    IMG=$1
    #    PNUM=$2
    #    read_image_partition_offset ${IMG} ${PNUM}
    #    # ext2 label is 16 bytes at offset 1144
    #    EXTFS_LABEL=$(/bin/dd if=${IMAGE} skip=$(( 1144 + ${OFFSET} )) bs=1 count=16 2>/dev/null)
    #}
    # read_image_partition_extfs_label ${IMAGE} ${SYS_PNUM}

    [ "x${INSTALLED}" != "x" ] || exit 1
}

write_mutable() {
    if ! mkfs.ext3 ${DEST_PARTITION} >/dev/null 2>&1 ; then
        echo "Failure creating filesystem for system state on : ${DEST_PARTITION}"
        exit 1
    fi

    STATE_LABEL="xc+"$(/bin/uuidgen | /usr/bin/cut -c 1-13)

    if ! e2label ${DEST_PARTITION} ${STATE_LABEL} >/dev/null 2>&1 ; then
        echo "Failure creating filesystem for system state on : ${DEST_PARTITION}"
        exit 1
    fi
}

write_boot() {
    SYS_PNUM=1
    FAT_PNUM=4

    MNTPOINT_SRC="$(mktemp -d /tmp/mnt-XXXXXX)"
    if ! mount_image_partition ${IMAGE} ${FAT_PNUM} ${MNTPOINT_SRC} ; then
        echo "Failure mounting partition ${FAT_PNUM} within image: ${IMAGE}"
        exit 1
    fi

    # Recovery CD doesn't have the msdos tools, so use the ones in the system image
    MNTPOINT_SYS="$(mktemp -d /tmp/mnt-XXXXXX)"
    if ! mount_image_partition ${IMAGE} ${SYS_PNUM} ${MNTPOINT_SYS} ; then
        echo "Failure mounting partition ${SYS_PNUM} within image: ${IMAGE}"
        exit 1
    fi

    MNTPOINT_ROOTFS="$(mktemp -d /tmp/mnt-XXXXXX)"
    if ! mount -o loop -t squashfs ${MNTPOINT_SYS}/rootfs ${MNTPOINT_ROOTFS} ; then
        echo "Failure mounting compressed image within image: ${IMAGE}"
        exit 1
    fi

    mount -o bind /dev/ ${MNTPOINT_ROOTFS}/dev

    if ! chroot ${MNTPOINT_ROOTFS} /sbin/mkdosfs -s 2 -F 16 -n ${FAT_LABEL} ${DEST_PARTITION} >/dev/null 2>&1 ; then
    #if ! chroot ${MNTPOINT_ROOTFS} /sbin/mkdosfs -s 2 -F 16 -n ${FAT_LABEL} ${DEST_PARTITION} ; then

        umount ${MNTPOINT_ROOTFS}/dev
        umount ${MNTPOINT_ROOTFS}
        rmdir ${MNTPOINT_ROOTFS}
        umount ${MNTPOINT_SYS}
        rmdir ${MNTPOINT_SYS}
        umount ${MNTPOINT_SRC}
        rmdir ${MNTPOINT_SRC}
        losetup -d ${IMAGE_LOOP_DEVICE}
        echo "Failure creating FAT filesystem on destination partition: ${DEST_PARTITION}"
        exit 1
    fi

    # Transfer filesystem contents to target partition:
    mount -o bind ${MNTPOINT_SRC} ${MNTPOINT_ROOTFS}/mnt
    for FILENAME in $(/bin/ls -1 ${MNTPOINT_ROOTFS}/mnt) ; do
        UPPER_FILENAME=$(/bin/echo "${FILENAME}" | /usr/bin/tr a-z A-Z)
        if ! chroot ${MNTPOINT_ROOTFS} /usr/bin/mcopy -i ${DEST_PARTITION} /mnt/${FILENAME} ::${UPPER_FILENAME} ; then
            umount ${MNTPOINT_ROOTFS}/mnt
            umount ${MNTPOINT_ROOTFS}/dev
            umount ${MNTPOINT_ROOTFS}
            rmdir ${MNTPOINT_ROOTFS}
            umount ${MNTPOINT_SYS}
            rmdir ${MNTPOINT_SYS}
            umount ${MNTPOINT_SRC}
            rmdir ${MNTPOINT_SRC}
            losetup -d ${IMAGE_LOOP_DEVICE}
            echo "Error transferring FAT partition contents to: ${DEST_PARTITION}"
            exit 1
        fi
    done

    # The bootloader config needs correcting for the partition layout
    # TODO: error handling
    if [ -e ${MNTPOINT_ROOTFS}/mnt/sysldisk.cfg ] ; then
        chroot ${MNTPOINT_ROOTFS} /usr/bin/mmove -i ${DEST_PARTITION} ::SYSLINUX.CFG ::SYSFLASH.CFG
        chroot ${MNTPOINT_ROOTFS} /usr/bin/mcopy -n -i ${DEST_PARTITION} /mnt/sysldisk.cfg ::SYSLINUX.CFG
    else
        chroot ${MNTPOINT_ROOTFS} /usr/bin/mcopy -i ${DEST_PARTITION} /mnt/syslinux.cfg ::SYSFLASH.CFG
        MNTPOINT_FAT="$(mktemp -d /tmp/mnt-XXXXXX)"
        mount -t msdos ${DEST_PARTITION} ${MNTPOINT_FAT}
        sed -i 's/current 1/current 5/' ${MNTPOINT_FAT}/syslinux.cfg
        sed -i 's/current 2/current 6/' ${MNTPOINT_FAT}/syslinux.cfg
        umount ${MNTPOINT_FAT}
        rmdir ${MNTPOINT_FAT}
    fi

    # Install the bootloader, using syslinux from the system image:
    mount -t tmpfs none ${MNTPOINT_ROOTFS}/var/tmp
    mount -o bind /proc ${MNTPOINT_ROOTFS}/proc
    if ! chroot ${MNTPOINT_ROOTFS} /usr/bin/syslinux ${DEST_PARTITION} ; then
        umount ${MNTPOINT_ROOTFS}/proc
        umount ${MNTPOINT_ROOTFS}/var/tmp
        umount ${MNTPOINT_ROOTFS}/mnt
        umount ${MNTPOINT_ROOTFS}/dev
        umount ${MNTPOINT_ROOTFS}
        rmdir ${MNTPOINT_ROOTFS}
        umount ${MNTPOINT_SYS}
        rmdir ${MNTPOINT_SYS}
        umount ${MNTPOINT_SRC}
        rmdir ${MNTPOINT_SRC}
        losetup -d ${IMAGE_LOOP_DEVICE}
        echo "Error installing syslinux bootloader on FAT partition : ${DEST_PARTITION}"
        exit 1
    fi

    umount ${MNTPOINT_ROOTFS}/proc
    umount ${MNTPOINT_ROOTFS}/var/tmp
    umount ${MNTPOINT_ROOTFS}/dev
    umount ${MNTPOINT_ROOTFS}/mnt
    umount ${MNTPOINT_ROOTFS}
    rmdir ${MNTPOINT_ROOTFS}
    umount ${MNTPOINT_SYS}
    rmdir ${MNTPOINT_SYS}

    umount ${MNTPOINT_SRC}
    rmdir ${MNTPOINT_SRC}
    losetup -d ${IMAGE_LOOP_DEVICE}
}

write_mbr() {
    if ! dd count=1 bs=410 if=${IMAGE} of=${DEST_PARTITION} ; then
        echo "Failure installing MBR on : ${DEST_PARTITION}"
        exit 1
    fi
    sync
}

case "$COMMAND" in
    system-image-0)
            write_system 0 read-only
        ;;
    system-image-0-rw)
            write_system 0 read-write
        ;;
    system-image-1)
            write_system 1 read-only
        ;;
    system-image-1-rw)
            write_system 1 read-write

        ;;
    system-image-2)
            write_system 2 read-only
        ;;
    system-image-2-rw)
            write_system 2 read-write
        ;;
    mutable-state)
            write_mutable
        ;;
    boot)
            write_boot
        ;;
    master-boot-record)
            write_mbr
        ;;
    *)
        echo "Unknown command: $COMMAND"
        exit 1
        ;;
esac

exit 0
