#!/bin/sh
# Copyright (c) 2016, Citrix Systems, Inc.
#
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
# 
#  1. Redistributions of source code must retain the above copyright
#     notice, this list of conditions and the following disclaimer.
#  2. Redistributions in binary form must reproduce the above copyright
#     notice, this list of conditions and the following disclaimer in the
#     documentation and/or other materials provided with the distribution.
#  3. Neither the name of the copyright holder nor the names of its 
#     contributors may be used to endorse or promote products derived from 
#     this software without specific prior written permission.
# 
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
# "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
# A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER
# OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
# PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
# LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING
# NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
# SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

set -e

PARTUTIL=/usr/sbin/part-util
LIBVHDIO=/usr/lib/libvhdio.so.1.0

die()
{
    echo "$@"
    exit 1
}

usage()
{
    echo "usage: $0 [-a | -d | -l] vhd [lib]"
    echo "-a add partition mappings"
    echo "-d del partition mappings"
    echo "-l list partition mappings"
    exit 1
}

parse_args()
{
    part_util=$PARTUTIL

    while [ $# -ge 1 ]; do
	case $1 in
	    -a) add="TRUE" && count="1$count";;
	    -d) del="TRUE" && count="1$count";;
	    -l) list="TRUE" && count="1$count";;
	    *) if [ -z "$vhd" ]; then vhd=$1;
	       elif [ -z "$lib" ]; then lib=$1;
	       else usage;
	       fi;;
	esac
	shift
    done

    [[ -z "$lib" ]] && lib=$LIBVHDIO
    [[ -z "$vhd" || "$count" != "1" ]] && usage
    return 0
}

# screen-scraping of fdisk... not used
fdisk_read_partitions()
{
    local data=$(LD_PRELOAD=$lib fdisk -l $vhd)

    local none=$(echo $data | grep "This doesn't look like a partition table")
    [[ -n "$none" ]] && partitions=0 && return 0

    partitions=4
    while [[ "$partitions" != "0" ]]; do
	local hit=$(echo $data | grep "${vhd}$partitions")
	[[ -n "$hit" ]] && break
	let partitions=$partitions-1
    done
}

part_util_read_partitions()
{
    partitions=$(LD_PRELOAD=$lib $part_util -c -i $vhd)
}

list_mappings()
{
    local parts=1
    while [[ $parts -le $partitions ]]; do
	echo ${vhd}$parts
	let parts=$parts+1
    done
}

add_mappings()
{
    local parts=1
    local path=$(realpath $vhd)
    while [[ $parts -le $partitions ]]; do
	[[ -e ${path}${parts} ]] || ln -s $(basename $path) ${path}$parts
	let parts=$parts+1
    done
}

del_mappings()
{
    local parts=1
    while [[ $parts -le $partitions ]]; do
	[[ -L ${vhd}$parts ]] && rm -f ${vhd}$parts
	let parts=$parts+1
    done
}

main()
{
    parse_args $@
    [[ -x $part_util ]] || die "can't find part-util"
    [[ -r $vhd && -r $lib ]] || die "can't find vhd or lib"

    part_util_read_partitions

    [[ -n "$add" ]] && add_mappings
    [[ -n "$del" ]] && del_mappings
    [[ -n "$list" ]] && list_mappings

    return 0
}

main $@
