#!/bin/bash

# 检查是否以 root 用户运行
if [ "$EUID" -ne 0 ]; then
    echo "此脚本需要以管理员权限运行，请使用 sudo 执行。"
    exit 1
fi

# 初始化变量
DOUAUTH_IS=0
DOUAUTH_PAM_CONFIG=0
if grep -E -q '^[[:space:]]*auth[[:space:]]+sufficient[[:space:]]+pam_ukeyroot\.so[[:space:]]*$' /etc/pam.d/system-auth; then
    DOUAUTH_PAM_CONFIG=1
fi
# 解析命令行选项
while getopts "i:ch" opt; do
  case $opt in
    i)
      if [ "$OPTARG" = "0" ] || [ "$OPTARG" = "1" ]; then
        DOUAUTH_IS=$OPTARG
      else
        echo "错误: -d 选项的值必须为 0 或 1。"
        exit 1
      fi
      ;;
    c)
        if [ $DOUAUTH_PAM_CONFIG -eq 0 ]; then
            echo "本地双因子认证未配置"
            exit 1
        else
            echo "本地双因子认证已配置"
            exit 1
        fi
      ;;
    h)
        echo "用法: $0 [-i 0|1] [-c]"
        exit 1
        ;;

  esac
done

# 检查是否没有传入任何选项参数
if [ $OPTIND -eq 1 ]; then
    echo "用法: $0 [-i 0|1] [-c]"
    exit 1
fi

OTP_PAM_sufficient=0
if grep -E -q '^[[:space:]]*auth[[:space:]]+sufficient[[:space:]]+pam_otp\.so[[:space:]]*$' /etc/pam.d/system-auth; then
    OTP_PAM_sufficient=1
fi

OTP_PAM_required=0
if grep -E -q '^[[:space:]]*auth[[:space:]]+required[[:space:]]+pam_otp\.so[[:space:]]*$' /etc/pam.d/system-auth; then
    OTP_PAM_required=1
fi

DOUAUTH_PAM_CONFIG=0
if grep -E -q '^[[:space:]]*auth[[:space:]]+sufficient[[:space:]]+pam_ukeyroot\.so[[:space:]]*$' /etc/pam.d/system-auth; then
    DOUAUTH_PAM_CONFIG=1
fi

#OTP未配置
if [ $OTP_PAM_sufficient -eq 0 ] && [ $OTP_PAM_required -eq 0 ] && [ $DOUAUTH_PAM_CONFIG -eq 0 ]; then

    if [ "$DOUAUTH_IS" -eq 1 ]; then
        sed -i '/^auth/s/auth[ ]*sufficient[ ]*pam_unix.so[ ]*nullok[ ]*try_first_pass/auth  required   pam_unix.so   nullok   try_first_pass/' /etc/pam.d/system-auth
        sed -i '/auth * required *  pam_unix.so  * nullok  * try_first_pass/aauth     sufficient    pam_ukeyroot.so' /etc/pam.d/system-auth
        sed -i '/auth   *  sufficient  *  pam_ukeyroot.so/aauth     sufficient    pam_ukey.so' /etc/pam.d/system-auth

    fi

elif [ $OTP_PAM_sufficient -eq 1 ] && [ $OTP_PAM_required -eq 0 ] && [ $DOUAUTH_PAM_CONFIG -eq 0 ]; then
    if [ "$DOUAUTH_IS" -eq 1 ]; then
        sed -i '/^auth[[:space:]]\+sufficient[[:space:]]\+pam_otp.so/a auth     sufficient    pam_ukeyroot.so' /etc/pam.d/system-auth
        sed -i '/auth   *  sufficient  *  pam_ukeyroot.so/aauth     sufficient    pam_ukey.so' /etc/pam.d/system-auth
        sed -i 's/^auth[[:space:]]\+sufficient[[:space:]]\+pam_otp\.so$/auth required pam_otp.so/g' /etc/pam.d/system-auth
    fi

elif [ $OTP_PAM_sufficient -eq 0 ] && [ $OTP_PAM_required -eq 1 ] && [ $DOUAUTH_PAM_CONFIG -eq 1 ]; then
    if [ "$DOUAUTH_IS" -eq 0 ]; then
        sed -i 's/^auth[[:space:]]\+required[[:space:]]\+pam_otp\.so$/auth sufficient pam_otp.so/g' /etc/pam.d/system-auth
        sed -i '/auth * sufficient *  pam_ukeyroot.so/d' /etc/pam.d/system-auth
        sed -i '/auth   *  sufficient  *  pam_ukey.so/d' /etc/pam.d/system-auth
    fi
elif [ $OTP_PAM_sufficient -eq 0 ] && [ $OTP_PAM_required -eq 0 ] && [ $DOUAUTH_PAM_CONFIG -eq 1 ]; then

    if [ "$DOUAUTH_IS" -eq 0 ]; then
        sed -i '/^auth/s/auth[ ]*required[ ]*pam_unix.so[ ]*nullok[ ]*try_first_pass/auth  sufficient   pam_unix.so   nullok   try_first_pass/' /etc/pam.d/system-auth
        sed -i '/auth * sufficient *  pam_ukeyroot.so/d' /etc/pam.d/system-auth
        sed -i '/auth   *  sufficient  *  pam_ukey.so/d' /etc/pam.d/system-auth
    fi
elif [ $OTP_PAM_sufficient -eq 1 ] && [ $OTP_PAM_required -eq 0 ] && [ $DOUAUTH_PAM_CONFIG -eq 1 ]; then
    sed -i '/auth[[:space:]]\+sufficient[[:space:]]\+pam_otp.so/d' /etc/pam.d/system-auth
    sed -i '/^auth/s/auth[ ]*sufficient[ ]*pam_unix.so[ ]*nullok[ ]*try_first_pass/auth  required   pam_unix.so   nullok   try_first_pass/' /etc/pam.d/system-auth
    sed -i '/auth[[:space:]]\+required[[:space:]]\+pam_unix.so[[:space:]]\+nullok[[:space:]]\+try_first_pass/aauth     required    pam_otp.so' /etc/pam.d/system-auth

fi

#-------------------------------------------------------------------------------------------------------------------
if grep -E -q '^[[:space:]]*auth[[:space:]]+sufficient[[:space:]]+pam_otp\.so[[:space:]]*$' /etc/pam.d/system-auth; then
    OTP_PAM_sufficient=1
fi

if grep -E -q '^[[:space:]]*auth[[:space:]]+required[[:space:]]+pam_otp\.so[[:space:]]*$' /etc/pam.d/system-auth; then
    OTP_PAM_required=1
fi

if grep -E -q '^[[:space:]]*auth[[:space:]]+sufficient[[:space:]]+pam_ukeyroot\.so[[:space:]]*$' /etc/pam.d/system-auth; then
    DOUAUTH_PAM_CONFIG=1
fi

if [ $OTP_PAM_sufficient -eq 0 ] && [ $OTP_PAM_required -eq 1 ] && [ $DOUAUTH_PAM_CONFIG -eq 0 ]; then
    sed -i 's/^auth[[:space:]]\+required[[:space:]]\+pam_otp\.so$/auth sufficient pam_otp.so/g' /etc/pam.d/system-auth

fi

if [ $OTP_PAM_sufficient -eq 1 ] && [ $OTP_PAM_required -eq 0 ] && [ $DOUAUTH_PAM_CONFIG -eq 0 ]; then
    sed -i '/^auth/s/auth[ ]*sufficient[ ]*pam_unix.so[ ]*nullok[ ]*try_first_pass/auth  required   pam_unix.so   nullok   try_first_pass/' /etc/pam.d/system-auth
    sed -i 's/^auth[[:space:]]\+required[[:space:]]\+pam_otp\.so$/auth sufficient pam_otp.so/g' /etc/pam.d/system-auth
fi

if [ $OTP_PAM_required -eq 1 ]; then
    sed -i '/^auth/s/auth[ ]*sufficient[ ]*pam_unix.so[ ]*nullok[ ]*try_first_pass/auth  required   pam_unix.so   nullok   try_first_pass/' /etc/pam.d/system-auth
fi