#!/bin/bash
# create-arista-veos-image.sh
# HS-Fulda - sebastian.rieger@informatik.hs-fulda.de
#
# changelog:
# V1.1    added injection of config defined in VM Maestro using config-drivex
# V1.1.1  fixed device mapping of extracted partitions, fixed problems with stale swi directory
# V1.1.2  rc.eos now supports e1000 and virtio as vnic types (virtio is supported in vEOS >=4.14.5F)
# V1.2    added dynamic handling of device mapping of extacted partitions
# V1.21   checking whether it safe to unmount working directories
# V1.3    added support to delete existing image with the same name and generating the default nova flavor
# V1.3.1  added support for newer glance releases (e.g. kilo) used in VIRL 1.0.0 
# V1.3.2  changed the extension of the bootloader iso to match the size of the partitions to be injected
# V1.4    support for variable VEOS image sizes (as requested by @Jade_Deane to use custom VEOS images)
# V1.4.1  fixed check for existing images
# V1.5    fixed detection of missing loop files for Ubuntu 16.04 in VIRL >=1.3, added default subtype creation
# V1.51   fixed default subtype creation to use ssh management and max 22 ports, changed ram to 2GB as recommended by Arista
# V1.52   fixed default subtype creation to support VIRL version <1.3
# V2.0    added support for new partition layout of vEOS-lab >= 4.21 images, using an integrated bootloader (as the earlier combined versions)

# usage
if [ ! $# -eq 3 ] ; then
  echo -e "usage: $0 <Aboot-veos-serial-version.iso> <vEOS-version.vmdk> <new glance image name>, e.g.:\n"
  echo "$0 Aboot-veos-serial-8.0.0.iso vEOS-lab-4.16.9M.vmdk vEOS"
  exit -1
fi

# sudo check
if [ ! $UID -eq 0 ] ; then
  echo "Insufficient privileges. Please consider using sudo -s."
  exit -1
fi

ABOOT_SERIAL_ISO=$1
ABOOT_SERIAL_ISO_BASENAME=$(basename -s .iso $1)
VEOS_VMDK=$2
VEOS_VMDK_BASENAME=$(basename -s .vmdk $2)
GLANCE_IMAGE_NAME=$3
GLANCE_IMAGE_RELEASE=$VEOS_VMDK_BASENAME-$ABOOT_SERIAL_ISO_BASENAME
TMP_NAME="vEOS-$GLANCE_IMAGE_RELEASE"
TIMESTAMP=$(date +%Y%m%d%H%M%S)

function wait_for_loop_dev_partitions() {
  LOOPDEV_PARTITION_1=$1
  LOOPDEV_PARTITION_2=$2

  echo -n "waiting for loop devs from kpartx ($LOOPDEV_PARTITION_1 and $LOOPDEV_PARTITION_2)"
  LOOP_DEV_RETRIES=0
  until dd if=/dev/mapper/$LOOPDEV_PARTITION_1 of=/dev/null bs=1k count=1 &>/dev/null && dd if=/dev/mapper/$LOOPDEV_PARTITION_2 of=/dev/null bs=1k count=1 &>/dev/null ; do
    echo -n "."
    LOOP_DEV_RETRIES=$(expr $LOOP_DEV_RETRIES + 1)
    if [ $LOOP_DEV_RETRIES -eq 10 ]; then
      echo
      echo "ERROR: timeout waiting for loop devs from kpartx"
      exit -1
    fi
    sleep 1
  done
  return 0
}

function remove_loop_devs() {
  LOOP_DEV_IMAGE=$1

  LOOP_DEV_DEL_RETRIES=0
  echo "removing loop devs from kpartx ($LOOP_DEV_IMAGE)"
  until kpartx -vd $LOOP_DEV_IMAGE ; do
    echo -n "."
    LOOP_DEV_DEL_RETRIES=$(expr $LOOP_DEV_DEL_RETRIES + 1)
    if [ $LOOP_DEV_DEL_RETRIES -eq 10 ]; then
      echo
      echo "ERROR: timeout waiting for loop dev removal from kpartx"
      exit -1
    fi
    sleep 1
  done
  echo
  return 0
}

function safe_unmount() {
  MOUNTPOINT=$1

  echo -n "Unmounting $MOUNTPOINT..."
  RETRY=0
  until umount $MOUNTPOINT &>/dev/null
  do
    echo -n "."
    sleep 1
    RETRY=$((RETRY+1))
    if [ "$RETRY" -ge "5" ] ; then
      echo
      echo "ERROR: unable to unmount working directory $MOUNTPOINT"
      exit -1
    fi
  done
  echo
  return 0
}

function inject_rc_eos() {
  echo
  echo "Injecting rc.eos startup script to get switch config..."
  echo "==========================================================="

  # inject rc.eos script to get switch config defined in VM Maestro (config-drive)
  INJECTION_TIMESTAMP=$1
  VEOS_VMDK_BASENAME_RC_EOS_PART=$2

  mkdir swi-$INJECTION_TIMESTAMP
  echo "loop mounting image..."
  mount -o loop $VEOS_VMDK_BASENAME_RC_EOS_PART swi-$INJECTION_TIMESTAMP
  cd swi-$INJECTION_TIMESTAMP
  echo "injecting rc.eos..."
  cat << EOF > rc.eos
#!/bin/sh
#
# startup script to get node configs from VM Maestro
#

echo "Getting switch config from config drive..."
echo "=========================================="
mkdir /config-drive

# check if config drive is CDROM or DISK, CDROM is required for vEOS >=4.21 combined images
blkid | grep /dev/sr0 | grep CDROM
if [ \$? -eq 0 ]; then
  # config_disk_type is cdrom
  mount /dev/sr0 /config-drive
else
  # config_disk_type is disk
  mount /dev/sdb1 /config-drive
fi

echo "Getting ip address for ma1 via dhcp..."
echo "=========================================="
MANAGEMENT_INTERFACE="ma1"
ip link show \$MANAGEMENT_INTERFACE
if [ \$? -ne 0 ]; then
  # if using virtio ma1 will not be up during Eos Init 1, hence we use eth0
  MANAGEMENT_INTERFACE="eth0"
fi
dhclient -r \$MANAGEMENT_INTERFACE
dhclient -1 -v \$MANAGEMENT_INTERFACE >/mnt/flash/dhclient.log
IP=\$(ip addr show \$MANAGEMENT_INTERFACE | grep inet | tr -s ' ' | cut -d ' ' -f 3 | sed s/"\/"/"\\\\\\\\\/"/g)
echo \$IP
sed s/"! ip of ma1 configured on launch"/"ip address \$IP"/g /config-drive/veos_config.txt >/mnt/flash/startup-config.tmp
cat /mnt/flash/startup-config.tmp
echo

echo "Copying switch config from config drive..."
echo "=========================================="
cp /mnt/flash/startup-config.tmp /mnt/flash/startup-config
EOF
  chmod 755 rc.eos
  cd ..

  echo "unmounting image..."
  safe_unmount swi-$INJECTION_TIMESTAMP

  rm -rf swi-$INJECTION_TIMESTAMP
  return 0
}

# check for an existing image with the same name and offer to delete it prior to creating a new one
CHECK_FOR_EXISTING_IMAGE=$(glance --os-image-api-version 1 image-show $GLANCE_IMAGE_NAME 2>&1)
if [ $? == 0 ] ; then
  glance --os-image-api-version 1 image-show $GLANCE_IMAGE_NAME
  echo
  echo
  read -r -p "There is already an image with the same name in glance. Do you want to overwrite it? [y/N] " RESPONSE
  if [[ $RESPONSE =~ ^([yY][eE][sS]|[yY])$ ]] ; then
    echo "Deleting existing image $GLANCE_IMAGE_NAME..."
    echo "==========================================================="

    glance --os-image-api-version 1 image-delete $GLANCE_IMAGE_NAME
  else
    echo "An image with the same name already exists. Either delete this image or choose another name."
    exit -1
  fi
fi

echo
echo "Detecting partition layout..."
echo "==========================================================="

# convert vmdk to raw
echo "Converting vmdk image to raw..."
qemu-img convert -O raw $2 $VEOS_VMDK_BASENAME.raw

echo
echo "vEOS image partition layout:"
fdisk -l $VEOS_VMDK_BASENAME.raw

PARTITION_LAYOUT=$(fdisk -l $VEOS_VMDK_BASENAME.raw -l -o device,id | tail -2 | tr -s " " | cut -d " " -f 2)
PARTITION_TYPE_PART1=$(echo "$PARTITION_LAYOUT" | sed '1q;d')
PARTITION_TYPE_PART2=$(echo "$PARTITION_LAYOUT" | sed '2q;d')

if [ "$PARTITION_TYPE_PART1" == "c" ] && [ "$PARTITION_TYPE_PART2" == "12" ]; then
  PARTITION_LAYOUT_TYPE="old"
  CONFIG_DISK_TYPE="disk"
  echo
  echo "Detected old partition layout. Bootloader partition needs to be injected in vEOS image..."

  # create a copy of Aboot bootloader and extend it to vEOS image size
  echo "Creating a copy of Aboot bootloader in $TMP_NAME.raw..."
  cp $1 $TMP_NAME.raw

  LOOPDEV=$(kpartx -av $VEOS_VMDK_BASENAME.raw)
  LOOPDEV_PART1=$(echo "$LOOPDEV" | sed '1q;d' | cut -d " " -f 3)
  LOOPDEV_PART2=$(echo "$LOOPDEV" | sed '2q;d' | cut -d " " -f 3)
  echo "Output from kpartx: $LOOPDEV"
  echo "PART1: $LOOPDEV_PART1"
  echo "PART2: $LOOPDEV_PART2"

  wait_for_loop_dev_partitions $LOOPDEV_PART1 $LOOPDEV_PART2

  echo
  echo "using dd to extract partitions..."
  dd if=/dev/mapper/$LOOPDEV_PART1 of=$VEOS_VMDK_BASENAME-p1.raw
  dd if=/dev/mapper/$LOOPDEV_PART2 of=$VEOS_VMDK_BASENAME-p2.raw

  remove_loop_devs $VEOS_VMDK_BASENAME.raw

  # inject rc.eos file to allow configuration from VM Maestro
  inject_rc_eos $TIMESTAMP $VEOS_VMDK_BASENAME-p1.raw

  echo
  echo "Injecting new partitions from vEOS vmdk in Aboot image..."
  echo "==========================================================="

  echo "$VEOS_VMDK_BASENAME.raw (original vEOS vmdk layout):"
  fdisk -l $VEOS_VMDK_BASENAME.raw

  echo "$TMP_NAME.raw (original ABoot bootloader layout):"
  fdisk -l $TMP_NAME.raw

  # calulate size of the two partitions
  PART1_START=$(fdisk -l $VEOS_VMDK_BASENAME.raw | grep "\.raw1" | tr -s " " | cut -d ' ' -f 3)
  PART1_END=$(fdisk -l $VEOS_VMDK_BASENAME.raw | grep "\.raw1" | tr -s " " | cut -d ' ' -f 4)
  PART1_LENGTH=$(expr $PART1_END - $PART1_START)

  echo "raw1 start=$PART1_START,end=$PART1_END,length=$PART1_LENGTH"

  PART2_START=$(fdisk -l $VEOS_VMDK_BASENAME.raw | grep "\.raw2" | tr -s " " | cut -d ' ' -f 2)
  PART2_END=$(fdisk -l $VEOS_VMDK_BASENAME.raw | grep "\.raw2" | tr -s " " | cut -d ' ' -f 3)
  PART2_LENGTH=$(expr $PART2_END - $PART2_START)

  echo "raw2 start=$PART2_START,end=$PART2_END,length=$PART2_LENGTH"

  # extend the bootloader iso to be able to append the two partitions
  EXTENSION_SIZE=$(ls -lk $VEOS_VMDK_BASENAME.raw | tr -s " " | cut -d " " -f 5)
  echo "extending image $TMP_NAME.raw to +$EXTENSION_SIZE"
  truncate -s +$EXTENSION_SIZE $TMP_NAME.raw

  echo "appending partitions from vmdk in the bootloader iso in $TMP_NAME.raw..."
  # append the two partitions from vmdk in the bootloader iso
  echo -e "n
p


+$PART1_LENGTH
t
2
c
a
2
n
p


+$PART2_LENGTH
t
3
12
w" | fdisk $TMP_NAME.raw

  echo "$TMP_NAME.raw (new combined image layout):"
  fdisk -l $TMP_NAME.raw

  # copy the partitions from vEOS vmdk to new image
  LOOPDEV=$(kpartx -av $TMP_NAME.raw)
  LOOPDEV_FIRST_PART=$(echo "$LOOPDEV" | sed '2q;d' | cut -d " " -f 3)
  LOOPDEV_SECOND_PART=$(echo "$LOOPDEV" | sed '3q;d' | cut -d " " -f 3)
  echo "Output from kpartx: $LOOPDEV"
  echo "PART2: $LOOPDEV_FIRST_PART"
  echo "PART3: $LOOPDEV_SECOND_PART"

  wait_for_loop_dev_partitions $LOOPDEV_FIRST_PART $LOOPDEV_SECOND_PART

  echo
  echo "copying the partitions from vEOS vmdk to new image..."
  dd if=$VEOS_VMDK_BASENAME-p1.raw of=/dev/mapper/$LOOPDEV_FIRST_PART
  dd if=$VEOS_VMDK_BASENAME-p2.raw of=/dev/mapper/$LOOPDEV_SECOND_PART

  remove_loop_devs $TMP_NAME.raw

  rm $VEOS_VMDK_BASENAME-p1.raw
  rm $VEOS_VMDK_BASENAME-p2.raw

elif [ "$PARTITION_TYPE_PART1" == "17" ] && [ "$PARTITION_TYPE_PART2" == "83" ]; then
  PARTITION_LAYOUT_TYPE="new"
  CONFIG_DISK_TYPE="cdrom"
  echo
  echo "Detected new combined partition layout. Bootloader already integrated in vEOS image ..."

  # extract bootloader partition
  LOOPDEV=$(kpartx -av $ABOOT_SERIAL_ISO)
  LOOPDEV_FIRST_PART=$(echo "$LOOPDEV" | sed '1q;d' | cut -d " " -f 3)
  echo "Output from kpartx: $LOOPDEV"
  echo "PART1: $LOOPDEV_FIRST_PART"

  # waiting for loop parts can be improved, using only one arg, better solution: get rid of kpartx completely and extract parts using dd etc.
  wait_for_loop_dev_partitions $LOOPDEV_FIRST_PART $LOOPDEV_FIRST_PART

  echo
  echo "copying the bootloader partition..."
  dd if=/dev/mapper/$LOOPDEV_FIRST_PART of=$ABOOT_SERIAL_ISO_BASENAME-p1.raw 

  remove_loop_devs $ABOOT_SERIAL_ISO

  # create a copy of vEOS image
  echo "Creating a copy of vEOS image in $TMP_NAME.raw..."
  cp $VEOS_VMDK_BASENAME.raw $TMP_NAME.raw

  # modify the content of the partitions of the image

  LOOPDEV=$(kpartx -av $TMP_NAME.raw)
  LOOPDEV_FIRST_PART=$(echo "$LOOPDEV" | sed '1q;d' | cut -d " " -f 3)
  LOOPDEV_SECOND_PART=$(echo "$LOOPDEV" | sed '2q;d' | cut -d " " -f 3)
  echo "Output from kpartx: $LOOPDEV"
  echo "PART1: $LOOPDEV_FIRST_PART"
  echo "PART2: $LOOPDEV_SECOND_PART"

  wait_for_loop_dev_partitions $LOOPDEV_FIRST_PART $LOOPDEV_SECOND_PART

  # inject rc.eos file to allow configuration from VM Maestro
  inject_rc_eos $TIMESTAMP /dev/mapper/$LOOPDEV_SECOND_PART

  # next lines currently commented out to use supplied bootloader in image

  # boot process can maybe be improved as Aboot is looking for /mnt/flash.conf and fails in new images

  #echo
  #echo "Overwrite Aboot bootloader in vEOS image..."
  #echo "==========================================================="
  #dd if=$ABOOT_SERIAL_ISO_BASENAME-p1.raw of=/dev/mapper/$LOOPDEV_FIRST_PART

  sleep 1

  remove_loop_devs $TMP_NAME.raw

  rm $ABOOT_SERIAL_ISO_BASENAME-p1.raw 

else

  echo
  echo "ERROR: Provided vEOS image uses unknown partition layout. Unable to extract and modify partitions."
  exit -1

fi

echo
echo "Convert new image to qcow2..."
echo "==========================================================="

# convert raw to qcow2
qemu-img convert -O qcow2 $TMP_NAME.raw $TMP_NAME.qcow2

echo
echo "Cleaning up..."
echo "==========================================================="

#cleanup
rm $TMP_NAME.raw
rm $VEOS_VMDK_BASENAME.raw

echo
echo "Importing image into glance..."
echo "==========================================================="
glance image-create --container-format bare --disk-format qcow2 --visibility public --name $GLANCE_IMAGE_NAME \
  --file $TMP_NAME.qcow2 --property hw_disk_bus=ide --property serial=1 \
  --property hw_vif_model=e1000 --property hw_cdrom_type=ide --property release="$GLANCE_IMAGE_RELEASE" --property subtype=$GLANCE_IMAGE_NAME --property config_disk_type=$CONFIG_DISK_TYPE

# create default flavor
CHECKING_FOR_EXISTING_FLAVOR=$(nova flavor-show $GLANCE_IMAGE_NAME.tiny 2>&1)
if [ $? == 1 ]; then
  echo "Creating default flavor $GLANCE_IMAGE_NAME.tiny..."
  echo "==========================================================="

  nova flavor-create --is-public true $GLANCE_IMAGE_NAME.tiny auto 1024 0 1
else
  echo "Default flavor $GLANCE_IMAGE_NAME.tiny already exists..."
fi
CHECKING_FOR_EXISTING_FLAVOR=$(nova flavor-show $GLANCE_IMAGE_NAME.small 2>&1)
if [ $? == 1 ]; then
  echo "Creating default flavor $GLANCE_IMAGE_NAME.small..."
  echo "==========================================================="

  nova flavor-create --is-public true $GLANCE_IMAGE_NAME.small auto 1536 0 1
else
  echo "Default flavor $GLANCE_IMAGE_NAME.small already exists..."
fi
CHECKING_FOR_EXISTING_FLAVOR=$(nova flavor-show $GLANCE_IMAGE_NAME.medium 2>&1)
if [ $? == 1 ]; then
  echo "Creating default flavor $GLANCE_IMAGE_NAME.medium..."
  echo "==========================================================="

  nova flavor-create --is-public true $GLANCE_IMAGE_NAME.medium auto 2048 0 1
else
  echo "Default flavor $GLANCE_IMAGE_NAME.medium already exists..."
fi

# create default subtype
cat << EOF > dynamic-subtype-$GLANCE_IMAGE_NAME.json.default.virl-above-1.3
{
  "dynamic-subtypes": [
    {
      "plugin_base": "generic",
      "device_type": "switch",
      "plugin_desc": "Arista vEOS",
      "cli_protocol": "ssh",
      "plugin_name": "$GLANCE_IMAGE_NAME",
      "cli_serial": 1,
      "interface_range": 22,
      "gui_visible": true,
      "interface_pattern": "Ethernet{0}",
      "baseline_image": "$GLANCE_IMAGE_NAME",
      "config_disk_type": "$CONFIG_DISK_TYPE",
      "hw_ram": 2048,
      "hw_vm_extra": "",
      "gui_icon": "iosvl2",
      "config_file": "/veos_config.txt",
      "interface_first": 1,
      "deprecated_use": "",
      "baseline_flavor": "$GLANCE_IMAGE_NAME.medium",
      "interface_management": "Management1"
    }
  ]
}
EOF
cat << EOF > dynamic-subtype-$GLANCE_IMAGE_NAME.json.default.virl-below-1.3
{
  "dynamic-subtypes": [
    {
      "plugin_base": "generic",
      "plugin_desc": "Arista vEOS",
      "cli_protocol": "ssh",
      "plugin_name": "$GLANCE_IMAGE_NAME",
      "cli_serial": 1,
      "interface_range": 22,
      "gui_visible": true,
      "interface_pattern": "Ethernet{0}",
      "baseline_image": "$GLANCE_IMAGE_NAME",
      "config_disk_type": "$CONFIG_DISK_TYPE",
      "hw_ram": 2048,
      "hw_vm_extra": "",
      "gui_icon": "iosvl2",
      "config_file": "/veos_config.txt",
      "interface_first": 1,
      "deprecated_use": "",
      "baseline_flavor": "$GLANCE_IMAGE_NAME.medium",
      "interface_management": "Management1"
    }
  ]
}
EOF
CHECKING_FOR_EXISTING_SUBTYPE=$(virl_uwm_client subtype-info --name $GLANCE_IMAGE_NAME 2>&1)
if [ $? == 255 ]; then
  echo "Creating default subtype $GLANCE_IMAGE_NAME..."
  echo "==========================================================="

  CHECKING_FOR_EXISTING_DEVICE_TYPE=$(virl_uwm_client subtype-info 2>&1 | grep u\'device_type\':)
  if [ $? == 1 ]; then
    # device_type attribute is not available in VIRL < 1.3, use JSON definition of subtype without this attribute
    echo "detected VIRL version < 1.3.0, selecting appropriate subtype to import"
    virl_uwm_client subtype-import --dynamic-subtypes @dynamic-subtype-$GLANCE_IMAGE_NAME.json.default.virl-below-1.3
  else
    echo "detected VIRL version > 1.3.0, selecting appropriate subtype to import"
    # device_type attribute is available, use JSON definition of subtype including this attribute
    virl_uwm_client subtype-import --dynamic-subtypes @dynamic-subtype-$GLANCE_IMAGE_NAME.json.default.virl-above-1.3
  fi
else
  echo "Default subtype $GLANCE_IMAGE_NAME already exists..."
fi
rm dynamic-subtype-$GLANCE_IMAGE_NAME.json.default.virl-above-1.3
rm dynamic-subtype-$GLANCE_IMAGE_NAME.json.default.virl-below-1.3

echo
echo "Image creation successful."
echo "==========================================================="
echo
echo "You can import and use the subtype $GLANCE_IMAGE_NAME in VM Maestro..."

#testing:
#
#  nova boot --image "Arista vEOS Disk" --flavor m1.small veos --nic net-id=abc7ad47-55fd-4396-8d31-91dd4d41a18a --nic net-id=abc7ad47-55fd-4396-8d31-91dd4d41a18a --nic net-id=abc7ad47-55fd-4396-8d31-91dd4d41a18a --nic net-id=abc7ad47-55fd-4396-8d31-91dd4d41a18a --nic net-id=abc7ad47-55fd-4396-8d31-91dd4d41a18a
#
#  using VM Maestro, the image can be chosen as "VM image", e.g., for an IOSv or IOSvL2 node