#!/bin/bash
#===============================================================================
# Copyright (c) 2012-2015 Wind River Systems, Inc.
# The right to copy, distribute, modify, or otherwise
# make use of this software may be licensed only pursuant
# to the terms of an applicable Wind River license agreement.
#
# Junxian.Xiao@windriver.com
#===============================================================================
DEVICE=""
CURRENT_VG=""
EXTEND_VG=0
EXTEND_LV=0
DEBUG=0

DEFAULT_LV_SIZE="45%VG"

#===============================================================================
# Functions
#===============================================================================
SHOW() {
    local level=$1 c p e="\n"

    case $level in
        ERROR)   c=31; p="[  ERROR] " ;;
        WARNING) c=32; p="[WARNING] " ;;
        INFO)    c=00; p="[   INFO] " ;;
        DEBUG)   c=00; p="[  DEBUG] " ; [ $DEBUG = 0 ] && return 0 ;;
        *) return 0 ;;
    esac

    shift ; echo -ne "\033[1;${c}m${PREFIX:-$p}$@${END:-$e}\033[0m"
}

EXIT() {
    SHOW $@
    SHOW INFO "Exit ..."
    do_clean_up
    exit 1
}

sys_umount_partition() {
    local dst=${1%/}

    [ -b "$dst" -o -d "$dst" ] || return 1
    SHOW DEBUG "Unmount $dst ..."
    for i in $(mount | grep "$dst" | awk '{print $1}') ; do
        SHOW INFO "Unmount $i ..."
        umount $i || SHOW INFO "Unmount $i failed!"
    done
    return 0
}

sys_get_current_device() {
    local curr_dev=$1 curr_vg=$2
    local root_uuid=$( grep -o '\<root=UUID=[^ ]*'  /proc/cmdline | \
                       cut -d "=" -f 3 | sed 's/"//g')
    local root_label=$(grep -o '\<root=LABEL=[^ ]*' /proc/cmdline | \
                       cut -d "=" -f 3 | sed 's/"//g')
    local root_lvm=$(  grep -o '\<root=LVM=[^ ]*'   /proc/cmdline | \
                       cut -d "=" -f 3 | sed 's/"//g')
    local root_sfr=$(  grep -o '\<root=SFR=[^ ]*'   /proc/cmdline | \
                       cut -d "=" -f 3 | sed 's/"//g')
    local root_dev=$(  grep -o '\<root=[^ ]*'       /proc/cmdline | \
                       cut -d "=" -f 2 | sed 's/"//g')
    local root_node vg dev

    [ -z "$root_lvm" ] && root_lvm=$root_sfr
    SHOW DEBUG "Boot root: $(grep -o '\<root=[^ ]*' /proc/cmdline)"

    if [ -n "$root_uuid" ] ; then
        root_node=$(blkid -U $root_uuid)
    elif [ -n "$root_label" ] ; then
        root_node=$(blkid -L $root_label)
    elif [ -n "$root_lvm" ] ; then
        vg=$(basename $root_lvm | awk -F'-' '{print $1}')
        root_node=$(pvs 2>/dev/null | grep "$vg" | tail -n1 | awk '{print $1}')
    elif [ -n "$root_dev" ] ; then
        root_node=$root_dev
    fi

    if [ -z "$root_node" ] ; then
        return 1
    elif echo $root_node | grep -q "mmcblk" ; then
        dev=${root_node:0:-2}
    else
        dev=${root_node:0:-1}
    fi

    eval $curr_vg=$vg
    eval $curr_dev=$dev
}

sys_get_device_vg() {
    local dev_node=$1 dev_vg=$2

    vg=$(pvs 2>/dev/null | grep "$dev_node" | tail -n1 | awk '{print $2}')
    [ -z "$vg" ] && EXIT ERROR "Cannot find volume group in $dev_node"
    eval $dev_vg=$vg
}

sys_get_device_size() {
    # return device total size in kB
    parted -s $1 -- unit kB print | grep "Disk $1" | \
        awk -F':' '{print $2}' | sed 's/kB//g'
}

sys_get_device_end() {
    # return last partition end offset in kB
    parted -s $1 -- unit kB print | grep "^[ ].*[0-9][0-9].*" | \
        tail -n1 | awk '{print $3}' | sed 's/kB//g'
}

sys_get_device_partition_number() {
    parted -s $1 print | grep "^[ ].*[0-9][0-9].*" | wc -l
}

sys_create_new_partition() {
    local dev_node=$1 new=$2 try_cnt=10 dev=""
    local dev_size=$(sys_get_device_size $dev_node)
    local dev_end=$(sys_get_device_end $dev_node)
    local part_num=$(sys_get_device_partition_number $dev_node)
    local mini_partition_size=$((100*1024)) #kB
    local free=$(($dev_size - $dev_end))

    # Check whether have enough size to create new partition
    SHOW INFO "Total size of $dev_node: $dev_size kB"
    SHOW INFO "Last partition end: $dev_end kB"
    SHOW INFO "Free space in $dev_node: $free kB"
    if [ -n "$free" -a $free -lt $mini_partition_size ] ; then
        EXIT ERROR "No enough free space left for a new partition"
    fi

    # Create new partition in the device
    dev_end=$((($dev_end + 999) / 1000))
    if parted -s $dev_node -- mkpart primary ext3 "${dev_end}MB" -1 ; then
        parted -s $dev_node -- set $(($part_num + 1)) lvm on
    else
        EXIT ERROR "Fail to create new partition in $dev_node"
    fi

    while [ $try_cnt -gt 0 ] ; do
        dev=$(ls ${dev_node}*$(($part_num + 1)))
        if [ -n "$dev" ] ; then
            SHOW INFO "New created partition: $dev"
            eval $new=$dev
            break
        fi
        SHOW DEBUG "Wait new partition node: $try_cnt"
        sleep 1
    done

    [ -z "$dev" ] && EXIT ERROR "Canont find new partition node"
}

sys_create_lvm_pv() {
    local pnode=$1

    umount $pnode 2>/dev/null
    SHOW INFO "Create physical volme $pnode"
    sys_umount_partition $pnode
    pvcreate -ff $pnode >/dev/null 2>&1 && return 0
    EXIT ERROR "Fail to create physical volme $pnode"
}

sys_extend_vg() {
    local vg=$1 pv=$2

    SHOW INFO "Extend volume group $vg to include $pv"
    vgextend $vg $pv >/dev/null 2>&1 && return 0
    EXIT ERROR "Fail to extend volume group $vg to include $pv"
}

sys_import_lvm_group() {
    local vg=$1

    [ -z "$vg" ] && return 1
    vgdisplay -c 2>/dev/null | grep -q "$vg" || return 0
    SHOW INFO "Import LVM volume group $vg"
    vgscan --mknodes --ignorelockingfailure >/dev/null 2>&1
    vgimport $vg >/dev/null 2>&1
    vgchange -ay --ignorelockingfailure $vg >/dev/null 2>&1
}

sys_export_lvm_group() {
    local vg=$1

    SHOW DEBUG "Try to export LVM volume group: $vg"
    [ -z "$vg" ] && return 1
    vgdisplay -c 2>&1 | grep -q "$vg is exported" && return 0
    SHOW INFO "Export LVM volume group: $vg"
    vgchange -an $vg >/dev/null 2>&1
    vgexport $vg >/dev/null 2>&1
    SHOW DEBUG "Volume groups after exported:"
}

sys_remove_lvm_logical_devices() {
    local vg=$1 device devices

    devices=$(dmsetup -c info | grep "${vg}-" | awk '{print $1}')

    SHOW INFO "Clear LVM logical devices"
    for device in $devices ; do
        grep -q "/dev/mapper/$device" /proc/cmdline && continue
        SHOW INFO "Removing LVM logical device $device"
        dmsetup remove -f $device >/dev/null 2>&1
        dmsetup -c info 2>/dev/null | grep -q "$device" || return 0
        SHOW INFO "Fail to remove LVM logical device $device"
    done
}

sys_extend_lvm_volume() {
    local vg=$1 size=$2 lv=""

    lv=$(lvdisplay -c 2>&- | grep ":$vg:" | tail -1 | \
         awk -F':' '{print $1}' | sed 's/\ //g')
    [ -z "$lv" ] && EXIT ERROR "Cannot find logical volume in group $vg"

    lv=$(basename $lv)
    SHOW DEBUG "Make logical device: ${vg}-${lv}"
    dmsetup mknodes ${vg}-${lv}

    SHOW INFO "Extend and resize logical volume to $size"
    SHOW INFO "Please wait for resizing the file system"
    SHOW INFO "Maybe need minutes or longer depends on the total size."
    echo "$size" | grep -q "%" && size="-l $size" || size="-L $size"
    SHOW DEBUG "lvextend $size -r -f /dev/${vg}/${lv}"
    lvextend $size -r -f /dev/${vg}/${lv} && return 0
    EXIT ERROR "Fail to extend /dev/${vg}/${lv} and resize it"

    # just double confirm to resize file system
    SHOW DEBUG "resize2fs /dev/mapper/${vg}-${lv}"
    [ -e /dev/mapper/${vg}-${lv} ] || return 0
    resize2fs /dev/mapper/${vg}-${lv} >/dev/null 2>&1
}

extend_lvm_group() {
    [ $EXTEND_VG = 1 ] || return 0

    # Create new partition in the device
    sys_create_new_partition $DEVICE NEW_PARTITION

    # Add phyical volume into volume group
    sys_create_lvm_pv $NEW_PARTITION
    sys_extend_vg $LVM_VG $NEW_PARTITION
}

extend_lvm_volume() {
    [ $EXTEND_LV = 1 ] || return 0

    # Extend LVM logical volume
    sys_extend_lvm_volume $LVM_VG $LV_SIZE
}

do_clean_up() {
    if [ -z "$CURRENT_DEVIE" ] ; then
        sys_umount_partition $DEVICE
        sys_export_lvm_group $LVM_VG
        sys_remove_lvm_logical_devices $LVM_VG
    fi
}

check_options() {
    if [ -z "$DEVICE" ] ; then
        sys_get_current_device DEVICE LVM_VG
        [ -b "$DEVICE" ] || EXIT ERROR "Fail to find current rootfs device"
        CURRENT_DEVIE=$DEVICE
        SHOW INFO "Current device: $DEVICE"

        # Fix backup GPT table is not at the end of the disk issue
        parted $DEVICE print -- Fix Fix >/dev/null 2>&1

        [ -z "$LVM_VG" ] && EXIT INFO "Nothing to do to extend LVM volume"
        sys_import_lvm_group $LVM_VG
    elif [ -n "$DEVICE" -a -b "$DEVICE" ] ; then
        CURRENT_DEVIE=""
        sys_umount_partition $DEVICE
        sys_get_device_vg $DEVICE LVM_VG
        sys_import_lvm_group $LVM_VG
        sys_umount_partition $LVM_VG
    else
        EXIT ERROR "Invalid block device: $DEVICE"
    fi

    [ -z "$LV_SIZE" ] && LV_SIZE=$DEFAULT_LV_SIZE

    SHOW DEBUG "----------------------------"
    SHOW DEBUG "OPTIONS: $OUT"
    SHOW DEBUG "----------------------------"
    SHOW DEBUG "DEVICE   : $DEVICE"
    SHOW DEBUG "VG NAME  : $LVM_VG"
    SHOW DEBUG "EXTEND_VG: $EXTEND_VG"
    SHOW DEBUG "EXTEND_LV: $EXTEND_LV"
    SHOW DEBUG "LV_SIZE  : $LV_SIZE"
    SHOW DEBUG "DEBUG    : $DEBUG"
    SHOW DEBUG "----------------------------"
}

show_help()
{
    cat <<EOF
LVM volume group and logical volume extention tool.

Usage:  sudo $(basename $0) OPTION ...

OPTIONs:

-d/--device device:
    Specify the device node which is to be extend (e.g. /dev/sdb).
    If it is not specified, extend current root file system on line.

-E/--extend-vg:
    Extend the LVM volume group by creating a new partition to use all
    the left space in the device and then adding it into the volume group.

-e[SIZE]/--extend-lv[=SIZE]:
    Resize LVM volume to SIZE, see also '-l/-L' option in lvextend command.
    If SIZE is not specified, extend LVM logical volume to 40%VG, and reserv
    the left spaces for LVM snapshots.

EOF
}

#===============================================================================
# Start Main
#===============================================================================
trap "EXIT INFO Canceled" SIGINT

OUT="`getopt -o d:Ee::vh -l device:,extend-vg,extend-lv::,verbose,help -- $@`"
eval set -- "$OUT"

while [ -n "$1" ]; do
    case "$1" in
        -d|--device)    DEVICE=$2 ; shift 2 ;;
        -E|--extend-vg) EXTEND_VG=1 ; shift 1 ;;
        -e|--extend-lv) EXTEND_LV=1 ; LV_SIZE=$2 ; shift 2 ;;
        -v|--verbose)   DEBUG=1 ; shift 1 ;;
        -h|--help)      show_help; exit 0 ;;
        --)             shift ; break ;;
        *)              EXIT ERROR "Not supported option $1" ;;
    esac
done

check_options
extend_lvm_group
extend_lvm_volume
do_clean_up
exit 0
