#!/bin/bash

###############################################################################
# Copyright © Huawei Technologies Co., Ltd. 2022-2022. All rights reserved.
# File name: disk_partition.sh
# Description: 适配NEC-Common默认分区，安装时根据分析器场景调整分区
###############################################################################

# 指定VG名、LV名及容量和日志路径
VG_NAME="vg_root"
declare -A COMPULSORY_LV_SIZE_DICT # [LV名]=LV容量
COMPULSORY_LV_SIZE_DICT=([root]=40 [home]=1 [usr]=10 [var]=20 [var_audit]=3 [var_log]=40 [var_tmp]=5 [tmp]=20 [opt]=100 [root_bak]=40 [usr_bak]=10 [var_bak]=20)
declare -A COMPULSORY_LV_MOUNT_PATH_DICT # [LV名]=LV挂载点
COMPULSORY_LV_MOUNT_PATH_DICT=([root]="/" [home]="/home" [usr]="/usr" [var]="/var" [var_audit]="/var/log/audit" [var_log]="/var/log" [var_tmp]="/var/tmp" [tmp]="/tmp" [opt]="/opt" [root_bak]="" [usr_bak]="" [var_bak]="")
declare -A SELECTIVE_LV_SIZE_DICT # [LV名]=LV容量
SELECTIVE_LV_SIZE_DICT=([home_bak]=1)
LOG_FILE='/root/install_config/disk_partition.log'

##################################################
# Description: 检查是否有残留vg_appdata并清除
# Parameters:  无
##################################################
function clean_vg() {
  lvs /dev/${VG_NAME}/opt > /dev/null 2>&1
  # 有opt分区说明已执行过安装，不进行清理
  if [[ $? -ne 0 ]]; then # 没有opt分区，执行清理
    vgs vg_appdata > /dev/null 2>&1
    if [[ $? -eq 0 ]]; then
      echo " " >> ${LOG_FILE}
      echo "-----------------------------------------------------" >> ${LOG_FILE}
      echo "remove vg_appdata." >> ${LOG_FILE}
      echo "-----------------------------------------------------" >> ${LOG_FILE}
      vgremove -y vg_appdata
    fi
  fi
}

##################################################
# Description: 检查磁盘空间是否满足分区要求，并识别出需要创建的LV
# Parameters:  无
##################################################
function volume_check() {
  # 修复可能存在的问题：虚拟机增加空间后未更新GPT表
  local GPT_fix_status=$(parted -l -s 2>&1 | grep "fix")
  if [[ x${GPT_fix_status} != x ]]; then
    printf "fix\n" | parted -l ---pretend-input-tty > /dev/null 2>&1
    if [[ $? -eq 0 ]]; then
      echo " " >> ${LOG_FILE}
      echo "-----------------------------------------------------" >> ${LOG_FILE}
      echo "[successfully]Update GPT for disks." >> ${LOG_FILE}
      echo "-----------------------------------------------------" >> ${LOG_FILE}
    else
      echo " " >> ${LOG_FILE}
      echo "-----------------------------------------------------" >> ${LOG_FILE}
      echo "[ERROR message]Failed to update GPT for disks." >> ${LOG_FILE}
      echo "-----------------------------------------------------" >> ${LOG_FILE}
      return 5
    fi
  fi

  # 找到LVM对应的磁盘名
  disk=$(fdisk -l | grep "Linux LVM" | tail -1 | awk '{print $1}')
  disk=${disk%?}
  if [[ x${disk} == x ]]; then
    echo " " >> ${LOG_FILE}
    echo "-----------------------------------------------------" >> ${LOG_FILE}
    echo "[ERROR message]There is no LVM in the disk." >> ${LOG_FILE}
    echo "-----------------------------------------------------" >> ${LOG_FILE}
    return 5
  fi

  # disk_free_space记录磁盘未分配空间大小
  disk_free_space=$(parted ${disk} unit gib print free | grep "Free Space" | tail -1 | awk '{print $3}')
  disk_free_space=${disk_free_space%G*}

  # vg_free_space记录VG未分配空间大小
  vg_free_space=$(vgdisplay ${VG_NAME} --unit g | grep "Free" | awk '{print $7}')

  # required_space记录LV需要的总空间，allocated_space记录已分配给LV的空间
  required_space=0
  allocated_space=0
  for lv_name in $(echo ${!COMPULSORY_LV_SIZE_DICT[*]}); do
    required_space_for_lv=${COMPULSORY_LV_SIZE_DICT[${lv_name}]}
    required_space=$(echo "${required_space} + ${required_space_for_lv}" | bc)

    lsblk | grep ${lv_name} > /dev/null 2>&1
    if [[ $? -eq 0 ]]; then
      # temp存放已分配给当前lv_name的空间，allocated_space相应增加
      temp=$(lvdisplay /dev/${VG_NAME}/${lv_name} --units g | grep "LV Size" | awk '{print $3}')
      if [[ $(echo "${temp} < ${required_space_for_lv}" | bc) -eq 1 ]]; then
        allocated_space=$(echo "${allocated_space} + ${temp}" | bc)
        lv_to_be_extended+=" ${lv_name}" # 若已分配空间小于需要分配的空间，说明需要扩大分区
      else
        allocated_space=$(echo "${allocated_space} + ${required_space_for_lv}" | bc)
      fi
    else
      # 若lv不存在，需要创建
      lv_to_be_created+=" ${lv_name}"
    fi
  done

  for lv_name in $(echo ${!SELECTIVE_LV_SIZE_DICT[*]}); do
    lsblk | grep ${lv_name} > /dev/null 2>&1
    if [[ $? -eq 0 ]]; then
      required_space_for_lv=${SELECTIVE_LV_SIZE_DICT[${lv_name}]}
      required_space=$(echo "${required_space} + ${required_space_for_lv}" | bc)
      # temp存放已分配给当前lv_name的空间，allocated_space相应增加
      temp=$(lvdisplay /dev/${VG_NAME}/${lv_name} --units g | grep "LV Size" | awk '{print $3}')
      if [[ $(echo "${temp} < ${required_space_for_lv}" | bc) -eq 1 ]]; then
        allocated_space=$(echo "${allocated_space} + ${temp}" | bc)
        lv_to_be_extended+=" ${lv_name}" # 若已分配空间小于需要分配的空间，说明需要扩大分区
      else
        allocated_space=$(echo "${allocated_space} + ${required_space_for_lv}" | bc)
      fi
    fi
  done

  echo " " >> ${LOG_FILE}
  echo "-----------------------------------------------------" >> ${LOG_FILE}
  echo "[INFO message]The disk has ${disk_free_space}GiB free space." >> ${LOG_FILE}
  echo "[INFO message]vg_root has ${vg_free_space}GiB free space." >> ${LOG_FILE}
  echo "[INFO message]Required space is ${required_space}GiB." >> ${LOG_FILE}
  echo "[INFO message]Allocated space is ${allocated_space}GiB." >> ${LOG_FILE}
  echo "-----------------------------------------------------" >> ${LOG_FILE}

  # 判断磁盘空间是否满足要求
  if [[ $(echo "${allocated_space} == ${required_space}" | bc) -eq 1 ]];then
    echo " " >> ${LOG_FILE}
    echo "-----------------------------------------------------" >> ${LOG_FILE}
    echo "[successfully]The disk partition already meets requirements." >> ${LOG_FILE}
    echo "-----------------------------------------------------" >> ${LOG_FILE}
    return 2
  else
    if [[ $(echo "${disk_free_space} + ${vg_free_space} + ${allocated_space} < ${required_space}" | bc) -eq 1 ]]; then
      echo "custom_check_disk_partition"
      echo " "
      echo "-----------------------------------------------------"
      echo "[failed]"
      echo "[ERROR message]There is not enough space in the disk."
      echo "-----------------------------------------------------"
      echo " "
      return 5
    else
      return 0
    fi
  fi

  return 0
}

##################################################
# Description: 检查空闲容量创建PV，并加入VG
# Parameters:  无
##################################################
function extend_vg() {
  # 检查是否有空闲容量
  total_sector=$(parted ${disk} unit s print | grep "^Disk ${disk}:" | awk '{print $3}')
  total_sector=${total_sector%?}
  end_sector=$(parted ${disk} unit s print | tail -2 | head -1 | awk '{print $3}')
  end_sector=${end_sector%?}

  if [[ $((${total_sector} - ${end_sector})) -gt 3000 ]]; then
    # parted创建实体分区
    new_start_sector=$((${end_sector} + 1))
    parted ${disk} mkpart primary ${new_start_sector}s 100% > /dev/null 2>&1
    if [[ $? -eq 0 ]]; then
      echo " " >> ${LOG_FILE}
      echo "-----------------------------------------------------" >> ${LOG_FILE}
      echo "[successfully]New partition has been parted." >> ${LOG_FILE}
      echo "-----------------------------------------------------" >> ${LOG_FILE}
    else
      echo " " >> ${LOG_FILE}
      echo "-----------------------------------------------------" >> ${LOG_FILE}
      echo "[ERROR message]Parted failed." >> ${LOG_FILE}
      echo "-----------------------------------------------------" >> ${LOG_FILE}
      return 5
    fi

    # 创建PV
    local new_part_number=$(parted ${disk} unit s print | tail -2 | head -1 | cut -c 2)
    local new_part_path=${disk}${new_part_number}
    pvcreate ${new_part_path} > /dev/null 2>&1

    # 将新PV加入VG
    local disk=$(pvscan | grep "${VG_NAME}" | awk '{print $2}')
    disk=${disk%?}
    new_pv_path=$(pvscan | grep "${disk}"| tail -1 | awk '{print $2}')
    vgextend ${VG_NAME} ${new_pv_path} > /dev/null 2>&1
  fi
}

##################################################
# Description: 创建分区
# Parameters:  无
##################################################
function create_partition() {
  systemctl daemon-reload  # 清除无用opt分区后，需重新加载systemd配置文件，否则创建新opt分区后会挂载失败
  for lv_name in ${lv_to_be_created}; do
    lv_mount_point=${COMPULSORY_LV_MOUNT_PATH_DICT[${lv_name}]}
    lvcreate -y -L ${COMPULSORY_LV_SIZE_DICT[${lv_name}]}G -n ${lv_name} ${VG_NAME} > /dev/null 2>&1
    mkfs -t ext4 /dev/${VG_NAME}/${lv_name} > /dev/null 2>&1
    if [[ $? -eq 0 ]]; then
      echo " " >> ${LOG_FILE}
      echo "-----------------------------------------------------" >> ${LOG_FILE}
      echo "[successfully]${lv_name} has been created." >> ${LOG_FILE}
      echo "-----------------------------------------------------" >> ${LOG_FILE}
    else
      echo " " >> ${LOG_FILE}
      echo "-----------------------------------------------------" >> ${LOG_FILE}
      echo "[ERROR message] Creating ${lv_name} failed." >> ${LOG_FILE}
      echo "-----------------------------------------------------" >> ${LOG_FILE}
      return 5
    fi
    if [[ x${lv_mount_point} != x ]]; then
      if [[ ! -d ${lv_mount_point} ]]; then
        mkdir ${lv_mount_point}
      fi
      mount /dev/${VG_NAME}/${lv_name} ${lv_mount_point}
      echo "/dev/mapper/${VG_NAME}-${lv_name}  ${lv_mount_point}  ext4  defaults  1 2" >> /etc/fstab
    fi
  done
}

##################################################
# Description: 调整分区大小
# Parameters:  无
##################################################
function extend_partition() {
  for lv_name in ${lv_to_be_extended}; do
    lvextend -L ${COMPULSORY_LV_SIZE_DICT[${lv_name}]}G /dev/${VG_NAME}/${lv_name} > /dev/null 2>&1 && resize2fs /dev/${VG_NAME}/${lv_name} > /dev/null 2>&1
    if [[ $? -eq 0 ]]; then
      echo " " >> ${LOG_FILE}
      echo "-----------------------------------------------------" >> ${LOG_FILE}
      echo "[successfully]${lv_name} has been extended." >> ${LOG_FILE}
      echo "-----------------------------------------------------" >> ${LOG_FILE}
    else
      echo " " >> ${LOG_FILE}
      echo "-----------------------------------------------------" >> ${LOG_FILE}
      echo "[ERROR message] Extending ${lv_name} failed." >> ${LOG_FILE}
      echo "-----------------------------------------------------" >> ${LOG_FILE}
      return 5
    fi
  done
}

##################################################
# Description: 删除多余分区
# Parameters:  无
##################################################
function remove_partition() {
  # 删除swap分区
  swap_status=$(swapon -s)
  os_version=$(grep EulerOS /etc/os-release)
  # 如存在swap分区，且是EulerOS，则进行删除（SUSE删除swap分区需重做内核mkinitrd）
  if [[ x${swap_status} != x && x${os_version} != x ]]; then
    swapoff /dev/vg_root/swap > /dev/null 2>&1
    sed "/swap/d" /etc/fstab > /etc/fstab.bak
    mv /etc/fstab.bak /etc/fstab
    chmod 644 /etc/fstab
    sed "s#rd.lvm.lv=vg_root/swap##g" /etc/default/grub > /etc/default/grub.bak
    mv /etc/default/grub.bak /etc/default/grub
    chmod 644 /etc/default/grub
    # 区分BIOS和UEFI
    if [[ -d /sys/firmware/efi ]]; then
      grub2-mkconfig -o /boot/efi/EFI/euleros/grub.cfg > /dev/null 2>&1  # UEFI
    else
      grub2-mkconfig -o /boot/grub2/grub.cfg > /dev/null 2>&1  # BIOS
    fi
    if [[ $? -eq 0 ]]; then
      echo " " >> ${LOG_FILE}
      echo "-----------------------------------------------------" >> ${LOG_FILE}
      echo "[successfully]grub2-mkconfig succeeded." >> ${LOG_FILE}
      echo "-----------------------------------------------------" >> ${LOG_FILE}
    else
      echo " " >> ${LOG_FILE}
      echo "-----------------------------------------------------" >> ${LOG_FILE}
      echo "[ERROR message]grub2-mkconfig failed." >> ${LOG_FILE}
      echo "-----------------------------------------------------" >> ${LOG_FILE}
      return 5
    fi
    lvremove -f /dev/vg_root/swap > /dev/null 2>&1
  fi
}

##################################################
# Description: 主函数
# Parameters:  无
##################################################
function main_disk_partition() {
  clean_vg

  remove_partition
  if [[ $? -eq 5 ]]; then
    return 5
  fi

  volume_check
  volume_check_return=$?
  if [[ ${volume_check_return} -eq 2 ]]; then
    return 0
  elif [[ ${volume_check_return} -eq 5 ]]; then
    return 5
  fi

  extend_vg

  create_partition
  if [[ $? -eq 5 ]]; then
    return 5
  fi

  extend_partition
  if [[ $? -eq 5 ]]; then
    return 5
  fi
}

function main() {
    typeset -i fail_num=0

    # step1: 产品安装前分区
    main_disk_partition
    if [[ $? -ne 0 ]]
    then
        ((fail_num=fail_num+1))
        echo "custom_check_disk_partition"
    else
        echo "[successfully] Disk partition successfully!" >> ${LOG_FILE}
        echo "custom_check_disk_partition" >> ${LOG_FILE}
    fi

    if [[ ${fail_num} -gt 0 ]]
    then
        echo "==============================DISK PARTITION FAILED=================================" >&2
        echo "===================================================================================="
        return 1
    fi
}
main $@
