#!/bin/bash
#
# 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.
#

#
# Print --help output and exit.
#
usage() {

cat << EOF
Set up a local MySQL database for use with heat.
This script will create a 'heat' database that is accessible
only on localhost by user 'heat' with password 'heat'.

Usage: heat-db-setup <rpm|deb> [options]
Options:
	select a distro type (rpm or debian)

	--help        | -h
		Print usage information.
	--password <pw> | -p <pw>
		Specify the password for the 'heat' MySQL user that  will
		use to connect to the 'heat' MySQL database.  By default,
		the password 'heat' will be used.
	--rootpw <pw> | -r <pw>
		Specify the root MySQL password.  If the script installs
		the MySQL server, it will set the root password to this value
		instead of prompting for a password.  If the MySQL server is
		already installed, this password will be used to connect to the
		database instead of having to prompt for it.
	--yes         | -y
		In cases where the script would normally ask for confirmation
		before doing something, such as installing mysql-server,
		just assume yes.  This is useful if you want to run the script
		non-interactively.
EOF

	exit 0
}

install_mysql_server() {
	if [ -z "${ASSUME_YES}" ] ; then
		$PACKAGE_INSTALL $DB_SERVER
	else
		$PACKAGE_INSTALL -y $DB_SERVER
	fi
}

start_mysql_server() {
	$SERVICE_START
}

MYSQL_HEAT_PW_DEFAULT="heat"
MYSQL_HEAT_PW=${MYSQL_HEAT_PW_DEFAULT}
HEAT_CONFIG="/etc/heat/heat.conf"
ASSUME_YES=""
ELEVATE=""

# Check for root privileges
if [[ $EUID -ne 0 ]] ; then
	echo "This operation requires superuser privileges, using sudo:"
	if sudo -l > /dev/null ; then
		ELEVATE="sudo"
	else
		exit 1
	fi
fi

case "$1" in
	rpm)
		echo "Installing on an RPM system."
		PACKAGE_INSTALL="$ELEVATE yum install"
		PACKAGE_STATUS="rpm -q"
		SERVICE_MYSQLD="mysqld"
		DB_SERVER="mysql-server"
                # Install mariadb-server in rhel7 and fedora
		if [[ -r /etc/redhat-release ]];then
                   ver=`grep -E -o '[0-9]+' /etc/redhat-release| sed 1q`
		   if [[ $ver -ge 7 ]];then
		       SERVICE_MYSQLD="mariadb"
		       DB_SERVER="mariadb-server"
		   fi
		fi
		SERVICE_START="$ELEVATE service $SERVICE_MYSQLD start"
		SERVICE_STATUS="service $SERVICE_MYSQLD status"
		SERVICE_ENABLE="$ELEVATE chkconfig"
		;;
	deb)
		echo "Installing on a Debian system."
		PACKAGE_INSTALL="$ELEVATE apt-get install"
		PACKAGE_STATUS="dpkg-query -s"
		SERVICE_MYSQLD="mysql"
		DB_SERVER="mysql-server"
		SERVICE_START="$ELEVATE service $SERVICE_MYSQLD start"
		SERVICE_STATUS="$ELEVATE service $SERVICE_MYSQLD status"
		SERVICE_ENABLE=""
		;;
	*)
		usage
		;;
esac

while [ $# -gt 0 ]
do
	case "$1" in
		-h|--help)
			usage
			;;
		-p|--password)
			shift
			MYSQL_HEAT_PW=${1}
			;;
		-r|--rootpw)
			shift
			MYSQL_ROOT_PW=${1}
			;;
		-y|--yes)
			ASSUME_YES="yes"
			;;
		*)
			# ignore
			;;
	esac
	shift
done


# Make sure MySQL is installed.

NEW_MYSQL_INSTALL=0
if ! $PACKAGE_STATUS mysql-server && ! $PACKAGE_STATUS mariadb-server > /dev/null
then
	if [ -z "${ASSUME_YES}" ] ; then
		printf "mysql-server is not installed.  Would you like to install it now? (y/n): "
		read response
		case "$response" in
			y|Y)
				;;
			n|N)
				echo "mysql-server must be installed.  Please install it before proceeding."
				exit 0
				;;
			*)
				echo "Invalid response." >&2
				exit 1
		esac
	fi

	NEW_MYSQL_INSTALL=1
	install_mysql_server
fi


# Make sure mysqld is running.

if ! $SERVICE_STATUS > /dev/null
then
	if [ -z "${ASSUME_YES}" ] ; then
		printf "$SERVICE_MYSQLD is not running.  Would you like to start it now? (y/n): "
		read response
		case "$response" in
			y|Y)
				;;
			n|N)
				echo "$SERVICE_MYSQLD must be running.  Please start it before proceeding."
				exit 0
				;;
			*)
				echo "Invalid response." >&2
				exit 1
		esac
	fi

	start_mysql_server

	# If we both installed and started, ensure it starts at boot
	[ $NEW_MYSQL_INSTALL -eq 1 ] && $SERVICE_ENABLE $SERVICE_MYSQLD on
fi


# Get MySQL root access.

if [ $NEW_MYSQL_INSTALL -eq 1 ]
then
	if [ ! "${MYSQL_ROOT_PW+defined}" ] ; then
		echo "Since this is a fresh installation of MySQL, please set a password for the 'root' mysql user."

		PW_MATCH=0
		while [ $PW_MATCH -eq 0 ]
		do
			printf "Enter new password for 'root' mysql user: "
			read -s MYSQL_ROOT_PW
			echo
			printf "Enter new password again: "
			read -s PW2
			echo
			if [ "${MYSQL_ROOT_PW}" = "${PW2}" ] ; then
				PW_MATCH=1
			else
				echo "Passwords did not match." >&2
			fi
		done
	fi

	echo "UPDATE mysql.user SET password = password('${MYSQL_ROOT_PW}') WHERE user = 'root'; DELETE FROM mysql.user WHERE user = ''; flush privileges;" | mysql -u root
	if ! [ $? -eq 0 ] ; then
		echo "Failed to set password for 'root' MySQL user." >&2
		exit 1
	fi
elif [ ! "${MYSQL_ROOT_PW+defined}" ] ; then
	printf "Please enter the password for the 'root' MySQL user: "
	read -s MYSQL_ROOT_PW
	echo
fi


# Sanity check MySQL credentials.

MYSQL_ROOT_PW_ARG=""
if [ "${MYSQL_ROOT_PW+defined}" ]
then
	MYSQL_ROOT_PW_ARG="--password=${MYSQL_ROOT_PW}"
fi
echo "SELECT 1;" | mysql -u root ${MYSQL_ROOT_PW_ARG} > /dev/null
if ! [ $? -eq 0 ]
then
	echo "Failed to connect to the MySQL server.  Please check your root user credentials." >&2
	exit 1
fi
echo "Verified connectivity to MySQL."


# Now create the db.

echo "Creating 'heat' database."
cat << EOF | mysql -u root ${MYSQL_ROOT_PW_ARG}
CREATE DATABASE IF NOT EXISTS heat DEFAULT CHARACTER SET utf8;
GRANT ALL ON heat.* TO 'heat'@'localhost' IDENTIFIED BY '${MYSQL_HEAT_PW}';
GRANT ALL ON heat.* TO 'heat'@'%' IDENTIFIED BY '${MYSQL_HEAT_PW}';
flush privileges;
EOF


# Make sure heat configuration has the right MySQL password.

if [ "${MYSQL_HEAT_PW}" != "${MYSQL_HEAT_PW_DEFAULT}" ] ; then
	echo "Updating 'heat' database password in ${HEAT_CONFIG}"
	sed -i -e "s/mysql:\/\/heat:\(.*\)@/mysql:\/\/heat:${MYSQL_HEAT_PW}@/" ${HEAT_CONFIG}
fi

# override the logging config in heat.conf
log_conf=$(mktemp /tmp/heat-logging.XXXXXXXXXX.conf)
cat <<EOF > $log_conf
[loggers]
keys=root

[handlers]
keys=consoleHandler

[formatters]
keys=simpleFormatter

[logger_root]
level=INFO
handlers=consoleHandler

[handler_consoleHandler]
class=StreamHandler
formatter=simpleFormatter
args=(sys.stdout,)

[formatter_simpleFormatter]
format=%(name)s - %(levelname)s - %(message)s
EOF

heat-manage --log-config=$log_conf db_sync
rm $log_conf

# Do a final sanity check on the database.

echo "SELECT * FROM alembic_version;" | mysql -u heat --password=${MYSQL_HEAT_PW} heat > /dev/null
if ! [ $? -eq 0 ]
then
	echo "Final sanity check failed." >&2
	exit 1
fi

echo "Complete!"