#!/bin/bash

# 服务器控制器脚本 - 监听端口5555
# 支持: Debian, Ubuntu, CentOS, RHEL, Fedora, Alpine

SERVER_PORT=5555
LOG_FILE="/var/log/controller_server.log"
CLIENTS_FILE="/var/lib/controller_clients.txt"
SCRIPT_DIR="/opt/controller_scripts"

# 颜色输出
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m' # No Color

log() {
    echo -e "$(date '+%Y-%m-%d %H:%M:%S') - $1" | tee -a "$LOG_FILE"
}

print_color() {
    echo -e "${2}${1}${NC}"
}

# 检查root权限
check_root() {
    if [[ $EUID -ne 0 ]]; then
        print_color "错误: 此脚本需要root权限运行" "$RED"
        exit 1
    fi
}

# 安装依赖
install_dependencies() {
    local distro=$(grep -oP '(?<=^ID=).+' /etc/os-release | tr -d '"')
    
    case $distro in
        debian|ubuntu)
            apt-get update
            apt-get install -y net-tools iproute2
            ;;
        centos|rhel|fedora)
            yum install -y net-tools iproute2 || dnf install -y net-tools iproute2
            ;;
        alpine)
            apk add net-tools iproute2
            ;;
        *)
            log "未知发行版，请手动安装net-tools和iproute2"
            ;;
    esac
}

# 初始化环境
initialize() {
    mkdir -p "$(dirname "$LOG_FILE")"
    mkdir -p "$(dirname "$CLIENTS_FILE")"
    mkdir -p "$SCRIPT_DIR"
    
    # 创建示例脚本
    create_sample_scripts
    
    touch "$LOG_FILE"
    touch "$CLIENTS_FILE"
    
    print_color "初始化完成!" "$GREEN"
}

create_sample_scripts() {
    # 关机脚本
    cat > "$SCRIPT_DIR/shutdown.sh" << 'EOF'
#!/bin/bash
echo "执行关机操作..."
shutdown -h now
EOF

    # 重启脚本
    cat > "$SCRIPT_DIR/reboot.sh" << 'EOF'
#!/bin/bash
echo "执行重启操作..."
reboot
EOF

    # 服务重启脚本
    cat > "$SCRIPT_DIR/restart_services.sh" << 'EOF'
#!/bin/bash
echo "重启网络服务..."
systemctl restart networking 2>/dev/null || systemctl restart network 2>/dev/null
echo "重启SSH服务..."
systemctl restart ssh 2>/dev/null || systemctl restart sshd 2>/dev/null
EOF

    # 无限重启脚本 (危险)
    cat > "$SCRIPT_DIR/reboot_loop.sh" << 'EOF'
#!/bin/bash
echo "设置无限重启 - 危险操作!"
(crontab -l 2>/dev/null; echo "*/5 * * * * /sbin/reboot") | crontab -
reboot
EOF

    # 取消无限重启
    cat > "$SCRIPT_DIR/cancel_loop.sh" << 'EOF'
#!/bin/bash
echo "取消无限重启..."
crontab -l | grep -v reboot | crontab -
echo "已取消计划重启任务"
EOF

    chmod +x "$SCRIPT_DIR"/*.sh
}

# 显示客户端列表
show_clients() {
    if [[ ! -s "$CLIENTS_FILE" ]]; then
        print_color "没有已连接的客户端" "$YELLOW"
        return
    fi
    
    print_color "已连接的客户端列表:" "$BLUE"
    print_color "序列号 | IP地址 | 主机名 | 系统信息 | 最后在线" "$BLUE"
    print_color "------------------------------------------------------------" "$BLUE"
    cat "$CLIENTS_FILE"
}

# 向客户端发送命令
send_command() {
    local serial=$1
    local command=$2
    
    if [[ -z "$serial" || -z "$command" ]]; then
        print_color "用法: send <序列号> <命令>" "$RED"
        return
    fi
    
    local client_info=$(grep "^$serial|" "$CLIENTS_FILE")
    if [[ -z "$client_info" ]]; then
        print_color "错误: 未找到序列号 $serial 的客户端" "$RED"
        return
    fi
    
    IFS='|' read -r serial ip hostname system last_seen <<< "$client_info"
    
    print_color "向客户端 $serial ($hostname) 发送命令: $command" "$GREEN"
    
    # 发送命令到客户端
    echo "COMMAND:$command" | nc -w 5 "$ip" 5556
    
    if [[ $? -eq 0 ]]; then
        print_color "命令发送成功!" "$GREEN"
    else
        print_color "命令发送失败!" "$RED"
    fi
}

# 广播命令到所有客户端
broadcast_command() {
    local command=$1
    
    if [[ -z "$command" ]]; then
        print_color "用法: broadcast <命令>" "$RED"
        return
    fi
    
    print_color "向所有客户端广播命令: $command" "$YELLOW"
    
    while IFS='|' read -r serial ip hostname system last_seen; do
        if [[ -n "$serial" ]]; then
            echo "向 $serial 发送命令..."
            echo "COMMAND:$command" | nc -w 3 "$ip" 5556
        fi
    done < "$CLIENTS_FILE"
}

# 执行本地脚本
execute_script() {
    local serial=$1
    local script_name=$2
    
    if [[ -z "$serial" || -z "$script_name" ]]; then
        print_color "用法: script <序列号> <脚本名>" "$RED"
        print_color "可用脚本:" "$BLUE"
        ls "$SCRIPT_DIR"/*.sh | xargs -n 1 basename
        return
    fi
    
    local script_path="$SCRIPT_DIR/$script_name"
    if [[ ! -f "$script_path" ]]; then
        print_color "错误: 脚本 $script_name 不存在" "$RED"
        return
    fi
    
    local client_info=$(grep "^$serial|" "$CLIENTS_FILE")
    if [[ -z "$client_info" ]]; then
        print_color "错误: 未找到序列号 $serial 的客户端" "$RED"
        return
    fi
    
    IFS='|' read -r serial ip hostname system last_seen <<< "$client_info"
    
    print_color "向客户端 $serial 发送脚本: $script_name" "$GREEN"
    
    # 发送脚本执行命令
    echo "SCRIPT:$script_name" | nc -w 5 "$ip" 5556
}

# 启动服务器
start_server() {
    print_color "启动控制器服务器在端口 $SERVER_PORT..." "$GREEN"
    print_color "服务器日志: $LOG_FILE" "$BLUE"
    print_color "客户端列表: $CLIENTS_FILE" "$BLUE"
    
    while true; do
        log "等待客户端连接..."
        nc -l -p $SERVER_PORT -c '
            client_ip=$(echo $SSH_CLIENT | awk "{print \$1}")
            read -r data
            timestamp=$(date "+%Y-%m-%d %H:%M:%S")
            
            if echo "$data" | grep -q "HEARTBEAT|"; then
                IFS="|" read -r heartbeat serial hostname system <<< "$data"
                # 更新或添加客户端信息
                grep -v "^$serial|" /var/lib/controller_clients.txt > /tmp/clients.tmp
                echo "$serial|$client_ip|$hostname|$system|$timestamp" >> /tmp/clients.tmp
                mv /tmp/clients.tmp /var/lib/controller_clients.txt
                echo "ACK:OK"
            else
                echo "ACK:UNKNOWN_COMMAND"
            fi
        '
        sleep 1
    done
}

# 显示帮助
show_help() {
    print_color "服务器控制器命令:" "$BLUE"
    echo "list       - 显示客户端列表"
    echo "send       - 向指定客户端发送命令"
    echo "broadcast  - 向所有客户端广播命令"
    echo "script     - 执行预定义脚本"
    echo "scripts    - 显示可用脚本"
    echo "start      - 启动服务器"
    echo "stop       - 停止服务器"
    echo "status     - 显示服务器状态"
    echo "help       - 显示此帮助"
}

# 主菜单
main_menu() {
    print_color "=== 跨平台服务器控制器 ===" "$BLUE"
    echo "1. 启动服务器"
    echo "2. 显示客户端列表"
    echo "3. 发送命令到客户端"
    echo "4. 广播命令到所有客户端"
    echo "5. 执行脚本"
    echo "6. 显示可用脚本"
    echo "7. 服务器状态"
    echo "8. 退出"
    
    read -p "请选择操作 [1-8]: " choice
    
    case $choice in
        1)
            start_server
            ;;
        2)
            show_clients
            ;;
        3)
            read -p "输入序列号: " serial
            read -p "输入命令: " command
            send_command "$serial" "$command"
            ;;
        4)
            read -p "输入广播命令: " command
            broadcast_command "$command"
            ;;
        5)
            read -p "输入序列号: " serial
            read -p "输入脚本名: " script
            execute_script "$serial" "$script"
            ;;
        6)
            print_color "可用脚本:" "$BLUE"
            ls "$SCRIPT_DIR"/*.sh | xargs -n 1 basename
            ;;
        7)
            print_color "服务器状态:" "$BLUE"
            pgrep -f "nc -l -p $SERVER_PORT" >/dev/null && \
                print_color "服务器运行中" "$GREEN" || \
                print_color "服务器未运行" "$RED"
            ;;
        8)
            print_color "再见!" "$GREEN"
            exit 0
            ;;
        *)
            print_color "无效选择!" "$RED"
            ;;
    esac
    
    echo
    main_menu
}

# 脚本入口
case "${1:-}" in
    start)
        check_root
        install_dependencies
        initialize
        start_server
        ;;
    *)
        check_root
        install_dependencies
        initialize
        main_menu
        ;;
esac