#!/bin/bash
# Copyright 2021 Red Hat, Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License"); you may
# not use this file except in compliance with the License. You may obtain
# a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
# License for the specific language governing permissions and limitations
# under the License.

#
# Script to upload files to a Artifact container for deployment via
# TripleO Heat Templates.
#

set -eu
set -o pipefail
SCRIPT_NAME=$(basename $0)
NBD_DEVICE=/dev/nbd0
IMAGE_FILE=""
MOUNT_DIR=""
if [ ! -a "${NBD_DEVICE}" ]; then
    modprobe nbd
fi

# supported LVM devices and mount points for whole-disk overcloud images
MOUNTS="/dev/mapper/vg-lv_var:/var \
        /dev/mapper/vg-lv_log:/var/log \
        /dev/mapper/vg-lv_audit:/var/log/audit \
        /dev/mapper/vg-lv_home:/home \
        /dev/mapper/vg-lv_tmp:/tmp \
        /dev/mapper/vg-lv_srv:/srv"
REVERSE_MOUNTS=""
for m in $MOUNTS; do
    REVERSE_MOUNTS="$m $REVERSE_MOUNTS"
done

mount_show_options() {
    echo "Usage: $SCRIPT_NAME"
    echo
    echo "Options:"
    echo "    -h, --help                  -- print this help."
    echo "    -a <file>                   -- Image file to mount."
    echo "    -m <directory>              -- Directory to mount image to."
    echo "    -n <nbd device>             -- NBD device to use."
    echo "                                   Defaults to /dev/nbd0"
    echo
    echo "Mount an overcloud image to a directory"
    echo
    exit $1
}

unmount_show_options() {
    echo "Usage: $SCRIPT_NAME"
    echo
    echo "Options:"
    echo "    -h, --help                  -- print this help."
    echo "    -m <directory>              -- Directory to unmount."
    echo "    -n <nbd device>             -- NBD device to disconnect."
    echo "                                   Defaults to /dev/nbd0"
    echo
    echo "Unmount a mounted overcloud image"
    echo
    exit $1
}

mount_volume () {
    if [ -b "$1" ]; then
        if [ ! -d $2 ]; then
            mkdir $2
        fi
        mount $1 $2
    fi
}

unmount_volume () {
    if mountpoint "$1"; then
        umount $1
    fi
}

remove_device () {
    if [ -b "$1" ]; then
        dmsetup remove $1
    fi
}

mount_image() {
    set -x

    if qemu-img info --output json $IMAGE_FILE |grep '"format": "raw"' ; then
        image_format='--format raw'
    elif qemu-img info --output json $IMAGE_FILE |grep '"format": "qcow2"' ; then
        image_format='--format qcow2'
    else
        image_format=''
    fi
    qemu-nbd $image_format --connect $NBD_DEVICE $IMAGE_FILE

    # search for the vg volume group, this is automatic in some environments
    vgscan
    # refresh for when this script is called with different values of $NBD_DEVICE
    vgchange --refresh

    # activate new logical volumes, this is automatic in some environments
    vgchange -ay

    root_device=""
    boot_device=""
    efi_device=""

    # wait for any sub-devices to appear
    timeout 5 sh -c "while ! ls ${NBD_DEVICE}p* ; do sleep 1; done" || true

    devices=$(ls -1 ${NBD_DEVICE}p*)
    device_count=$(echo $devices | wc -w)
    if [ $device_count == "0" ]; then
        # if there are no partition devices, assume one root device
        root_device=${NBD_DEVICE}
    elif [ $device_count == "1" ]; then
        # if there is one partition device, assume it is the root device
        root_device=${devices}
        devices=""
    fi

    for device in ${devices}; do
        lsblk --nodeps -P --output-all $device
        label=$(lsblk --nodeps --noheadings --output LABEL $device)
        part_type_name=$(lsblk --nodeps --noheadings --output PARTTYPENAME $device)
        part_label=$(lsblk --nodeps --noheadings --output PARTLABEL $device)

        if [[ ${part_type_name} == "BIOS boot" ]] || [[ ${part_type_name} == "PowerPC PReP boot" ]]; then
            # Ignore unmountable partition
            continue
        fi

        # look for EFI partition to mount at /boot/efi
        if [ -z "$efi_device" ]; then
            if [[ ${part_type_name} == "EFI System" ]]; then
                efi_device=$device
                continue
            fi
        fi

        # look for partition to mount as /boot, only the RHEL guest image is known
        # to have this
        if [ -z "$boot_device" ]; then
            if [[ ${label} == "boot" ]]; then
                boot_device=$device
                continue
            fi
        fi

        if [ -z "$root_device" ]; then
            root_device=$device
            continue
        fi
    done

    if [ -z "$root_device" ]; then
        echo "ERROR: No root device found to mount"
        exit 1
    else
        if [ -b "/dev/mapper/vg-lv_root" ]; then
            # a whole-disk overcloud with lvm volumes
            # for example, overcloud-hardened-uefi-full.qcow2
            mount /dev/mapper/vg-lv_root $MOUNT_DIR
            for m in $MOUNTS; do
                device=${m%:*}
                path=${m#*:}
                mount_volume $device $MOUNT_DIR$path
            done
        else
            # a simple root partition
            mount $root_device $MOUNT_DIR
        fi
    fi
    if [ ! -z "$boot_device" ]; then
        # mount to /boot
        mount $boot_device $MOUNT_DIR/boot
    fi
    if [ ! -z "$efi_device" ]; then
        # mount to /boot/efi
        mount $efi_device $MOUNT_DIR/boot/efi
    fi
}

unmount_image() {

    set -x

    if mountpoint "$MOUNT_DIR"; then
        for m in $REVERSE_MOUNTS; do
            path=${m#*:}
            unmount_volume $MOUNT_DIR$path
        done
        unmount_volume $MOUNT_DIR/boot/efi
        unmount_volume $MOUNT_DIR/boot
        unmount_volume $MOUNT_DIR
    fi

    # `--activate n` makes LVs inactive, they must be set
    # inactive so that the nbd device can be disconnected.
    # Ref bug: https://bugs.launchpad.net/tripleo/+bug/1950137
    vgchange --activate n vg || true
    qemu-nbd --disconnect $NBD_DEVICE
    vgchange --refresh vg || true

    for m in $REVERSE_MOUNTS; do
        device=${m%:*}
        remove_device $device
    done
    remove_device vg-lv_root
}


if [ $SCRIPT_NAME == "tripleo-unmount-image" ]; then
    TEMP=`getopt -o hm:n: -l help -n $SCRIPT_NAME -- "$@"`
    if [ $? != 0 ]; then
        echo "Terminating..." >&2
        exit 1
    fi
    eval set -- "$TEMP"

    while true ; do
        case "$1" in
            -h|--help) unmount_show_options 0 >&2;;
            -m) MOUNT_DIR=$2 ; shift 2;;
            -n) NBD_DEVICE=$2 ; shift 2;;
            --) shift ; break;;
            *) echo "Error: unsupported option $1." ; exit 1;;
        esac
    done
    if [ -z "${MOUNT_DIR}" ]; then
        unmount_show_options 1
    fi
    unmount_image
else
    TEMP=`getopt -o ha:m:n: -l help -n $SCRIPT_NAME -- "$@"`
    if [ $? != 0 ]; then
        echo "Terminating..." >&2
        exit 1
    fi
    eval set -- "$TEMP"


    while true ; do
        case "$1" in
            -h|--help) mount_show_options 0 >&2;;
            -a) IMAGE_FILE=$2 ; shift 2;;
            -m) MOUNT_DIR=$2 ; shift 2;;
            -n) NBD_DEVICE=$2 ; shift 2;;
            --) shift ; break;;
            *) echo "Error: unsupported option $1." ; exit 1;;
        esac
    done
    if [ -z "${MOUNT_DIR}" ] || [ -z "${IMAGE_FILE}" ]; then
        mount_show_options 1
    fi
    mount_image
fi