#!/bin/sh -e
#
# Copyright (C) 2007-2009 Canonical, Ltd.
# Author: Jamie Strandboge <jamie@canonical.com>
# License: GPLv3
#
# Tested on Jaunty
#
# apt-get install kpartx nbd-client
#

ustconf="$HOME/.uqt-vm-tools.conf"
if [ -s "$ustconf" ]; then
    . "$ustconf"
else
    echo "Could not find '$ustconf'"
    exit 1
fi

. $UQT_VM_TOOLS/libvm.sh
abort_if_root

original=""
clone=""
script=""
ssh_opt="-o StrictHostKeyChecking=no"
convert=""
nbd_device=""
update_errors=""

help() {
    cat << EOM
USAGE:
  vm-clone <original> <new> [<script>]
  vm-clone -p PREFIX [-c] [-a ARCH] NEWPREFIX [<script>]

Specifying a 'PREFIX' will clone all machines created with vm-new with that
PREFIX that are in '$vm_release_list'. Eg:
$ vm-clone -p clean sec

will clone sec-<release>-<arch> VMs for each clean-<release>-<arch>.

Specifying '-c' will convert a qcow to raw instead of using nbd.

WARNING: vm-clone does NOT check to make sure you have enough space on disk,
so make sure you do.
EOM
}

tmpdir=`mktemp -d`
trap "rmdir $tmpdir" EXIT HUP INT QUIT TERM

mount_qemu() {
    max=16
    count=

    for i in `seq 0 $max` ; do
        if ! nbd-client -c /dev/nbd$i ; then
            count=$i
            break
        fi
    done

    if [ -z "$count" ]; then
        echo "Could not find available nbd device. Aborting" >&2
        exit 1
    fi
    nbd_device=/dev/nbd$count

    disk="$1"
    mountpoint="$2"
    if [ -e "/sys/module/nbd/parameters/max_part" ]; then
        cur_max_part=`cat /sys/module/nbd/parameters/max_part`
        if [ "$cur_max_part" = "0" ]; then
            sudo rmmod nbd || {
                cat << EOM
/sys/module/nbd/parameters/max_part is $cur_max_part. This setting should be
$max or higher. I tried to unload the module, but could not. Please
release any nbd devices with:
$ for i in `seq 0 $max`; do sudo nbd-client -c /dev/nbd\$i && sudo nbd-client -d /dev/nbd\$i ; done

Then unload the nbd module and reload it:
$ sudo rmmod nbd
$ sudo modprobe nbd max_part=$max

Aborting
EOM
                exit 1
            }
        fi
    fi

    lsmod | grep -q '^nbd ' || {
        echo "Inserting nbd module" >&2
        sudo modprobe nbd max_part=$max
    }

    sudo qemu-nbd -n --connect=$nbd_device "$disk" || {
        echo "Could not mount '$disk' (qemu-nbd failed). Skipping" >&2
        return
    }
    sleep 2

    # mount the non-swap partitions
    sudo fdisk -l $nbd_device | while read line ; do
        if ! echo "$line" | grep -q "83  Linux" || echo "$line" | grep -q "swap" ; then
            #echo "Not mounting '$line'" >&2
            continue
        fi
        partition=`echo "$line" | cut -d ' ' -f 1`
        device=`echo $partition | cut -d '/' -f 3`
        mp="$mountpoint/$device"
        mkdir -p "$mp" || {
            echo "Could not create '$mp'" >&2
            continue
        }
        sudo mount $partition $mp
    done

}


mount_raw() {
    disk="$1"
    mountpoint="$2"

    loopdev=`sudo kpartx -av "$disk" | head -1 | cut -d ' ' -f 8 | cut -d '/' -f 3`
    sleep 2

    # mount the non-swap partitions
    count=0
    sudo fdisk -l "$disk" 2>/dev/null | while read line ; do
        echo "$line" | grep -q "8[23]  Linux" && count=$((count+1))
        if ! echo "$line" | grep -q "83  Linux" || echo "$line" | grep -q "swap" ; then
            #echo "Not mounting '$line'"
            continue
        fi
        partition="/dev/mapper/${loopdev}p$count"
        device=`basename $partition`
        mp="$mountpoint/$device"
        mkdir -p "$mp" || {
            echo "Could not create '$mp'"
            continue
        }
        sudo mount $partition $mp
    done

}

mount_disks() {
    for d in "$@" ; do
        if [ "$d" = "--file" ]; then
            continue
        fi

        mountpoint=$tmpdir/`basename $d`
        tmpdisk="$d"

        if [ -z "$convert" ] && file $d | grep -q 'Qcow'; then
            echo "Mounting image..."
            mount_qemu "$tmpdisk" $mountpoint
        else
            if file $d | grep -q 'Qcow' ; then
                # convert to raw
                echo "Converting to raw..."
                tmpdisk="${d}.converted_to_raw"
                qemu-img convert "$d" -O raw "$tmpdisk"
            fi

            echo "Mounting image..."
            mount_raw "$tmpdisk" $mountpoint
        fi
    done
}

umount_disks() {
    tmp_args="$@"
    for i in `ls -1 $tmpdir` ; do
        diskdir="$tmpdir/$i"
        if [ ! -d "$diskdir" ]; then
            continue
        fi
        for j in `ls -1 $diskdir` ; do
            mp="$diskdir/$j"
            if [ -d "$mp" ]; then
                echo "Unmounting image..."
                sudo umount $mp && rmdir $mp || {
                    echo "Could not umount '$tmpdir/$i'"
                    continue
                }
            fi
        done
        rmdir $diskdir || {
            echo "Could not remove '$diskdir'"
            continue
        }
    done

    for d in $tmp_args ; do
        if [ "$d" = "--file" ]; then
            continue
        fi
        tmpimg="${d}"

        if [ ! -z "$nbd_device" ]; then
            sudo qemu-nbd --disconnect "${tmpimg}"
            sudo nbd-client -d $nbd_device || {
                echo "ERROR: Could not disconnect $nbd_device! Aborting" >&2
                exit 1
            }
        else
            if [ -s "${d}.converted_to_raw" ]; then
                tmpimg="${d}.converted_to_raw"
            fi
            sudo kpartx -dv "${tmpimg}"

            if [ -s "${d}.converted_to_raw" ]; then
                # convert back to qcow
                echo "Converting to qcow2..."
                qemu-img convert -f raw "${d}.converted_to_raw" -O qcow2 "${d}"
                rm -f "${d}.converted_to_raw"
            fi
        fi
    done
}

arch=
prefix=
while getopts "ca:p:" opt
do
    case "$opt" in
        a) arch="$OPTARG";;
        c) convert="yes";;
        p) prefix="$OPTARG";;
        ?) help
           exit 1
           ;;
    esac
done
shift $(($OPTIND - 1))

which qemu-img >/dev/null || {
    echo "Could not find qemu-img. Aborting" >&2
    exit 1
}

if [ -z "$1" ]; then
    help
    exit 1
elif [ -z "$prefix" ]; then
    original="$1"
    shift
else
    original="$prefix"
fi

if [ -z "$1" ]; then
    help
    exit 1
else
    clone="$1"
    shift
fi

if [ ! -z "$1" ]; then
    script="$1"
    shift
    if [ ! -x "$script" ]; then
        echo "Specified script '$script' is not executable"
        exit 1
    fi
fi

update_machine() {
    local original="$1"
    local clone="$2"
    local release="$3"
    # check if original and clone exists
    if ! vm_exists $original ; then
        echo "'$original' does not exist. Skipping" >&2
        return
    fi

    if vm_exists $clone ; then
        echo "'$clone' already exists. Skipping" >&2
        return
    fi

    disks=`get_disks $original`
    if [ -z "$disks" ]; then
        echo "Could not find any disks for '$original'. Skipping" >&2
        return
    fi

    # just grab the first disk for the path
    tmp=`echo "$disks" | awk '{ print $1 }' | sed "s/$original/$clone"/g`
    newpath=`dirname "$tmp"`
    mkdir "$newpath" || return

    # compose file args
    file_args=""
    for d in $disks ; do
        tmp=`echo "$d" | cut -d ' ' -f 1 | sed "s/$original/$clone"/g`
        file_args="$file_args --file $tmp"
    done

    if [ -z "$file_args" ]; then
        echo "Could not create '--file' arguments. Skipping" >&2
        return
    fi

    virt-clone --connect ${vm_connect} --original "$original" --name "$clone" $file_args || {
        rm -rf "$newpath"
        return
    }

    mount_disks $file_args

    prevdir=`pwd`
    for i in `ls -1 $tmpdir` ; do
        diskdir="$tmpdir/$i"
        if [ ! -d "$diskdir" ]; then
            continue
        fi
        for j in `ls -1 $diskdir` ; do
            mp="$diskdir/$j"
            if [ ! -d "$mp/etc" ]; then
                continue
            fi
            echo "Updating files..."
            sudo sh -c "echo $clone > $mp/etc/hostname" || echo "Couldn't update '$mp/etc/hostname'" >&2
            sudo sed -i "s/$original/$clone/g" $mp/etc/dhcp3/dhclient.conf || echo "Couldn't update '$mp/etc/dhcp3/dhclient.conf'" >&2
            sudo sed -i "s/$original/$clone/g" $mp/etc/hosts || echo "Couldn't update '$mp/etc/hosts'" >&2

            if [ "$release" = "gutsy" ] || [ "$release" = "hardy" ]; then
                echo "Updating udev..."
                mac=`virsh dumpxml $clone 2>/dev/null | grep 'mac address' | cut -d "'" -f 2`
                sudo sh -c "echo SUBSYSTEM==\"net\", DRIVERS==\"?*\", ATTRS{address}==\"$mac\", NAME=\"eth0\" > $mp/etc/udev/rules.d/70-persistent-net.rules"
            fi

            if [ ! -z "$script" ]; then
                sudo cp $script $mp/$script
            fi
        done
    done

    umount_disks $file_args

    # Get rid of old ssh key as it may be wrong
    ssh-keygen -R $clone
    ssh-keygen -R $clone.

    # start it up and try to login
    virsh --connect ${vm_connect} start "$clone"
    clone_hostname=`vm_wait $clone`
    if [ -z "$clone_hostname" ]; then
        update_errors="${update_errors}'$clone' is not pingable'\n"
    elif ssh $ssh_opt -t root@${clone_hostname} cat /etc/hostname | grep -q "$clone"; then
        echo "/etc/hostname updated"
        if [ ! -z "$script" ]; then
            if ssh $ssh_opt -t root@${clone_hostname} /${script} ; then
                echo "Script successfully executed"
            else
                echo "Script not successful"
                update_errors="${update_errors}'$script' on '$clone' not successful\n"
            fi
        fi
    else
        update_errors="${update_errors}Could not ssh into '$clone'\n"
    fi

    virsh --connect ${vm_connect} shutdown "$clone"
}

# check that we have all the needed binaries
error=
for b in qemu-nbd nbd-client kpartx ; do
    which $b >/dev/null || {
        echo "Could not find '$b'" >&2
        error="yes"
    }
done
if [ "$error" = "yes" ]; then
    echo "Required binaries not found. Aborting" >&2
    exit 1
fi

ionice -c 2 -n 7 -p $$
if [ -z "$prefix" ]; then
    update_machine $original $clone
else
    archs=$vm_archs
    if [ ! -z "$arch" ]; then
        archs="$arch"
    fi

    for release in $vm_release_list
    do
        for arch in $archs ; do
            echo "Cloning ${original}-${release}-${arch} to ${clone}-${release}-${arch}"
            update_machine ${original}-${release}-${arch} ${clone}-${release}-${arch} $release
        done
    done
fi

if [ -z "$update_errors" ]; then
    echo "SUCCESS"
else
    echo "Completed with errors:"
    echo -e "$update_errors"
fi

exit
