~canonical-platform-qa/ubuntu-qa-tools/practisubunit

#!/bin/sh -e
#
# Copyright (C) 2007-2011 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

# Find out our current directory
if [ -z "$UQT_VM_TOOLS" ]; then
    UQT_VM_TOOLS=$(dirname $(readlink -f $0))
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
    }

    # Wait until disk is available
    sleep 3
    for i in `seq 0 9` ; do
         if sudo fdisk -l $nbd_device | grep -q $nbd_device ; then
             break
         elif [ $i = "9" ]; then
             echo "ERROR: Disk could not be mounted! Aborting" >&2
             exit 1
         else
             echo "Waiting for disk to appear at nbd device...retry #$i..."
             sleep 3
         fi
    done

    # 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 -i 'Qcow'; then
            echo "Mounting image..."
            mount_qemu "$tmpdisk" $mountpoint
        else
            if file $d | grep -q -i '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 -l $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}" || {
                    echo "Could not delete partition mapping '$tmpimg'"
                    continue
                }

            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
}

symlink_pristine() {
    d="$1"
    doit="$2"
    if [ "$convert_to_snapshot" = "yes" ]; then
        dir=`dirname "$d"`
        bn=`basename "$d" .qcow2`
        pristine="$dir/$bn.pristine.qcow2"
        if [ -e "$pristine" ]; then
            if [ "$doit" = "no" ] && [ -e "$d" ]; then
                echo "Removing symlink to $pristine from '$d'"
                rm -f "$d"
            elif [ "$doit" = "" ] && [ ! -e "$d" ]; then
                echo "Creating symlink to $pristine from '$d'"
                ln -s "$pristine" "$d"
            else
                echo "ERROR: cannot 'symlink_pristine $d $doit'" >&2
                return 1
            fi
        fi
    fi
}

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

    # TODO: implement cloning of machines that use snapshots
    convert_to_snapshot="no"
    for d in $disks ; do
        real=`readlink -f "$d"`
        if ! echo "$real" | egrep -q '.qcow2$' ; then
            continue
        fi
        dir=`dirname "$real"`
        bn=`basename "$real" .qcow2`
        pristine="$dir/$bn.pristine.qcow2"
        if [ -e "$pristine" ]; then
            convert_to_snapshot="yes"
            if [ -e "$real" ]; then
                echo "Found existing snapshot. Either remove it or commit it"
                echo "Aborted clone of '$original' to '$clone'" >&2
                return 1
            fi
        fi
    done

    # 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`
        if [ "$convert_to_snapshot" = "yes" ]; then
            symlink_pristine "$d"
        fi
        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"
        if [ "$convert_to_snapshot" = "yes" ]; then
            for d in $disks ; do
                symlink_pristine "$d" no || true
            done
        fi
        return
    }

    # Update xml to use qcow2 if the cloned disks are qcow2
    do_convert=
    for d in $disks ; do
        tmp=`echo "$d" | cut -d ' ' -f 1 | sed "s/$original/$clone"/g`
        if file $tmp | grep -q -i 'Qcow' ; then
            do_convert="yes"
        fi
    done

    if [ "$do_convert" = "yes" ]; then
        in_disk=
        converted_line=
        unconverted_line=
        virsh -c ${vm_connect} dumpxml "$clone" 2>/dev/null | while read line ; do
            if [ -z "$in_disk" ] && echo "$line" | grep -q "<disk type='file'" ; then
                in_disk="yes"
            elif [ "$in_disk" = "yes" ] && echo "$line" | grep -q "driver name='qemu'" ; then
                unconverted_line="$line"
                if echo "$line" | grep -q "type='raw'" ; then
                    converted_line=`echo "$line" | sed 's/raw/qcow2/'`
                fi
                continue
            elif [ "$in_disk" = "yes" ] && echo "$line" | grep -q "<source file=" ; then
                converted=
                for d in $disks ; do
                    tmp=`echo "$d" | cut -d ' ' -f 1 | sed "s/$original/$clone"/g`
                    if echo "$line" | grep -q "$tmp" ; then
                        # Ok, we found a matching disk, now output the appropriate line
                        if [ -n "$converted_line" ]; then
                            echo "Converting xml to use 'qcow2' for '$tmp'"
                            echo "$converted_line" >> "$tmpdir/converted_xml"
                            converted="yes"
                            break
                        fi
                    fi
                done
                if [ -z "$converted" ]; then
                    echo "$unconverted_line" >> "$tmpdir/converted_xml"
                fi
                in_disk=
                converted_line=
                unconverted_line=
            fi
            echo "$line" >> "$tmpdir/converted_xml"
        done
        cat "$tmpdir/converted_xml"
        virsh -c ${vm_connect} define "$tmpdir/converted_xml"
        rm -f "$tmpdir/converted_xml"
    fi

    # remove the temporary symlinks
    if [ "$convert_to_snapshot" = "yes" ]; then
        for d in $disks ; do
            symlink_pristine "$d" no
        done
    fi

    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/dhcp*/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

            # Remove persistent network interface names as cloning changed our mac address
            if [ -e "$mp/etc/udev/rules.d/70-persistent-net.rules" ]; then
                echo "Updating udev..."
                sudo mv -f $mp/etc/udev/rules.d/70-persistent-net.rules $mp/etc/udev/rules.d/70-persistent-net.rules.bak || echo "Couldn't update '$mp/etc/udev/rules.d/70-persistent-net.rules'" >&2
            fi

            echo "Generating new ssh host keys..."
            sudo rm -f $mp/etc/ssh/*_key $mp/etc/ssh/*_key.pub
            # Using dpkg-reconfigure openssh-server here tries to start
            # an sshd daemon on the host, which isn't a very nice thing to
            # do, and fails if there is already one running. Manually
            # generate an rsa and a dsa key for now.
            # TODO: natty and higher should also have a ecdsa key
            sudo ssh-keygen -q -f "$mp/etc/ssh/ssh_host_rsa_key" -N '' -t rsa
            sudo ssh-keygen -q -f "$mp/etc/ssh/ssh_host_dsa_key" -N '' -t dsa

            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.
    ssh-keygen -R $clone.local

    # 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"

    if [ "$convert_to_snapshot" = "yes" ]; then
        echo "For the snapshot, please use:"
        echo "vm-use-snapshots $clone"
    fi
}

# 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