#!/bin/env bash
##Description: precheck before tidb dr swtich over 

function f_syntax
{
    [ ! -z "$1" ] && echo ">>[$(date "+%Y%m%d.%H%M%S")] $1"
    echo " "
    echo "ksh ${ShellName}  [-clsname <clsname>] [-hostname <hostname>]"
    echo "  [-ip <ip>] [-port <port>] [-user <user>] [-passwd <passwd>]  "
    echo "  [-ignoreusers <user1 user2 user3>] "
    echo "  [-task <backup,revoke>]"
    echo " -priv_list <all|read|write>"
    echo " "
    echo " #sample:"
    echo " #sh ${ShellName} -clsname sa2 -IP 172.16.xxx -port 4100 -user abcuser1 -passwd xxx -USERLIST ccb,root,abcuser1 -task backupuser,backupgrant -filename -filename -force "
    echo " #sh ${ShellName} -clsname sa2 -ip 172.16.xxx -port 4100 -user abcuser1 -passwd xxx -userlist ccb1,test1,user,drainer -task backup,revoke -bakfile -rvkfile -force -revoke_userpriv write"
    exit 22
}

function f_PrintLog
{
    _LogFlag=$(echo "$1"|tr "a-z" "A-Z")
    [ $(echo "$1"|wc -L) -eq 4 ] && _LogFlag=$(echo "${_LogFlag}_")
    _LogMessage=$2
    declare LogFile="${LogFile:-"/tmp/${PShellName:-${ShellName}}.log"}"
    touch "${LogFile}" 2>/dev/null
    if [ -w "${LogFile}" ]
    then
        echo "$(date "+%Y%m%d.%H%M%S"):[${_LogFlag}]:${UserName}@${HostName}:${2}"|tee -a "${LogFile}";
    else
        echo "$(date "+%Y%m%d.%H%M%S"):[${_LogFlag}]:${UserName}@${HostName}:${2}";
    fi
}


###############################################################################
##Init##
ShellName="$(echo $0|awk -F / '{print $NF}')";ShellOption="$@";
WorkDir="$(echo $0|sed s/${ShellName}//g)";[ -z "${WorkDir}" ] && WorkDir=${PWD};cd ${WorkDir};WorkDir=${PWD}
LogDir="${WorkDir%/}/logs";mkdir -p ${LogDir};chmod ugo+rwx ${LogDir} 2>/dev/null
LogFile="${LogDir}/${ShellName}.$(date "+%Y%m%d.%H%M%S").log";LogLevel=3;
touch "${LogFile}" 2>/dev/null;chmod ugo+rw "${LogFile}" 2>/dev/null
declare Shellcksum=$(cksum ${WorkDir%/}/${ShellName} 2>/dev/null|awk '{printf "%d",$1}')
declare -i UserID=$(id -u);declare UserName=$(whoami);declare HostName=$(hostname);
###############################################################################
##Get inopt
[ -z "${ShellOption}" ] && f_syntax;
for inopt in ${ShellOption}
do
    case $(echo $inopt|tr a-z A-Z) in
    -HOSTNAME|-HOST) CurOpt="-HOSTNAME";continue;;
    -CLSNAME|-C) CurOpt="-CLSNAME";continue;;
    -USERLIST|-USERS) CurOpt="-USERLIST";continue;;
    -IP) CurOpt="-IP";continue;;
    -PORT) CurOpt="-PORT";continue;;
    -USER|-U) CurOpt="-USER";continue;;
    -PASSWORD|-PASSWD) CurOpt="-PASSWORD";continue;;
    -TASK|-TASKS) CurOpt="-TASK";continue;;
    -FILENAME|-F) CurOpt="-FILENAME";continue;;
    -BAKFILE) CurOpt="-BAKFILE";continue;;
    -RVKFILE) CurOpt="-RVKFILE";continue;;
    -REVOKE_USERPRIV) CurOpt="-REVOKEPRIV";continue;;
    -FORCE|--FORCE) CurOpt="-FORCE";declare FORCE="FORCE";continue;;
    -HELP|--HELP|-H) f_syntax;;
    -*) CurOpt="";continue;;
    esac
    case "${CurOpt}" in
    -HOSTNAME) declare LocalHost="${inopt}";continue;;
    -CLSNAME) declare clsname="${inopt}";continue;;
    -USERLIST) declare userlist="${inopt}";continue;;
    -IP) declare tidbip="${inopt}";continue;;
    -PORT) declare tidbport="${inopt}";continue;;
    -USER) declare tidbuser="${inopt}";continue;;
    -PASSWORD) declare tidbpasswd="${inopt}";continue;;
    -TASK) declare tasks=$(echo "${tasks} ${inopt}"|tr "A-Z," "a-z ");continue;;
    -FILENAME) declare filename="${inopt}";continue;;
    -BAKFILE) declare bakfile="${inopt}";continue;;
    -RVKFILE) declare rvkfile="${inopt}";continue;;
    -REVOKEPRIV) declare revokepriv="${inopt}";continue;;
    esac
done

f_PrintLog "Info" "Start to run Script ${ShellName}, the Shellcksum is ${Shellcksum}"
f_PrintLog "Debug" "Script Options:${ShellOption}"

function f_runsqlfile
{
    _sqlfile=$1

    f_PrintLog "Info" "START TO RUN SQL:${_sqlfile}"
    echo "*************************************************************"
    cat ${_sqlfile}
    echo "*************************************************************"
    f_PrintLog "Info" "PLS CHECK ALL COMMANDS IN FILE, SLEEP 10......"
    sleep 10
    echo "mysql -h${tidbip} -P${tidbport} -u${tidbuser} -p${tidbpasswd:-123} < ${_sqlfile}"
    mysql -h${tidbip} -P${tidbport} -u${tidbuser} -p${tidbpasswd:-123} -vvv < ${_sqlfile}
    return $?
}


function f_usergrantsbak
{
    f_PrintLog "Info" "Start to run f_usergrantsbak."
    #提供用户名清单\备份文件名,默认备份全部用户权限
    local _filename=$1
    local _userlist=$2
    
    #backup file
    if [ ! -f ${_filename} ]
    then 
        touch ${_filename}
    elif [ -w ${_filename} ]
    then
        f_PrintLog "Warn" "backup file ${_filename} exist, clean content now."
        sleep 3
        ls -l ${_filename}
        [ "${FORCE}" = "FORCE" ] && echo "--start backup user grants. $(date +%Y%m%d.%H%M%S)">${_filename}
    else
        ls -l ${_filename}
        f_PrintLog "Error" "file ${_filename} can not write, quit!!!"
        exit 1
    fi

    #get user info
    if [ -z "${_userlist}" ]
    then
        _execsql="select concat(user,'@',host) from  mysql.user"
        echo "mysql -h${tidbip} -P${tidbport} -u${tidbuser} -p123 -NBe "${_execsql}""
        _userlist=$(mysql -h${tidbip} -P${tidbport} -u${tidbuser} -p${tidbpasswd:-123} -NBe "${_execsql}")
    else
        _userlist=$(echo "${_userlist}"|tr "," " ")
    fi
    f_PrintLog "Info" "start to backup user grants, userlist:${_userlist}."
    
    #backup grants info
    for _user in ${_userlist}
    do
        if [ -z "$(echo ${_user}|grep "@")" ] 
        then 
            _user="${_user}@'%'"
        else
            _user="${_user%%@*}@'${_user##*@}'"
        fi

        echo "--${_user}" | tee -a "${_filename}"
        _execsql="SHOW GRANTS FOR ${_user}"
        echo "mysql -h${tidbip} -P${tidbport} -u${tidbuser} -p${tidbpasswd:-123} -NBe ${_execsql} "
        mysql -h${tidbip} -P${tidbport} -u${tidbuser} -p${tidbpasswd:-123}  -NBe "${_execsql}" |tee -a "${_filename}"
    done

    #add ;
    sed -i 's/^GRANT.*/&;/g' ${_filename}
    f_PrintLog "Info" "backup user grants finish, backup file is ${_filename}"
    ls -l ${_filename}
}

function f_generatesqlfile
{
f_PrintLog "Info" "Start to run f_generatesqlfile."
f_PrintLog "Info" "parameters: $1 $2 $3 $4"

python -c "
import re,sys

args = sys.argv
infile = args[1]
outfile = args[2]
userlist = args[3]
rvktype = args[4]

userlist_lst = userlist.split(',')
revoke_lst = []

for username in userlist_lst:
    with open(infile, mode='r') as f:
        for line in f:
            if not re.search(username, line):
                continue
            elif line.startswith('--'):
                revoke_lst.append(line)
                continue
            elif line.startswith('GRANT USAGE ON'):
                continue
            else:
                line_user = re.findall(r'[\']*\w+[\']*@[\'\w\%]*', line)
                if not re.search(username, line_user[0]):
                    continue

                revoke_all = re.sub(r'(GRANT )([\w, .*]+)( TO )([ \']*\w+[\']*@.*)', r'REVOKE \2 FROM \4', line)
                revoke_read = re.sub(r'(GRANT [\w, .*]+ ON )(.* TO [ \']*\w+[\']*@.*)', r'REVOKE select FROM \2', line)
                grant_read = re.sub(r'(GRANT [\w, ]+)( ON .* TO .*)', r'GRANT select \2', line, re.IGNORECASE)
            
            if rvktype == 'all':
                revoke_lst.append(revoke_all.replace('WITH GRANT OPTION', ''))
            elif rvktype == 'write':
                revoke_lst.append(revoke_all.replace('WITH GRANT OPTION', ''))
                revoke_lst.append(grant_read)
            elif rvktype == 'read':
                revoke_lst.append(revoke_read.replace('WITH GRANT OPTION', ''))

with open(outfile, mode='w') as f:
     f.write(''.join(revoke_lst))
" "${1}" "${2}" "${3}" "${4}"
f_PrintLog "Info" "Start generate revoke sqlfile finish, sqlfile: ${2}"
ls -l ${2}
#cat ${2}
}

function f_userrevoke
{
    f_PrintLog "Info" "Start to run f_userrevoke."
    #提供用户名清单\全部回收\保留只读\保留连库
    #all revoke all 
    #read revoke select
    #write revoke all+grant select 
    _rvkfile=$1
    _userlist=$2
    _rvkpriv=$3

    case "${_rvkpriv}" in 
        all|read|write) f_PrintLog "Info" "start revoke ${_revokepriv} from ${_userlist}";;
        *)f_PrintLog "Error" "revoke priveleges must be <all|read|write>";;
    esac

    f_PrintLog "Info" "Make and clean grant file ${_rvkfile}"
    echo "">${_rvkfile}_TMP
    echo "">${_rvkfile}

    #backup current user grants
    f_usergrantsbak "${_rvkfile}_TMP" "${_userlist}"
    
    #generate revoke sqls
    f_generatesqlfile "${_rvkfile}_TMP" "${_rvkfile}" "${_userlist}" "${_rvkpriv}"
    
    #run sqlfile
    f_runsqlfile "${_rvkfile}"
    _rt_code=$?
    if [ ${_rt_code} -eq 0 ]
    then
        f_PrintLog "Info" "revoke user privileges finish!"
    else
        f_PrintLog "Error" "revoke user privileges failed!!!"
        exit 1
    fi

    sleep 30
    # f_runsqlfile "${_rvkfile}_TMP"
}

function f_usergrant
{
    f_PrintLog "Info" "Start to run f_usergrant."
    #提供用户名清单\权限备份文件
    return
}

###################
####    Main   ####
###################

#backup must have a file
[ -z "${bakfile}" -a -n "${filename}" ] && bakfile=${filename}
[ -z "${bakfile}" ] && bakfile="tidb_usergrants_${tidbip}_$(date +%Y%m%d).sql"
if [ -s "${bakfile}" -a "${FORCE}" != "FORCE" -a -n "$(echo ${tasks} | grep -i backup)" ]
then
    ls -l ${bakfile}
    f_PrintLog "Error" "file ${bakfile} already exists, for backup task pls use -force option or empty the file!"
    f_syntax
fi

if [ -n "$(echo ${tasks} | grep -i backup)" -a -z "$(echo ${tasks} | grep -i "^backup")" ]
then
    f_PrintLog "Error" "backup must be the first task in tasklist:${tasks}!"
    f_syntax
fi


#revoke must have a file
[ -z "${rvkfile}" -a -n "${filename}" -a -z "$(echo ${tasks} | grep -i backup)" ] && rvkfile=${filename}
[ -z "${rvkfile}" -a -n "$(echo ${tasks} | grep -i revoke)" ] && rvkfile="tidb_userrevoke_${tidbip}_$(date +%Y%m%d).sql"
#revoke must have a user
#revoke must have a revoke scope(all|read|write) 
if [ -z "${userlist}" -a -n "$(echo ${tasks} | grep -i revoke)" ]
then
    f_PrintLog "Error" "action revoke need a userlist, pls provide them through -userlist <user1,user2>!"
    f_syntax
elif [ -z "${revokepriv}" -a -n "$(echo ${tasks} | grep -i revoke)" ]
then
    f_PrintLog "Error" "action revoke need a revoke type, pls provide it through -revoke_userpriv <read|write|all>!"
    f_syntax
fi

#
for task in ${tasks}
do 
    case "${task}" in 
    backup) f_usergrantsbak "${bakfile}" "${userlist}";;
    #grant) f_usergrant ;;
    revoke) f_userrevoke "${rvkfile}" "${userlist}" "${revokepriv}";;
    *) f_PrintLog "Error" "Invalid task:${task}, pls input the right tasks!!!" ;;
    esac
done