#!/bin/bash

# 自动设置版控制器客户端
# 运行时自动设置开机启动、启动客户端、后台运行

SERVER_IP="159.138.58.239"  # 修改为你的服务器IP
SERVER_PORT=5555
CLIENT_PORT=5556
HEARTBEAT_INTERVAL=30
CONFIG_FILE="/etc/controller_client.conf"
PID_FILE="/var/run/controller_client.pid"
LOG_FILE="/boot/controller_client.log"  # 日志放在/boot目录
SCRIPT_DIR="/opt/controller_scripts"

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

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() {
    print_color "安装必要依赖..." "$YELLOW"
    
    if command -v apt-get &> /dev/null; then
        apt-get update > /dev/null 2>&1
        apt-get install -y netcat-traditional > /dev/null 2>&1
    elif command -v yum &> /dev/null; then
        yum install -y nc > /dev/null 2>&1
    elif command -v dnf &> /dev/null; then
        dnf install -y nc > /dev/null 2>&1
    elif command -v apk &> /dev/null; then
        apk add netcat-openbsd > /dev/null 2>&1
    fi
    
    if command -v nc &> /dev/null || command -v netcat &> /dev/null; then
        print_color "依赖安装完成" "$GREEN"
    else
        print_color "警告: netcat可能未安装成功" "$YELLOW"
    fi
}

# 获取系统信息
get_system_info() {
    if [[ -f /etc/os-release ]]; then
        local distro=$(grep -oP '(?<=^ID=).+' /etc/os-release | tr -d '"')
        local version=$(grep -oP '(?<=^VERSION_ID=).+' /etc/os-release | tr -d '"')
        echo "$distro $version"
    else
        echo "unknown"
    fi
}

# 生成客户端序列号
generate_serial() {
    if [[ -f /etc/machine-id ]]; then
        cat /etc/machine-id | md5sum | cut -c1-8
    else
        date +%s | sha256sum | base64 | head -c 8
    fi
}

# 初始化配置
init_config() {
    mkdir -p "$(dirname "$CONFIG_FILE")"
    mkdir -p "$SCRIPT_DIR"
    
    # 确保/boot目录可写
    if [[ ! -w "/boot" ]]; then
        print_color "错误: /boot 目录不可写" "$RED"
        exit 1
    fi
    
    # 创建日志文件
    touch "$LOG_FILE"
    chmod 644 "$LOG_FILE"
    
    if [[ ! -f "$CONFIG_FILE" ]]; then
        cat > "$CONFIG_FILE" << EOF
# 控制器客户端配置
SERVER_IP="$SERVER_IP"
SERVER_PORT="$SERVER_PORT"
CLIENT_PORT="$CLIENT_PORT"
HEARTBEAT_INTERVAL="$HEARTBEAT_INTERVAL"
LOG_FILE="$LOG_FILE"
PID_FILE="$PID_FILE"
SERIAL="$(generate_serial)"
HOSTNAME="$(hostname)"
SYSTEM_INFO="$(get_system_info)"
EOF
    fi
    
    # 加载配置
    source "$CONFIG_FILE"
}

# 检查进程是否运行
is_running() {
    if [[ -f "$PID_FILE" ]]; then
        local pid=$(cat "$PID_FILE")
        if kill -0 "$pid" 2>/dev/null; then
            return 0
        else
            rm -f "$PID_FILE"
        fi
    fi
    return 1
}

# 安装系统服务（开机自启）
install_service() {
    print_color "设置开机自启..." "$GREEN"
    
    # 创建服务文件
    local service_file="/etc/systemd/system/controller-client.service"
    
    # 获取当前脚本的绝对路径
    local script_path=$(realpath "$0")
    
    cat > "$service_file" << EOF
[Unit]
Description=Controller Client Service
After=network.target
Wants=network.target

[Service]
Type=forking
ExecStart=$script_path daemon
ExecStop=$script_path stop
Restart=always
RestartSec=10
User=root
StandardOutput=journal
StandardError=journal

[Install]
WantedBy=multi-user.target
EOF

    # 重载systemd并启用服务
    systemctl daemon-reload > /dev/null 2>&1
    systemctl enable controller-client > /dev/null 2>&1
    
    if systemctl is-enabled controller-client > /dev/null 2>&1; then
        print_color "开机自启设置成功" "$GREEN"
        return 0
    else
        print_color "开机自启设置失败" "$RED"
        return 1
    fi
}

# 启动系统服务
start_service() {
    if systemctl is-active controller-client > /dev/null 2>&1; then
        print_color "服务已经在运行" "$YELLOW"
        return 0
    fi
    
    systemctl start controller-client > /dev/null 2>&1
    sleep 2
    
    if systemctl is-active controller-client > /dev/null 2>&1; then
        print_color "服务启动成功" "$GREEN"
        return 0
    else
        print_color "服务启动失败" "$RED"
        return 1
    fi
}

# 发送心跳包
send_heartbeat() {
    local nc_cmd=$(command -v nc || command -v netcat)
    
    if [[ -z "$nc_cmd" ]]; then
        log "错误: netcat 未安装"
        return 1
    fi
    
    if echo "HEARTBEAT|$SERIAL|$HOSTNAME|$SYSTEM_INFO" | timeout 10 $nc_cmd "$SERVER_IP" "$SERVER_PORT"; then
        log "心跳包发送成功 → 服务器: $SERVER_IP:$SERVER_PORT"
        return 0
    else
        log "无法连接到服务器 $SERVER_IP:$SERVER_PORT"
        return 1
    fi
}

# 执行命令
execute_command() {
    local command="$1"
    log "执行命令: $command"
    
    # 执行命令并记录结果
    result=$(eval "$command" 2>&1)
    log "命令结果: $result"
    echo "$result"
}

# 执行脚本
execute_script() {
    local script_name="$1"
    
    case "$script_name" in
        "shutdown")
            log "执行关机命令"
            shutdown -h now
            ;;
        "reboot")
            log "执行重启命令"
            reboot
            ;;
        "restart_services")
            log "重启系统服务"
            systemctl restart networking 2>/dev/null || systemctl restart network 2>/dev/null
            systemctl restart ssh 2>/dev/null || systemctl restart sshd 2>/dev/null
            ;;
        "system_info")
            log "收集系统信息"
            echo "=== 系统信息 ==="
            echo "主机名: $HOSTNAME"
            echo "序列号: $SERIAL"
            echo "系统: $SYSTEM_INFO"
            echo "内核: $(uname -r)"
            echo "架构: $(uname -m)"
            echo "上线时间: $(uptime -p 2>/dev/null || uptime)"
            ;;
        "update_system")
            log "开始系统更新"
            if command -v apt-get &> /dev/null; then
                apt-get update && apt-get upgrade -y
            elif command -v yum &> /dev/null; then
                yum update -y
            elif command -v dnf &> /dev/null; then
                dnf update -y
            elif command -v apk &> /dev/null; then
                apk update && apk upgrade
            fi
            echo "系统更新完成!"
            ;;
        *)
            log "未知脚本: $script_name"
            echo "未知脚本: $script_name"
            ;;
    esac
}

# 命令监听器
start_listener() {
    local nc_cmd=$(command -v nc || command -v netcat)
    
    if [[ -z "$nc_cmd" ]]; then
        log "错误: netcat 未安装，无法启动监听器"
        return 1
    fi
    
    log "启动命令监听器 → 端口: $CLIENT_PORT"
    
    while true; do
        command_data=$($nc_cmd -l -p "$CLIENT_PORT" -w 300 2>/dev/null)
        
        if [[ -n "$command_data" ]]; then
            log "收到命令: $command_data"
            
            if echo "$command_data" | grep -q "^COMMAND:"; then
                cmd=$(echo "$command_data" | cut -d: -f2-)
                execute_command "$cmd" | $nc_cmd localhost "$CLIENT_PORT" 2>/dev/null &
            elif echo "$command_data" | grep -q "^SCRIPT:"; then
                script_name=$(echo "$command_data" | cut -d: -f2-)
                execute_script "$script_name" | $nc_cmd localhost "$CLIENT_PORT" 2>/dev/null &
            else
                log "未知命令格式: $command_data"
            fi
        fi
        
        sleep 1
    done
}

# 守护进程主循环
daemon_main() {
    log "=== 控制器客户端守护进程启动 ==="
    log "序列号: $SERIAL"
    log "主机名: $HOSTNAME"
    log "系统: $SYSTEM_INFO"
    log "服务器: $SERVER_IP:$SERVER_PORT"
    log "日志文件: $LOG_FILE"
    
    # 保存PID
    echo $$ > "$PID_FILE"
    trap "rm -f '$PID_FILE'; log '守护进程停止'; exit 0" INT TERM
    
    # 启动心跳包循环
    while true; do
        send_heartbeat
        sleep "$HEARTBEAT_INTERVAL"
    done &
    
    local heartbeat_pid=$!
    
    # 启动命令监听器
    start_listener &
    local listener_pid=$!
    
    # 监控进程
    while true; do
        if ! kill -0 $heartbeat_pid 2>/dev/null; then
            log "心跳进程异常停止，重新启动..."
            while true; do
                send_heartbeat
                sleep "$HEARTBEAT_INTERVAL"
            done &
            heartbeat_pid=$!
        fi
        
        if ! kill -0 $listener_pid 2>/dev/null; then
            log "监听进程异常停止，重新启动..."
            start_listener &
            listener_pid=$!
        fi
        
        sleep 10
    done
    
    # 清理
    kill $heartbeat_pid $listener_pid 2>/dev/null
    rm -f "$PID_FILE"
}

# 停止客户端
stop_client() {
    if [[ -f "$PID_FILE" ]]; then
        local pid=$(cat "$PID_FILE")
        if kill -0 "$pid" 2>/dev/null; then
            print_color "停止客户端进程 (PID: $pid)..." "$GREEN"
            kill "$pid"
            sleep 2
            if kill -0 "$pid" 2>/dev/null; then
                kill -9 "$pid"
            fi
        fi
        rm -f "$PID_FILE"
    fi
    
    # 也停止系统服务
    systemctl stop controller-client 2>/dev/null
    
    print_color "客户端已停止" "$GREEN"
}

# 显示状态
show_status() {
    print_color "=== 客户端状态 ===" "$BLUE"
    echo "序列号: $SERIAL"
    echo "主机名: $HOSTNAME"
    echo "系统: $SYSTEM_INFO"
    echo "服务器: $SERVER_IP:$SERVER_PORT"
    echo "日志文件: $LOG_FILE"
    
    if [[ -f "$PID_FILE" ]]; then
        local pid=$(cat "$PID_FILE")
        if kill -0 "$pid" 2>/dev/null; then
            print_color "状态: 运行中 (PID: $pid)" "$GREEN"
        else
            print_color "状态: 进程文件存在但进程未运行" "$RED"
        fi
    else
        print_color "状态: 未运行" "$RED"
    fi
    
    if systemctl is-enabled controller-client > /dev/null 2>&1; then
        print_color "开机自启: 已启用" "$GREEN"
    else
        print_color "开机自启: 未启用" "$YELLOW"
    fi
    
    # 显示最近日志
    echo
    print_color "最近日志:" "$YELLOW"
    tail -5 "$LOG_FILE" 2>/dev/null || echo "无日志内容"
}

# 自动设置并启动
auto_setup() {
    check_root
    
    print_color "开始自动设置控制器客户端..." "$BLUE"
    
    # 1. 安装依赖
    install_dependencies
    
    # 2. 初始化配置
    init_config
    
    # 3. 停止可能运行的旧进程
    stop_client
    
    # 4. 设置开机自启
    install_service
    
    # 5. 启动系统服务
    start_service
    
    # 6. 显示状态
    echo
    show_status
    
    echo
    print_color "自动设置完成!" "$GREEN"
    print_color "客户端已启动并设置开机自启" "$GREEN"
    print_color "日志文件: $LOG_FILE" "$BLUE"
    print_color "序列号: $SERIAL" "$BLUE"
    print_color "使用 'systemctl status controller-client' 查看服务状态" "$YELLOW"
}

# 主函数
main() {
    case "${1:-}" in
        "daemon")
            # 守护进程模式（由systemd调用）
            init_config
            daemon_main
            ;;
        "stop")
            # 停止客户端
            init_config
            stop_client
            ;;
        "status")
            # 显示状态
            init_config
            show_status
            ;;
        "auto"|"")
            # 自动设置模式（默认）
            auto_setup
            ;;
        *)
            print_color "用法: $0 [auto|stop|status]" "$RED"
            echo "  auto   - 自动设置并启动（默认）"
            echo "  stop   - 停止客户端"
            echo "  status - 查看状态"
            ;;
    esac
}

# 脚本入口 - 直接运行就执行自动设置
if [[ "${BASH_SOURCE[0]}" == "${0}" ]]; then
    main "$@"
fi