#!/bin/bash

LOG_FILE='./es_recovery.log'

_color_red=$(TERM=xterm-256color tput setaf 1)
_color_normal=$(TERM=xterm-256color tput sgr0)
_color_blue=$(TERM=xterm-256color tput setaf 4)
_color_green=$(TERM=xterm-256color tput setaf 2)
_color_yellow=$(TERM=xterm-256color tput setaf 3)

function log_important()
{
    local ts=$(date -Iseconds)
    # to log file
    echo -e "$ts [IMPORTANT] $@" >> $LOG_FILE
    # to console (stderr)
    >&2 echo -e "${_color_yellow}$ts [IMPORTANT] $@${_color_normal}"
    return 0
}

function log_info()
{
    local ts=$(date -Iseconds)
    # to log file
    echo -e "$ts [INFO] $@" >> $LOG_FILE
    # to console (stderr)
    >&2 echo -e "$ts [INFO] $@"
    return 0
}

function log_error()
{
    local ts=$(date -Iseconds)
    # to log file
    echo -e "$ts [ERROR] $@" >> $LOG_FILE
    # to console (stderr)
    >&2 echo -e "${_color_red}$ts [ERROR] $@${_color_normal}"
    return 0
}

function log_debug()
{
    local ts=$(date -Iseconds)
    if [ $_DEBUG -eq 1 ]; then
        # to log file
        echo -e "$ts [DEBUG] $@" >> $LOG_FILE
        # to console (stderr)
        >&2 echo -e "$ts [DEBUG] $@"
    fi
    return 0
}

function print_msg(){
    echo
    echo -n ${_color_yellow}
    printf '%.0s=' {1..100};echo
    cat << EOF
$@
EOF
    printf '%.0s=' {1..100};echo
    echo -n ${_color_normal}
}

function interactive_input() {
    local _inputFlag=">>>"
    read -p "${_color_green}$_inputFlag ${1}:${_color_normal} " _str
    local _ret=$2
    eval $_ret='$_str'
    return 0
}

function install_dependents_pkg(){
    if ! rpm -q jq &> /dev/null;then
        log_info "install Packages sshpass .."
        rpm -ivh Packages/sshpass-1.06-2.el7.x86_64.rpm
    fi
    if ! rpm -q sshpass &> /dev/null;then
        log_info "install Packages jq .."
        rpm -ivh Packages/jq-1.6-2.el7.x86_64.rpm Packages/oniguruma-6.8.2-1.el7.x86_64.rpm
    fi
}

function check_middleware_proxy(){
    if ! /usr/sbin/ss -lnpt | grep 30003 &> /dev/null;then
        log_error "middleware proxy需要开启，通过proxy访问容器ES实例，请在compose中调整 replicas为1 ..."
        exit 1000
    fi
    log_info "middleware proxy已开启, 监听端口30003"
}

function check_es_status(){
    if curl -s $ES_URL/_cat/health | grep green &> /dev/null;then
        log_info "Elasticsearch status is green, check OK ..."
    elif curl -s $ES_URL/_cat/health | grep red &> /dev/null;then
        log_info "Elasticsearch status check is red, recovery 异常 ..."
        exit 1001
    elif curl -s $ES_URL/_cat/health | grep yellow &> /dev/null;then
        if [ $DEPLOY_MODE == "single" ];then
            ## 单机状态ES 因有一个副本配置，不能进行分配到其他节点，故只有50% 的分片活跃，也就是主分片数量
            local result=$(curl -s $ES_URL/_cat/health | awk '$NF != "50.0%" {print $0}')
            if [ ! -z "$result" ];then
                log_important "Elasticsearch 单机部署模式，状态异常，请进行检查，result: $result ..."
            fi
            log_info "Elasticsearch status yellow, check OK ..."
        fi
    fi
}

function create_snapshot_warehouse(){
    if [ $(curl -sIL -w "%{http_code}\n" -o /dev/null --location --request GET "$ES_URL/_snapshot/macrowing") != 200 ];then
    	log_info "Elasticsearch snapshots warehouse macrowing not exist, create it ... "
        curl -L -X POST $ES_URL/_snapshot/macrowing -H 'Content-Type: application/json' -d '
        {
            "type": "fs",
            "settings": {
                "location": "edoc2_backup"
            }
        }
        '
    fi
}

function get_mac_backup_dir(){
    local mac_backup_dir=$(docker inspect $(docker ps | grep elasticsearch | awk '{print $1}') --format '{{range .Mounts}}{{if eq .Destination "/esdata"}}{{.Source}}{{end}}{{end}}')
    if [ -n "$mac_backup_dir" ];then
        local _ret=$1
        eval $_ret='$mac_backup_dir'
    else
        log_error "获取宿主机ES备份目录失败，请检查执行脚本机器是否运行ES容器。。"
        exit
    fi
}

function list_all_backup_snapshots(){
    log_info "list_all_backup_snapshots ..."
    local snapshot_json=$(curl -s ${ES_URL}/_snapshot/macrowing/_all)
    local len=$(echo ${snapshot_json} | jq '.snapshots|length')
    local end=$(($len-1))
    local snap_arr=()
    log_info "backup snapshots name and state ..."
    if [ "${len}" != 0 ];then
        for index in $(seq 0 ${end})
        do
            local name=$(echo ${snapshot_json} | jq ".snapshots[${index}].snapshot" | sed 's/"//g')
            local state=$(echo ${snapshot_json} | jq '.snapshots[0].state' | sed 's/"//g')
            snap_arr=( "${snap_arr[@]}" "$name:$state")
        done
        PS3="${_color_green}>>> 选择恢复使用的快照前面的 号码: ${_color_normal}"
        select i in ${snap_arr[@]}
        do
            log_info "选择使用恢复的快照文件为: ${i%:*}"
            local _ret=$1
            eval $_ret='${i%:*}'
            break
        done
    else
        log_error "not found snapshot file , 请确保备份的快照文件拷贝到 ES数据存储目录 的backup 目录下"
    fi
}

function check_recovery_status(){
    log_info "check_recovery_status ..."
    local indices=( $(curl -s ${ES_URL}/_cat/indices | awk '{print $3}' | xargs) )
    local finished=false
    while [[ "$finished" == "false" ]]; do
        for i in ${!indices[@]}
        do
            local recovery_status_json=
            recovery_status_json=$(curl -s ${ES_URL}/${indices[$i]}/_recovery)
            local len=$(echo ${recovery_status_json} | jq ".${indices[$i]}.shards|length")
            local end=$(($len-1)) 
            if [ ${len} -ne 0 ];then
                for shard in $(seq 0 ${end})
                do
                   if [ "$(echo ${recovery_status_json} | jq ".${indices[$i]}.shards[${shard}].stage" | sed 's/"//g')" != "DONE" ];then
                        log_important "${indices[$i]} index recovery not finished ..."
                        finished=false
                        break 2
                   fi
                done
            fi
            log_info "${indices[$i]} index recovery is finished ..."
            if [ $i -eq $((${#indices[@]}-1)) ];then
                finished=true
            fi
        done
        if [ "$finished" == "false" ];then
            log_info "sleep 30s, wait index recovery finished ... "
            sleep 30
        fi
    done
}

function select_es_recovery_mode(){
    print_msg "可以选择的恢复方式：
    1. partial_index
         只进行部分索引的恢复
    2. all_index
         所有的索引都进行恢复"

    PS3="${_color_green}>>> 输入你选择ES恢复模式前面的 号码: ${_color_normal}"
    select i in 'partial_index' 'all_index'
    do
        log_info "选择当前数据恢复的模式为: $i"
        local _ret=$1
        eval $_ret='$i'
        break
    done
}

function select_es_deploy_mode(){
    print_msg "选择当前需要恢复集群的部署方式：
    1. single
         ES 只是单机部署
    2. cluster
         多个ES节点组成集群"

    PS3="${_color_green}>>> 输入你选择ES部署模式方式前面的 号码: ${_color_normal}"
    select i in 'single' 'cluster'
    do
        log_info "选择当前集群的部署方式为: $i"
        local _ret=$1
        eval $_ret='$i'
        break
    done
}

# Turns a comma separated string into array
# Usage: css_to_array "a, b, c" arr
function css_to_array()
{
    local __arr=$2
    eval $__arr='( ${1//,/ } )'
    return 0
}

function assert_success(){
    if [ $? -ne 0 ]; then
        log_error "$@"
        return 1
    else
        log_info "$@"
        return 0
    fi
}

function close_index(){
    log_info "close index $1 ..."
    #curl -XPUT ${ES_URL}/_settings -H "Content-Type: application/json" -d '{"index.blocks.read_only_allow_delete": null}'
    if [ $(curl -sIL -w "%{http_code}\n" -o /dev/null -X POST "${ES_URL}/${1}/_close") == 200 ];then
        log_info "close index $1 is SUCCESS ..."
    else
        log_error "close index is FAILED ..."
        exit
    fi
}

function recovery_from_snapshot(){
    log_info "recovery_from_snapshot ..."
    create_snapshot_warehouse ${ES_URL}
    local recovery_mode
    select_es_recovery_mode recovery_mode
    local snapshot_name 
    list_all_backup_snapshots snapshot_name
    if [ "$recovery_mode"x == "partial_index"x ];then
        print_msg "当前系统使用的索引有：
        file_1
        synwatchfile_1
        hotword_1
        folder_1
        eformoptlog_1
        useranalyze_1
        filelog_1
        signlog_1
        orglog_1"
        interactive_input "请输入需要恢复的索引名字，多个索引可以使用 英文逗号隔开：" recovery_index
        close_index $recovery_index
        curl -s -X POST "${ES_URL}/_snapshot/macrowing/${snapshot_name}/_restore?wait_for_completion=true" -H 'Content-Type: application/json' -d '
        {
            "indices": "'"${recovery_index}"'"
        }'
        assert_success "$recovery_index 索引恢复检查 ..." || exit  
    elif [ "$recovery_mode"x == "all_index"x ];then
        close_index "_all"
        curl -s -X POST "${ES_URL}/_snapshot/macrowing/${snapshot_name}/_restore?wait_for_completion=true"
        assert_success "all_index 索引恢复检查 ..." || exit
    fi
}

function get_confirm_input(){
    interactive_input "是否以上条件都已满足(yes|no)" _input
    case $_input in
        "yes"|"Yes"|"YES"|"y"|"Y")
            log_info "你选择了 Yes ..."
        ;;
        "no"|"No"|"NO"|"n"|"N")
            log_info "你选择了 No，将退出脚本执行"
            exit 0
        ;;
        *)
            log_info "你的输入有误，请输入yes or no..."
            exit
        ;;
    esac
}

function recovery_or_recreate(){
    print_msg "选择是否在新的ES实例上恢复：
    1. origin
         在原有故障ES实例上恢复
    2. new 
         在新的ES实例上恢复"

    PS3="${_color_green}>>> 输入上面选项前面的 号码: ${_color_normal}"
    select i in 'origin' 'new'
    do
        log_info "选择当前集群的部署方式为: $i"
        local _ret=$1
        eval $_ret='$i'
        break
    done
}

function single_print_info(){
    local _sel
    recovery_or_recreate _sel
    if [ "$_sel"x == "origin"x ];then
        print_msg "单机ES恢复在原有ES实例上恢复的条件：
        1. 关闭运行的ecm 实例
            根据环境情况移除，如：docker stack rm indrive
        2. 拷贝先前的备份快照文件到 $ES_SNAPSHOT_BACKUP_DIR/backup 目录下
        3. 拷贝的文件确保 容器内的进程用户 作为其他人 也可读写
        4. 恢复过程会关闭所要恢复的索引"
    elif [ "$_sel"x == "new"x ];then
        print_msg "单机ES恢复在新ES实例上恢复的条件：
        1. 拷贝先前的备份快照文件到新集群的 $ES_SNAPSHOT_BACKUP_DIR/backup 目录下
        2. 拷贝的文件确保 容器内的进程 作为其他人 也可读写"
    fi
    get_confirm_input $_sel
}

function cluster_print_info(){
    local _sel
    recovery_or_recreate _sel
    if [ "$_sel"x == "origin"x ];then
        print_msg "cluster ES恢复在原有ES集群上恢复的条件：
        1. 关闭运行的ecm 实例
            根据环境情况移除，如：docker stack rm indrive
        2. 通过共享文件系统，如NFS 使 $ES_SNAPSHOT_BACKUP_DIR/backup 挂载到NFS同一个共享目录
        3. 拷贝先前的备份快照文件到： $ES_SNAPSHOT_BACKUP_DIR/backup 目录下，即：共享目录下
        4. 拷贝的文件确保 容器内的进程用户 作为其他人 也可读写
        5. 恢复过程会关闭所要恢复的索引"
    elif [ "$_sel"x == "new"x ];then
        print_msg "cluster ES恢复在新ES集群上恢复的条件：
        1. 通过共享文件系统，如NFS 使 $ES_SNAPSHOT_BACKUP_DIR/backup 挂载到NFS同一个共享目录
        2. 拷贝先前的备份快照文件到： $ES_SNAPSHOT_BACKUP_DIR/backup 目录下，即：共享目录下
        3. 拷贝的文件确保 容器内的进程 作为其他人 也可读写"
    fi

    get_confirm_input $_sel 
}

LOCAL_IP=192.168.251.179
ES_URL=http://$LOCAL_IP:30003

function main(){
    check_middleware_proxy
    install_dependents_pkg
    get_mac_backup_dir ES_SNAPSHOT_BACKUP_DIR
    select_es_deploy_mode DEPLOY_MODE
    create_snapshot_warehouse

    if [ "$DEPLOY_MODE"x == "single"x ];then
        single_print_info
        recovery_from_snapshot
    elif [ "$DEPLOY_MODE"x == "cluster"x ];then
        cluster_print_info
        recovery_from_snapshot
    fi

    check_recovery_status
    check_es_status
}

main
