#!/bin/bash

#########################################
# This script is used to upgrade ZStack #
#########################################


[ -f /usr/bin/zstack-ctl ] && zstack_home=`zstack-ctl getenv ZSTACK_HOME | awk -F '=' '{ print $2 }'`
[ -z $zstack_home ] && zstack_home="/usr/local/zstack/apache-tomcat/webapps/zstack"

# Helper function
REPO_ONLY='n'
ADD_REPO='n'
FORCE_UPGRADE_DB='n'
FORCE_UPGRADE_REPO='n'
PRODUCT_NAME='Cloud'
UPGRADE_LOG="/tmp/${PRODUCT_NAME,,}-upgrade.log"

call_help() {
  echo "Usage: $0 [-h] [-a] [-r] [-F] LOCAL_ISO_PATH|REMOTE_ISO_LINK"
  echo ""
  echo "Optional arguments:"
  echo -e "  -h\tShow help message"
  echo -e "  -a/--add_repo\tAdd a new repo"
  echo -e "  -r\tOnly upgrade ${PRODUCT_NAME} local repo"
  echo -e "  -F\tDo force database upgrating."
  echo -e "    \tNOTE: only use -F when you know exactly what it does"
  exit 0
}

fail(){
    sync
    tput cub 6
    echo -e "$(tput setaf 1) FAIL$(tput sgr0)" | tee -a $UPGRADE_LOG
    echo -e "$(tput setaf 1) Reason: $*\n$(tput sgr0)" | tee -a $UPGRADE_LOG
    echo "The detailed upgrade log could be found in $UPGRADE_LOG"
    exit 1
}

pass(){
    echo -e "$(tput setaf 2) PASS$(tput sgr0)" | tee -a $UPGRADE_LOG
    sync
}

# Upgrade ZStack using local ISO file or remote ISO link
OPTS=`getopt -o ahrfF --long add_repo -- "$@"`
eval set -- "$OPTS"
while true; do
  case "$1" in
    -r) REPO_ONLY='y'; shift 1;;
    -a) ADD_REPO='y'; shift 1;;
    --add_repo) ADD_REPO='y'; shift 1;;
    -F) FORCE_UPGRADE_DB='y'; shift 1;;
    -h|--help) call_help; exit 0;;
    *) shift; break;;
  esac
done

[ $# != 1 ] && call_help
ISO_PATH=$1
ISO_ARCH='x86_64'
ISO_VER=''

[[ "${ISO_PATH}" =~ .*aarch64.*$ ]] && ISO_ARCH='aarch64'
[[ "${ISO_PATH}" =~ .*arm64.*$ ]] && ISO_ARCH='aarch64'
[[ "${ISO_PATH}" =~ .*mips64el.*$ ]] && ISO_ARCH='mips64el'
[[ "${ISO_PATH}" =~ .*loongarch64.*$ ]] && ISO_ARCH='loongarch64'

# Helper function - Echo Title
STEP="1"
echo_title() {
  echo ""
  echo $(tput bold)" ${STEP}. $*"$(tput sgr0) | tee -a $UPGRADE_LOG
  STEP=`expr $STEP + 1`
}

# Get and check ISO
PWD=`pwd`
get_check_iso() {
  echo_title "Download Remote ISO / Check Local ISO..."
  # Remote ISO Link - ZStack Staff Only
  if [[ $ISO_PATH =~ ^http://.*/latest/*$ ]]; then
    CONTEXT=`curl ${ISO_PATH}/ 2>/dev/null | grep -e "<a href"`
    REGEX='(<th><a\ +href=)([^\"]+.iso)(>.*</th>)'
    if [[ $CONTEXT =~ $REGEX ]]; then
      ISO_PATH=${ISO_PATH}/${BASH_REMATCH[2]}
    else
      fail "Cannot get ZStack ISO download link from $ISO_PATH"
    fi
    wget -c $ISO_PATH -O ZStack-x86_64-DVD.iso
    if [ $? -eq 0 ]; then
      ISO_PATH=${PWD}/ZStack-x86_64-DVD.iso
    else
      fail "Cannot download ZStack ISO from $ISO_PATH"
    fi
  # Remote ISO Link
  elif [[ $ISO_PATH =~ ^http://.*.iso$  || \
      $ISO_PATH =~ ^https://.*.iso$ || \
      $ISO_PATH =~ ^ftp://.*.iso$ ]]; then
    wget -c $ISO_PATH -O ZStack-x86_64-DVD.iso
    if [ $? -eq 0 ]; then
      ISO_PATH=${PWD}/ZStack-x86_64-DVD.iso
    else
      fail "Cannot download ZStack ISO from $ISO_PATH"
    fi
  # Local ISO File - Absolute Path
  elif [[ $ISO_PATH =~ ^/.*.iso.*$ ]]; then
    ISO_PATH=${ISO_PATH}
  # Local ISO File - Relative Path
  else
    ISO_PATH=${PWD}/${ISO_PATH}
  fi

  if [ ! -f ${ISO_PATH} ]; then
    fail "${ISO_PATH} doesn't exists!"
  fi
  
  # check root space
  root_space=`df -h -BG|awk '/\/$/{print $4}'|tr -d "G"`
  iso_size=`du -sh $ISO_PATH |awk '{print $1}'|tr -d "G"`
  ret=`awk -v num1=$root_space -v num2=$iso_size 'BEGIN{print(num1>num2)?"0":"1"}'`
  if [ $ret -eq 1 ];then
    fail "Not enough space to apply iso."
  fi

  pass
}

make_tmp_repo() {
echo_title "Make temporary YUM repo for dependences..."
mkdir -p /opt/tmp-iso
umount /opt/tmp-iso/ >/dev/null 2>&1
[ -d /opt/tmp-iso ] && mount ${ISO_PATH} /opt/tmp-iso/ >/dev/null 2>&1
cat > /etc/yum.repos.d/zstack-tmp.repo << EOF
[zstack-tmp]
name=zstack-tmp
baseurl=file:///opt/tmp-iso/
gpgcheck=0
enabled=0
EOF

pass
}

# Install TUI/Expert Mode dependences in case of upgrading from older ISO
install_deps() {
  echo_title "Install Expert Mode dependences..."
  pkg_list="expect gcc python-devel rsync"
  rpm -q ${pkg_list}
  [ $? -eq 0 ] && return 0
  yum -y install --disablerepo=* --enablerepo=zstack-tmp ${pkg_list}
  if [ $? -ne 0 ]; then
    echo "Failed to install Expert Mode dependences etc..."
    exit -1
  fi
  rm -f /etc/yum.repos.d/zstack-tmp.repo
}

upgrade_repo() {
echo_title "Upgrade local repo & Create new repo files..."
which rsync >/dev/null 2>&1
  [ $? -ne 0 ] && rpm -ivh /opt/tmp-iso/Packages/rsync* >>$UPGRADE_LOG 2>&1

umount -l /opt/zstack-dvd/Extra/qemu-kvm-ev >/dev/null 2>&1
rsync -av --delete /opt/tmp-iso/ /opt/zstack-dvd/
[ $? -eq 0 ] && umount -l /opt/tmp-iso/
[ $? -eq 0 ] && rm -rf /opt/tmp-iso/

# create repo files
cat > /etc/yum.repos.d/zstack-local.repo << EOF
[zstack-local]
name=zstack-local
baseurl=file:///opt/zstack-dvd/\$basearch/\$YUM0
gpgcheck=0
enabled=1
EOF

cat > /etc/yum.repos.d/ceph.repo << EOF
[ceph]
name=Ceph
baseurl=file:///opt/zstack-dvd/\$basearch/\$YUM0/Extra/ceph
gpgcheck=0
enabled=0
EOF

cat > /etc/yum.repos.d/uek4-ocfs2.repo << EOF
[uek4-ocfs2]
name=UEK4-OCFS2
baseurl=file:///opt/zstack-dvd/\$basearch/\$YUM0/Extra/uek4
gpgcheck=0
enabled=0
EOF

cat > /etc/yum.repos.d/galera.repo << EOF
[mariadb]
name = MariaDB
baseurl=file:///opt/zstack-dvd/\$basearch/\$YUM0/Extra/galera
gpgcheck=0
enabled=0
EOF

cat > /etc/yum.repos.d/qemu-kvm-ev.repo << EOF
[qemu-kvm-ev]
name=Qemu KVM EV
baseurl=file:///opt/zstack-dvd/\$basearch/\$YUM0/Extra/qemu-kvm-ev
gpgcheck=0
enabled=0
EOF

cat > /etc/yum.repos.d/virtio-win.repo << EOF
[virtio-win]
name=virtio-win
baseurl=file:///opt/zstack-dvd/\$basearch/\$YUM0/Extra/virtio-win
gpgcheck=0
enabled=0
EOF

yum clean all
}

do_ks_post() {
  echo_title "Do Kickstart Post - For Expert Mode Only"

# config history
cat > "/etc/profile.d/history.sh" << EOF
shopt -s histappend
HISTTIMEFORMAT='%F %T '
HISTSIZE="5000"
HISTFILESIZE=5000
PROMPT_COMMAND="history -a"
export HISTTIMEFORMAT HISTSIZE HISTFILESIZE PROMPT_COMMAND
EOF
source /etc/profile.d/history.sh

  # do not run the following commands twice
  grep "NAutoVTs=2" /etc/systemd/logind.conf
  [ $? -eq 0 ] && return

  # modify logind.conf, so that only two ttys are avaiable
  echo "NAutoVTs=2" >> /etc/systemd/logind.conf
  echo "ReserveVT=2" >> /etc/systemd/logind.conf

  # set tuned profile to virtual-host
  echo virtual-host > /etc/tuned/active_profile

  # install npyscreen and psutil
  cd /opt/zstack-dvd/Extra/pip/npyscreen
  python setup.py install
  cd /opt/zstack-dvd/Extra/pip/psutil
  python setup.py install

  # hide native yum repos
  mkdir -p /opt/zstack-dvd/Extra/native-repos
  mv /etc/yum.repos.d/CentOS-*.repo /opt/zstack-dvd/Extra/native-repos

  # configs
  sed -i "s/ONBOOT.*/ONBOOT\=yes/g" /etc/sysconfig/network-scripts/ifcfg-*
  sed -i "s/SELINUX\=enforcing/SELINUX\=disabled/g" /etc/selinux/config
  sed -i "s/#UseDNS yes/UseDNS no/g" /etc/ssh/sshd_config
  echo 'IndexOptions NameWidth=*' >> /etc/httpd/conf/httpd.conf
  rm -f /etc/httpd/conf.d/welcome.conf
  echo -e 'export TMOUT=600\nreadonly TMOUT' >> /etc/profile
  if [ -f /etc/libvirt/qemu/networks/autostart/default.xml ]; then
    unlink /etc/libvirt/qemu/networks/autostart/default.xml
  fi

  # services
  chkconfig NetworkManager off
  chkconfig firewalld off
  chkconfig iptables on
  chkconfig network on
  chkconfig httpd on
  chkconfig atd on
}

install_zstack_manager() {
  echo_title "Install/Reinstall zstack-manager.rpm..."
  rpm -e zstack-enterprise-host >/dev/null 2>&1
  rpm -e zstack-community-host >/dev/null 2>&1
  rpm -e zstack-host >/dev/null 2>&1
  yum install --disablerepo=* --enablerepo=zstack-local zstack-*manager -y
if [ ! -f /etc/yum.repos.d/zstack-local.repo ]; then
cat > /etc/yum.repos.d/zstack-local.repo << EOF
[zstack-local]
name=zstack-local
baseurl=file:///opt/zstack-dvd/
gpgcheck=0
enabled=1
EOF
fi
  yum reinstall --disablerepo=* --enablerepo=zstack-local zstack-*manager -y
  if [ $? -ne 0 ]; then
    echo "Failed to install zstack manager"
    exit -1
  fi
}

upgrade_zstack() {
  echo_title "Upgrade ZStack..."
  /bin/cp -f /opt/zstack-dvd/zstack-installer.bin /opt/
  if [ ! -f /opt/zstack-*installer.bin ]; then
    echo "zstack installer.bin doesn't exists in /opt"
    exit -1
  fi

  # whether mevoco-1.x before upgrade
  zstack_home=`zstack-ctl getenv ZSTACK_HOME | awk -F '=' '{ print $2 }'`
  [ -z $zstack_home ] && zstack_home="/usr/local/zstack/apache-tomcat/webapps/zstack"
  if [ -f $zstack_home/WEB-INF/lib/mevoco-1\.*.jar ]; then
    NEED_INSTALL_TRIAL_LICENSE='y'
  fi

  installer=`ls -tc1 /opt/zstack-*installer.bin | head -n 1`
  base_upgrade_command="bash ${installer} -u"
  if [ x"$FORCE_UPGRADE_DB" = x'y' ]; then
    base_upgrade_command="${base_upgrade_command} -F"
  fi

  management_ip=`zstack-ctl show_configuration | awk '/management.server.ip/{ print $3 }'`
  if [ ! -z "$management_ip" ]; then
    base_upgrade_command="${base_upgrade_command} -I ${management_ip}"
  fi
  $base_upgrade_command

  if [ $? -eq 0 ]; then
    # if mevoco-1.x before upgrade, then install-license manually, and restart mn
    if [ x"$NEED_INSTALL_TRIAL_LICENSE" = x'y' -a ! -f /var/lib/zstack/license/license.txt -a -f /var/lib/zstack/license/zstack_trial_license ]; then
      zstack-ctl install_license --license /var/lib/zstack/license/zstack_trial_license
      zstack-ctl restart_node
    fi
    [ ! -f /usr/lib/systemd/system/getty@.service ] && return
    rm -f /etc/systemd/system/getty.target.wants/getty@tty1.service
    rm -f /usr/lib/systemd/system/getty@.service
    systemctl enable getty@tty1.service
    systemctl enable getty@tty2.service
    systemctl enable getty@tty3.service
    systemctl daemon-reload
    systemctl restart getty@tty1.service
    systemctl restart getty@tty2.service
    systemctl restart getty@tty3.service
  else
    echo "Failed to upgrade ZStack!"
    exit -1
  fi
}

echo $(tput bold)"=== Begin to Upgrade ${PRODUCT_NAME}  ==="$(tput sgr0) | tee $UPGRADE_LOG

# upgrate zstack local repo
get_check_iso
make_tmp_repo
INNER_UPGRADE_SCRIPT="/opt/tmp-iso/upgrade_repo.sh"
if [ -f "$INNER_UPGRADE_SCRIPT" ]; then
  bash "$INNER_UPGRADE_SCRIPT" "$STEP" "$REPO_ONLY" "$FORCE_UPGRADE_DB" "$ISO_PATH" "$FORCE_UPGRADE_REPO" "$ADD_REPO" || exit 1
else
  upgrade_repo
  do_ks_post
  install_zstack_manager
  if [ x"$REPO_ONLY" = x'n' ]; then
    upgrade_zstack
  fi
fi

echo -e $(tput bold)"\n=== Successfully Upgrade ${PRODUCT_NAME}! ==="$(tput sgr0) | tee -a $UPGRADE_LOG