#!/bin/sh -e

EFIDIR=$(grep ^ID= /etc/os-release | sed -e 's/^ID=//' -e 's/rhel/redhat/' -e 's/\"//g')
if [ -d /sys/firmware/efi/efivars/ ]; then
    grubdir=`echo "/boot/efi/EFI/${EFIDIR}/" | sed 's,//*,/,g'`
else
    grubdir=`echo "/boot/grub2" | sed 's,//*,/,g'`
fi

PACKAGE_VERSION="2.12"
PACKAGE_NAME="GRUB"
self=`basename $0`
bindir="/usr/bin"
grub_mkpasswd="${bindir}/grub2-mkpasswd-pbkdf2"

# Locale: choose Chinese or English for user-visible messages.
_is_zh_lang() {
    _l="${LC_ALL:-${LC_MESSAGES:-${LANG:-C}}}"
    case "$_l" in
        zh_CN*|zh_SG*) return 0 ;;
        *) return 1 ;;
    esac
}

# t "English" "中文" — prints one of two strings (no newline).
t() {
    if _is_zh_lang; then
        printf %s "$2"
    else
        printf %s "$1"
    fi
}

# Usage: usage
# Print the usage.
usage () {
    if _is_zh_lang; then
    cat <<EOF
用法: $0 [选项]
$0 将提示您为 GRUB 引导加载程序设置密码。密码将写入名为 user.cfg 的文件，
该文件默认位于 GRUB 目录：${grubdir}。

  -h, --help                     显示本说明并退出
  -v, --version                  显示版本信息并退出
  -o, --output_path <目录>       将 user.cfg 写入用户指定的目录

报告问题：https://bugzilla.redhat.com
EOF
    else
    cat <<EOF
Usage: $0 [OPTION]
$0 prompts the user to set a password on the grub bootloader. The password
is written to a file named user.cfg which lives in the GRUB directory
located by default at ${grubdir}.

  -h, --help                     print this message and exit
  -v, --version                  print the version information and exit
  -o, --output_path <DIRECTORY>  put user.cfg in a user-selected directory

Report bugs at https://bugzilla.redhat.com.
EOF
    fi
}

argument () {
    opt=$1
    shift

    if test $# -eq 0; then
        printf "$(t "%s: option requires an argument -- \`%s'\n" "%s：选项 \`%s' 需要参数\n")" "$self" "$opt" 1>&2
        exit 1
    fi
    echo $1
}

# Ensure that it's the root user running this script
if [ "${EUID}" -ne 0 ]; then
    echo "$(t "The grub bootloader password may only be set by root." "仅 root 用户可设置 GRUB 引导加载程序密码。")"
    usage
    exit 2
fi

# Check the arguments.
while test $# -gt 0
do
    option=$1
    shift

    case "$option" in
    -h | --help)
	usage
	exit 0 ;;
    -v | --version)
	echo "$self (${PACKAGE_NAME}) ${PACKAGE_VERSION}"
	exit 0 ;;
    -o | --output)
        OUTPUT_PATH=`argument $option "$@"`; shift ;;
    --output=*)
        OUTPUT_PATH=`echo "$option" | sed 's/--output=//'` ;;
    -o=*)
        OUTPUT_PATH=`echo "$option" | sed 's/-o=//'` ;;
    esac
done

# set user input or default path for user.cfg file
if [ -z "${OUTPUT_PATH}" ]; then
    OUTPUT_PATH="${grubdir}"
fi

if [ ! -d "${OUTPUT_PATH}" ]; then
    printf "$(t "%s does not exist.\n" "%s 不存在。\n")" "${OUTPUT_PATH}"
    usage
    exit 2;
fi

ttyopt=$(stty -g)
fixtty() {
      stty ${ttyopt}
}

trap fixtty EXIT

getsaltpass() {
	local P0
	local P1
	P0="$1" && shift
	P1="$1" && shift
	P2="$1" && shift

	( echo ${P0} ; echo ${P1} ) | \
    	LC_ALL=C ${grub_mkpasswd} -a ${P2} | \
    	grep -v '[eE]nter password:' | \
    	sed -e "s/PBKDF2 hash of your password is //"
}

verifyusercfgoldpasswd() {
    # get old password salt
    expectsalt=`cat ${grubdir}/user.cfg | cut -d "." -f 5`
    # get expect password
    expectpass=`cat ${grubdir}/user.cfg`
    prefix="GRUB2_PASSWORD="

    stty -echo
    echo -n "$(t "Enter Current password: " "请输入当前密码：")"
    read PASSWORD_CURRENT
    echo

    needcheckpass="${prefix}$(getsaltpass "${PASSWORD_CURRENT}" "${PASSWORD_CURRENT}" "${expectsalt}")"
    if [ "$expectpass" != "$needcheckpass" ]; then
        echo "$(t "Authentication failed" "身份验证失败")"
        exit 1
    fi

    stty ${ttyopt}
}

verifygrubcfgoldpasswd() {
    # get old password line
    expectpass=`cat ${grubdir}/grub.cfg | grep "password_pbkdf2 root grub.pbkdf2.sha512" | cut -d " " -f 3`
    # if not get password, try a quotation mark match
    if [ -z "$expectpass" ];then
        expectpass=`cat ${grubdir}/grub.cfg | grep "password_pbkdf2 root \"grub.pbkdf2.sha512" | cut -d " " -f 3 | cut -d "\"" -f 2`
    fi
    if [ -z "$expectpass" ];then
        expectpass=`cat ${grubdir}/grub.cfg | grep "password_pbkdf2 root 'grub.pbkdf2.sha512" | cut -d " " -f 3 | cut -d "'" -f 2`
    fi
    if [ -n "$expectpass" ];then
        # get old password salt
        expectsalt=`echo ${expectpass} | cut -d "." -f 5`
        stty -echo
        echo -n "$(t "Enter Current password: " "请输入当前密码：")"
        read PASSWORD_CURRENT
        echo

        needcheckpass="$(getsaltpass "${PASSWORD_CURRENT}" "${PASSWORD_CURRENT}" "${expectsalt}")"
        if [ "$expectpass" != "$needcheckpass" ]; then
            echo "$(t "Authentication failed" "身份验证失败")"
            exit 1
        fi
    fi

}

if [ -e ${grubdir}/user.cfg ];then
    verifyusercfgoldpasswd
else
    verifygrubcfgoldpasswd
fi

checkcomplexity() {
    set +e
    USERNAME=`cat ${grubdir}/grub.cfg | grep "set superusers=" | cut -d "\"" -f 2 |tail -1`
    local P1="$1" && shift
    if [ "$P1" = "$USERNAME" ];then
        echo "$(t "The password contains the user name in some form" "密码不能与超级用户名称相同")"
        exit 1
    fi
    # password len >= 8
    strlen=`echo "$P1" | grep -E '^(.{8,}).*$'`
    if [ -z "$strlen" ];then
        echo "$(t "The password is shorter than 8 characters" "密码长度不足 8 个字符")"
        exit 1
    fi
    # lowercase
    strlow=`echo "$P1" | grep -E --color '^(.*[a-z]+).*$'`
    # uppercase
    strupp=`echo $P1 | grep -E --color '^(.*[A-Z]).*$'`
    # special character
    strts=`echo $P1 | grep -E --color '^(.*\W).*$'`
    # num
    strnum=`echo $P1 | grep -E --color '^(.*[0-9]).*$'`
    complexity=0
    if [ -n "$strlow" ];then
        complexity=`expr $complexity + 1`
    fi
    if [ -n "$strupp" ];then
        complexity=`expr $complexity + 1`
    fi
    if [ -n "$strts" ];then
        complexity=`expr $complexity + 1`
    fi
    if [ -n "$strnum" ];then
        complexity=`expr $complexity + 1`
    fi
    if [ $complexity -lt 3 ];then
        echo "$(t "The password contains less than 3 character classes" "密码字符类别不足 3 种")"
        exit 1
    fi
    set -e
}

stty -echo

# prompt & confirm new grub2 root user password
echo -n "$(t "Enter password: " "请输入新密码：")"
read PASSWORD
echo
stty ${ttyopt}
checkcomplexity $PASSWORD
stty -echo
echo -n "$(t "Confirm password: " "请再次确认密码：")"
read PASSWORD_CONFIRM
echo
stty ${ttyopt}
checkcomplexity $PASSWORD_CONFIRM

getpass() {
    local P0
    local P1
    P0="$1" && shift
    P1="$1" && shift

    ( echo ${P0} ; echo ${P1} ) | \
        LC_ALL=C ${grub_mkpasswd} | \
        grep -v '[eE]nter password:' | \
        sed -e "s/PBKDF2 hash of your password is //"
}

MYPASS="$(getpass "${PASSWORD}" "${PASSWORD_CONFIRM}")"
if [ -z "${MYPASS}" ]; then
      printf "$(t "%s: error: empty password\n" "%s：错误：密码为空\n")" "${self}" 1>&2
      exit 1
fi

# on the ESP, these will fail to set the permissions, but it's okay because
# the directory is protected.
install -m 0600 /dev/null "${OUTPUT_PATH}/user.cfg" 2>/dev/null || :
chmod 0600 "${OUTPUT_PATH}/user.cfg" 2>/dev/null || :
echo "GRUB2_PASSWORD=${MYPASS}" > "${OUTPUT_PATH}/user.cfg"

if ! grep -q "^### BEGIN /etc/grub.d/01_users ###$" "${OUTPUT_PATH}/grub.cfg"; then
    echo "$(t "WARNING: The current configuration lacks password support!" "警告：当前 GRUB 配置未启用密码相关片段！")"
    echo "$(t "Update your configuration with grub2-mkconfig to support this feature." "请运行 grub2-mkconfig 重新生成配置以支持该功能。")"
fi

