#!/bin/bash

CONFIGS=/var/cache/frankendisk
VERBOSE=0

# Default disk geometry
# FIXME: Vary this suitably for given disks
HEADS=255
SECTORS=63
CYLINDER_SIZE=$(( ${HEADS} * ${SECTORS} ))

# Exit codes:
# 1	Already exists
# 2	Failed to set up loop device
# 3	File wasn't regular file or block device
# 4	Failed to read block device size
# 5	Not a valid frankendisk configuration

######################
# Support and logging

function help
{
    cat <<EOF
Usage: frankendisk -C [options] <name> <device|file> [<device|file> ...]
       frankendisk -A|-D|-X|-E [options] <name>
       frankendisk -L [options]
Options:
        -C, --create	        Create new disk from devices
        -A, --assemble          Assemble previously-created disk
        -D, --dismantle         Dismantle previously-assembled disk
        -X, --destroy           Forget all configuration about disk
	-L, --list		Show all configured disks
	-E, --examine		Show details of disk
	-v, --verbose		Increase verbosity level
EOF
    exit 0
}

function vlog
{
    if [ ${VERBOSE} -ge $1 ]; then
	shift
	echo = "$@"
    fi
}

function test_config
{
    NAME=$1

    vlog 3 Checking for valid config: directory
    if [ ! -d ${CONFIGS}/${NAME} ]; then
	return 1
    fi

    vlog 3 Checking for valid config: devices file
    if [ ! -f ${CONFIGS}/${NAME}/devices ]; then
	return 2
    fi

    return 0
}

function good_config
{
    NAME=$1

    test_config ${NAME}
    RV=$?

    case RV in
	0)
	    return 0
	    ;;
	1)
	    echo >&2 Configuration ${NAME} not known
	    exit 5
	    ;;
	2)
	    echo >&2 Configuration ${NAME} is not valid
	    exit 5
	    ;;
    esac
}

function get_block_count
{
    # FIXME: probably breaks on block devices larger than 2TiB
    DEVICE=$1
    if ! blockdev --getsz ${DEVICE}; then
	echo >&2 Error reading block size for ${DEVICE}
	echo 0
    fi
}

function write_dev_size
{
    NAME=$1
    TYPE=$2
    DEVICE=$3

    SIZE=$(get_block_count ${DEVICE})
    vlog 2 Device ${DEVICE} is ${SIZE} blocks, type ${TYPE}
    if [ ${SIZE} -eq 0 ]; then
	return 4
    fi
    echo >>${CONFIGS}/${NAME}/sizes ${TYPE} ${DEVICE} ${SIZE}

    # Compute the amount of padding needed
    PADDING=$(( ${SIZE} % ${CYLINDER_SIZE} ))
    PADDING=$(( ${CYLINDER_SIZE} - ${PADDING} ))
    if [ ${PADDING} -eq ${CYLINDER_SIZE} ]; then PADDING=0; fi
    if [ ${PADDING} -gt 0 ]; then
	echo >>${CONFIGS}/${NAME}/sizes P padding ${PADDING}
    fi
}

function canonicalise_path
{
    NAME=$1
    if [ ${NAME:0:1} != / ]; then
        # If the filename doesn't start with a /
	# then it's relative, and we should prepend the PWD
	echo ${PWD}/${NAME}
    else
	# If the filename does start with a /
	# then it's absolute, and we just use it verbatim
	echo ${NAME}
    fi
}

#############################
# Functions to tidy up stuff

function clean_config
{
    NAME=$1
    rm -r ${CONFIGS}/${NAME}
}

function unloop
{
    NAME=$1

    if [ -f ${CONFIGS}/${NAME}/sizes ]; then
	while read TYPE DEVICE SIZE; do
	    if [ ${TYPE} = L ]; then
		vlog 2 Unconfiguring loop device ${DEVICE}
		losetup -d ${DEVICE}
	    fi
	done < ${CONFIGS}/${NAME}/sizes
    fi
}

#################
# Main functions

function create_disk
{
    NAME=$1
    shift

    vlog 1 Creating new disk ${NAME}

    # Check for duplicate name
    if [ -e ${CONFIGS}/${NAME} ]; then
	echo >&2 Frankendisk ${NAME} already exists
	exit 1
    fi

    # Make directory & start config files
    mkdir -p ${CONFIGS}/${NAME}
    touch ${CONFIGS}/${NAME}/devices

    # Check sanity
    vlog 1 Checking devices exist and are sane
    for DEVICE in "$@"; do
	DEVICE=$(canonicalise_path ${DEVICE})
	vlog 2 Device at ${DEVICE}
	if [ ! -f ${DEVICE} -a ! -b ${DEVICE} ]; then
	    echo >&2 ${DEVICE} is neither a block device nor a file
	    clean_config ${NAME}
	    exit 3
	fi
	# Add to the list of devices
	echo ${DEVICE} >>${CONFIGS}/${NAME}/devices
    done

    # Create boot sector
    vlog 1 Creating boot sector
    dd if=/dev/zero of=${CONFIGS}/${NAME}/bs-m count=1 bs=512

    return 0
}

function assemble_disk
{
    NAME=$1

    vlog 1 Assembling new disk ${NAME}

    good_config ${NAME}

    # Read control file
    # Read old partition table for flags

    # Clean up any floating junk
    rm -f ${CONFIGS}/${NAME}/sizes

    # Create the boot sector(s)
    BOOTSECTOR=$(losetup -f -s ${CONFIGS}/${NAME}/bs-m)
    if ! write_dev_size ${NAME} b ${BOOTSECTOR}; then
	echo >&2 Failed to set up boot sector for disk ${NAME}
	exit 4
    fi
    vlog 1 Master boot sector device created as ${BOOTSECTOR}

    # For each device:
    while read DEVICE; do
	DEVTYPE=B
        # Convert files to loop devices
	if [ -f ${DEVICE} ]; then
	    vlog 1 ${DEVICE} is a file: setting up loop device
	    LOOPDEV=$(losetup -f)
	    LOOPDEV_LIST="${LOOPDEVLIST} ${LOOPDEV}"
	    vlog 1 "  ${DEVICE} = ${LOOPDEV}"
	    losetup ${LOOPDEV} ${DEVICE}
	    if [ $? -ne 0 ]; then
		echo >&2 Failure setting up loop device ${LOOPDEV} for source file ${DEVICE}
		unloop ${NAME}
		clean_config ${NAME}
		exit 2
	    fi
	    ORIGDEVICE=${DEVICE}
	    DEVICE=${LOOPDEV}
	    DEVTYPE=L
	fi

        # Get length of device & store
	if ! write_dev_size ${NAME} ${DEVTYPE} ${DEVICE}; then
	    AS=
	    if [ ${ORIGDEVICE} != ${DEVICE} ]; then
		AS=" (as ${DEVICE})"
	    fi
	    echo >&2 Failed to set up device ${ORIGDEVICE}${AS} for disk ${NAME}
	    exit 4
	fi
    done < ${CONFIGS}/${NAME}/devices

    # Write dmtable and partition table config
    vlog 1 Constructing configuration files

    DMTABLE=${CONFIGS}/${NAME}/dmtable
    rm -f ${DMTABLE}

    PARTTABLE=${CONFIGS}/${NAME}/parttable
    rm -f ${PARTTABLE}

    POSITION=0
    while read TYPE DEVICE SIZE; do
	case ${TYPE} in
	    P)
		# Padding
		DLINE="${POSITION} ${SIZE} ${PADTYPE}"
		PLINE=
		;;
	    b)
		# Boot sector
		DLINE="${POSITION} ${SIZE} linear ${DEVICE} 0"
		PLINE=
		PADTYPE=zero
		;;
	    *)
		# Loop or block device
		DLINE="${POSITION} ${SIZE} linear ${DEVICE} 0"
		PLINE="${POSITION},${SIZE},83,-"
		PADTYPE=error
		;;
	esac
	echo ${DLINE} >>${DMTABLE}
	vlog 2 dmtable: ${DLINE}

	if [ ${PLINE} ]; then
	    vlog 2 parttable: ${PLINE}
	    echo ${PLINE} >>${PARTTABLE}
	fi
	POSITION=$(( ${POSITION} + ${SIZE} ))
    done < ${CONFIGS}/${NAME}/sizes

    vlog 1 Setting up dm device
    dmsetup create ${NAME} < ${CONFIGS}/${NAME}/dmtable

    # Write partition table to the boot sector
    vlog 1 Writing partition table
    sfdisk /dev/mapper/${NAME} -uS --force < ${PARTTABLE}

    return 0
}

function dismantle_disk
{
    NAME=$1
    vlog 1 Dismantling disk ${NAME}
    good_config ${NAME}

    # Call dmsetup to stop device
    vlog 1 Stopping the DM device
    dmsetup remove /dev/mapper/${NAME}
    # Clear the loop devices we were using
    vlog 1 Unmounting loop devices
    unloop ${NAME}
    rm -f ${CONFIGS}/${NAME}/sizes
    rm -f ${CONFIGS}/${NAME}/dmtable
    rm -f ${CONFIGS}/${NAME}/parttable

    return 0
}

function destroy_disk
{
    NAME=$1
    good_config ${NAME}

    # Do dismantle disk
    dismantle_disk ${NAME}
    # Delete config files
    clean_config ${NAME}
}

function list_disks
{
    # For every disk, make sure it's a good configuration, and print
    # stuff about it
    for DIR in ${CONFIGS}/*; do
	if [ -d ${DIR} ]; then
	    NAME=$(basename ${DIR})
	    if test_config ${NAME}; then
		ACTIVE=inactive
		if [ -f ${CONFIGS}/${NAME}/sizes \
		    -a -b /dev/mapper/${NAME} ]; then
		    ACTIVE=active
		fi
		echo ${NAME} ${ACTIVE}
	    fi
	fi
    done
}

function examine_disk
{
    # FIXME: Show details of configuration in inactive disks, 
    NAME=$1
    good_config ${NAME}

    if [ ! -f ${CONFIGS}/${NAME}/sizes \
	-o ! -b /dev/mapper/${NAME} ]; then
	echo Not active
	return
    fi

    FIRST=1

    while read -u3 TYPE REALDEVICE SIZE; do
	if [ ${FIRST} -eq 1 ]; then
	    DEVICE="[boot]"
	    FIRST=0
	else
	    read DEVICE
	fi
	if [ ${DEVICE} = ${REALDEVICE} ]; then
	    echo ${DEVICE} ${SIZE}
	else
	    echo ${DEVICE}=${REALDEVICE} ${SIZE}
	fi
    done < ${CONFIGS}/${NAME}/devices 3< ${CONFIGS}/${NAME}/sizes
}

# Config structure:
#
#/var/cache/frankendisk/$NAME
#  devices - list of files/block devices to use, in order
#  sizes - computed length of file sizes (in bytes) -- for block devices
#  bs-m - master boot sector
#  bs-e0, bs-e1, bs-e2, ... - Extended partition sectors

######################

# Find and read config file

OPERATION=HELP

while [ $# -gt 0 ]; do
    case $1 in
	-h|--help)
	    help
	    ;;

	-C|--create)
	    OPERATION=CREATE
	    ;;

	-A|--assemble)
	    OPERATION=ASSEMBLE
	    ;;

	-D|--dismantle)
	    OPERATION=DISMANTLE
	    ;;

	-X|--destroy)
	    OPERATION=DESTROY
	    ;;

	-L|--list)
	    OPERATION=LIST
	    ;;

	-E|--examine)
	    OPERATION=EXAMINE
	    ;;

	-v|--verbose)
	    VERBOSE=$((${VERBOSE}+1))
	    ;;

	*)
	    break 2
	    ;;
    esac
    shift
done

NAME=$1
shift

case ${OPERATION} in
    CREATE)
	create_disk ${NAME} "$@"
	;;

    ASSEMBLE)
	if [ $# -gt 1 ]; then help; fi
	assemble_disk ${NAME}
	;;

    DISMANTLE)
	if [ $# -gt 1 ]; then help; fi
	dismantle_disk ${NAME}
	;;

    DESTROY)
	if [ $# -gt 1 ]; then help; fi
	destroy_disk ${NAME}
	;;

    EXAMINE)
	if [ $# -gt 1 ]; then help; fi
	examine_disk ${NAME}
	;;

    LIST)
	list_disks
	;;

    *)
	help
	;;
esac

exit 0
