#!/bin/bash

################################################################################
# Arch Linux 服务器安全加固脚本
# 版本: v2.0
# 日期: 2025-01-19
# 
# 针对Arch Linux优化的工具选择:
# - 防火墙: nftables (现代化，替代iptables) 或 firewalld (更易用)
# - 入侵防护: sshguard (轻量级，原生nftables支持) 或 crowdsec (社区驱动)
# - 包管理: pacman (Arch原生)
# - 安全工具: 优先使用官方仓库和AUR
################################################################################

export TERM=xterm-color
set +H
set +e

# 颜色定义
if command -v tput &> /dev/null; then
    RED=$(tput setaf 1 2>/dev/null || echo '\033[0;31m')
    GREEN=$(tput setaf 2 2>/dev/null || echo '\033[0;32m')
    YELLOW=$(tput setaf 3 2>/dev/null || echo '\033[1;33m')
    BLUE=$(tput setaf 4 2>/dev/null || echo '\033[0;34m')
    CYAN=$(tput setaf 6 2>/dev/null || echo '\033[0;36m')
    MAGENTA=$(tput setaf 5 2>/dev/null || echo '\033[0;35m')
    NC=$(tput sgr0 2>/dev/null || echo '\033[0m')
else
    RED='\033[0;31m'
    GREEN='\033[0;32m'
    YELLOW='\033[1;33m'
    BLUE='\033[0;34m'
    CYAN='\033[0;36m'
    MAGENTA='\033[0;35m'
    NC='\033[0m'
fi

# 日志和备份目录
LOG_FILE="/var/log/arch_hardening_$(date +%Y%m%d_%H%M%S).log"
BACKUP_DIR="/root/security_backup_$(date +%Y%m%d_%H%M%S)"
ROLLBACK_SCRIPT="$BACKUP_DIR/rollback.sh"

# 全局变量
FIREWALL_CHOICE=""
IPS_CHOICE=""

################################################################################
# 工具函数
################################################################################

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

log_info() {
    echo -e "${BLUE}[INFO]${NC} $1" | tee -a "$LOG_FILE"
}

log_warning() {
    echo -e "${YELLOW}[WARNING]${NC} $1" | tee -a "$LOG_FILE"
}

log_error() {
    echo -e "${RED}[ERROR]${NC} $1" | tee -a "$LOG_FILE"
}

log_success() {
    echo -e "${GREEN}[SUCCESS]${NC} $1" | tee -a "$LOG_FILE"
}

log_critical() {
    echo -e "${RED}[CRITICAL]${NC} $1" | tee -a "$LOG_FILE"
}

log_section() {
    echo -e "\n${CYAN}========================================${NC}" | tee -a "$LOG_FILE"
    echo -e "${CYAN}$1${NC}" | tee -a "$LOG_FILE"
    echo -e "${CYAN}========================================${NC}\n" | tee -a "$LOG_FILE"
}

check_root() {
    if [[ $EUID -ne 0 ]]; then
        log_error "此脚本必须以root权限运行"
        exit 1
    fi
}

backup_file() {
    local file=$1
    
    if [[ ! -f "$file" ]]; then
        log_warning "文件不存在，跳过备份: $file"
        return 1
    fi
    
    mkdir -p "$BACKUP_DIR"
    
    local backup_name="$(basename $file).$(date +%Y%m%d_%H%M%S).bak"
    local backup_path="$BACKUP_DIR/$backup_name"
    
    if cp -p "$file" "$backup_path" 2>/dev/null; then
        if cmp -s "$file" "$backup_path"; then
            log_success "✓ 备份成功: $file -> $backup_name"
            echo "cp -p '$backup_path' '$file'" >> "$ROLLBACK_SCRIPT"
            return 0
        else
            log_error "✗ 备份验证失败: $file"
            rm -f "$backup_path"
            return 1
        fi
    else
        log_error "✗ 备份失败: $file"
        return 1
    fi
}

confirm_action() {
    local message=$1
    if [[ -t 0 ]]; then
        echo -e "${YELLOW}${message} [y/N]: ${NC}\c"
        read -n 1 -r
        echo
        if [[ ! $REPLY =~ ^[Yy]$ ]]; then
            return 1
        fi
        return 0
    else
        log_info "非交互模式: $message - 跳过"
        return 1
    fi
}

safe_execute() {
    local cmd="$1"
    local description="$2"
    local critical="${3:-false}"
    
    log_info "执行: $description"
    
    if eval "$cmd" >> "$LOG_FILE" 2>&1; then
        log_success "✓ 成功: $description"
        return 0
    else
        local exit_code=$?
        
        if [[ "$critical" == "true" ]]; then
            log_critical "✗ 关键操作失败: $description (退出码: $exit_code)"
            log_critical "脚本终止，可使用回滚: bash $ROLLBACK_SCRIPT"
            exit $exit_code
        else
            log_warning "⚠ 警告: $description 失败 (退出码: $exit_code)"
            return $exit_code
        fi
    fi
}

check_port_available() {
    local port=$1
    
    log_info "检查端口 $port 可用性..."
    
    if ss -tuln 2>/dev/null | grep -q ":$port " || netstat -tuln 2>/dev/null | grep -q ":$port "; then
        log_error "✗ 端口 $port 已被占用:"
        ss -tlnp 2>/dev/null | grep ":$port " || netstat -tlnp 2>/dev/null | grep ":$port "
        return 1
    fi
    
    log_success "✓ 端口 $port 可用"
    return 0
}

verify_ssh_keys_exist() {
    log_info "验证SSH密钥配置..."
    
    local has_keys=false
    local key_locations=()
    
    for user_home in /home/* /root; do
        local authorized_keys="$user_home/.ssh/authorized_keys"
        
        if [[ -f "$authorized_keys" ]]; then
            local key_count=$(grep -c "^ssh-" "$authorized_keys" 2>/dev/null || echo 0)
            
            if [[ $key_count -gt 0 ]]; then
                has_keys=true
                key_locations+=("$authorized_keys ($key_count 个密钥)")
                log_success "✓ 找到 $key_count 个SSH密钥: $authorized_keys"
            fi
        fi
    done
    
    if [[ "$has_keys" == true ]]; then
        log_success "✓ SSH密钥验证通过"
        printf '%s\n' "${key_locations[@]}"
        return 0
    else
        log_error "✗ 未找到任何SSH密钥配置"
        return 1
    fi
}

verify_ssh_config() {
    local ssh_config="/etc/ssh/sshd_config"
    
    log_info "验证SSH配置语法..."
    
    if sshd -t -f "$ssh_config" 2>&1 | tee -a "$LOG_FILE"; then
        log_success "✓ SSH配置语法正确"
        return 0
    else
        log_error "✗ SSH配置语法错误"
        return 1
    fi
}

################################################################################
# Arch Linux 特有函数
################################################################################

# 检查是否为Arch Linux
check_arch_linux() {
    if [[ ! -f /etc/arch-release ]]; then
        log_error "此脚本仅支持Arch Linux系统"
        exit 1
    fi
    
    log_success "✓ 检测到Arch Linux系统"
}

# 安装AUR助手（如果需要）
install_aur_helper() {
    if command -v yay &> /dev/null; then
        log_info "AUR助手已安装: yay"
        return 0
    fi
    
    if command -v paru &> /dev/null; then
        log_info "AUR助手已安装: paru"
        return 0
    fi
    
    if confirm_action "是否安装AUR助手 yay？（用于安装AUR包）"; then
        log_info "安装yay..."
        
        # 需要非root用户安装yay
        local build_user="nobody"
        if id "nobody" &>/dev/null; then
            cd /tmp
            safe_execute "pacman -S --needed --noconfirm git base-devel" "安装构建依赖" false
            safe_execute "git clone https://aur.archlinux.org/yay.git" "克隆yay仓库" false
            safe_execute "chown -R nobody:nobody /tmp/yay" "设置权限" false
            safe_execute "cd yay && sudo -u nobody makepkg -si --noconfirm" "构建yay" false
            cd -
        fi
    fi
}

################################################################################
# 主要加固功能
################################################################################

collect_system_info() {
    log_section "收集系统信息"
    
    log_info "操作系统: Arch Linux $(uname -r)"
    log_info "内核版本: $(uname -r)"
    log_info "主机名: $(hostname)"
    log_info "IP地址: $(hostname -I 2>/dev/null | awk '{print $1}' || echo "无法获取")"
    log_info "当前用户: $(whoami)"
    log_info "包管理器: pacman $(pacman --version | head -1)"
    
    echo "系统信息收集完成" >> "$LOG_FILE"
}

update_system() {
    log_section "系统更新"
    
    if confirm_action "是否执行系统更新？"; then
        # 更新pacman数据库
        safe_execute "pacman -Sy" "同步包数据库" false
        
        # 升级系统
        safe_execute "pacman -Su --noconfirm" "升级系统" false
        
        # 清理孤立包
        safe_execute "pacman -Rns \$(pacman -Qtdq) --noconfirm" "清理孤立包" false 2>/dev/null || true
        
        # 清理缓存
        safe_execute "pacman -Sc --noconfirm" "清理包缓存" false
        
        log_success "系统更新完成"
    else
        log_warning "跳过系统更新"
    fi
}

manage_users() {
    log_section "用户管理和SSH密钥配置"
    
    if ! confirm_action "是否创建专用管理用户？"; then
        log_warning "跳过用户创建"
        return 0
    fi
    
    local username=""
    
    if [[ -t 0 ]]; then
        echo -e "${YELLOW}请输入新用户名 (建议：admin、operator、sysadmin): ${NC}\c"
        read -p "" username
    else
        username="admin"
    fi
    
    if [[ -z "$username" ]] || [[ ! "$username" =~ ^[a-z_][a-z0-9_-]*$ ]]; then
        log_error "无效的用户名: $username"
        return 1
    fi
    
    if id "$username" &>/dev/null 2>&1; then
        log_info "用户 $username 已存在"
        if ! confirm_action "用户已存在，是否继续配置？"; then
            return 0
        fi
    else
        log_info "创建用户: $username"
        safe_execute "useradd -m -G wheel -s /bin/bash '$username'" "创建用户并添加到wheel组" false
        
        if [[ -t 0 ]]; then
            echo -e "${YELLOW}为用户 $username 设置密码: ${NC}"
            passwd "$username"
        else
            local temp_password=$(openssl rand -base64 32)
            echo "$username:$temp_password" | chpasswd
            log_warning "临时密码: $temp_password"
        fi
    fi
    
    # 配置SSH密钥
    local user_home=$(eval echo "~$username")
    local ssh_dir="$user_home/.ssh"
    local authorized_keys="$ssh_dir/authorized_keys"
    
    if [[ -t 0 ]]; then
        echo -e "${CYAN}选择SSH密钥配置方式:${NC}"
        echo "1) 粘贴现有公钥（推荐）"
        echo "2) 自动生成新密钥对"
        echo "3) 跳过"
        echo -e "${YELLOW}选择 [1-3]: ${NC}\c"
        read -p "" key_option
    else
        key_option="3"
    fi
    
    case $key_option in
        1)
            echo -e "${YELLOW}粘贴SSH公钥: ${NC}"
            read -p "" public_key
            
            if [[ -n "$public_key" ]] && [[ "$public_key" =~ ^ssh- ]]; then
                mkdir -p "$ssh_dir"
                echo "$public_key" >> "$authorized_keys"
                chown -R "$username:$username" "$ssh_dir"
                chmod 700 "$ssh_dir"
                chmod 600 "$authorized_keys"
                log_success "✓ SSH公钥已添加"
            fi
            ;;
        2)
            mkdir -p "$ssh_dir"
            chown "$username:$username" "$ssh_dir"
            
            local key_file="$ssh_dir/id_ed25519"
            
            if [[ ! -f "$key_file" ]]; then
                log_info "生成ED25519密钥对（更安全）..."
                su - "$username" -c "ssh-keygen -t ed25519 -f '$key_file' -N '' -C '$username@$(hostname)'"
                
                if [[ -f "$key_file" ]]; then
                    log_success "✓ SSH密钥对生成成功"
                    echo -e "${CYAN}━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━${NC}"
                    echo -e "${YELLOW}公钥内容:${NC}"
                    cat "$key_file.pub"
                    echo -e "${CYAN}━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━${NC}"
                fi
            fi
            ;;
    esac
    
    # 配置sudo权限
    configure_sudo_access "$username"
    
    if id "$username" &>/dev/null; then
        log_success "✓ 用户 $username 配置完成"
    fi
}

configure_sudo_access() {
    local username=$1
    
    log_info "配置sudo权限..."
    
    # 确保wheel组可以使用sudo
    if ! grep -q "^%wheel ALL=(ALL:ALL) ALL" /etc/sudoers; then
        backup_file "/etc/sudoers"
        echo "%wheel ALL=(ALL:ALL) ALL" >> /etc/sudoers
    fi
    
    if [[ -t 0 ]]; then
        echo -e "${CYAN}选择sudo配置:${NC}"
        echo "1) 需要密码（推荐）"
        echo "2) 无需密码"
        echo "3) 部分命令无需密码"
        echo -e "${YELLOW}选择 [1-3]: ${NC}\c"
        read -p "" sudo_option
    else
        sudo_option="1"
    fi
    
    local sudoers_file="/etc/sudoers.d/$username"
    
    case $sudo_option in
        1)
            echo "$username ALL=(ALL:ALL) ALL" > "$sudoers_file"
            log_success "✓ sudo需要密码"
            ;;
        2)
            echo "$username ALL=(ALL) NOPASSWD:ALL" > "$sudoers_file"
            log_warning "⚠ sudo无需密码"
            ;;
        3)
            cat > "$sudoers_file" << EOF
$username ALL=(ALL) NOPASSWD: /usr/bin/systemctl restart *, /usr/bin/systemctl status *
$username ALL=(ALL) NOPASSWD: /usr/bin/pacman -Syu, /usr/bin/pacman -Sy
$username ALL=(ALL) ALL
EOF
            log_success "✓ 部分命令无需密码"
            ;;
    esac
    
    chmod 440 "$sudoers_file"
    
    if visudo -c -f "$sudoers_file" >/dev/null 2>&1; then
        log_success "✓ sudoers配置验证通过"
    else
        log_error "✗ sudoers配置错误，已删除"
        rm -f "$sudoers_file"
    fi
}

harden_ssh() {
    log_section "SSH安全加固"
    
    local ssh_config="/etc/ssh/sshd_config"
    
    if [[ ! -f "$ssh_config" ]]; then
        log_error "SSH配置文件不存在"
        return 1
    fi
    
    backup_file "$ssh_config"
    
    # 禁用root登录
    if confirm_action "是否禁用SSH root登录？"; then
        sed -i 's/^#*PermitRootLogin.*/PermitRootLogin no/' "$ssh_config"
        log_success "✓ 已禁用root登录"
    fi
    
    # 禁用密码认证
    if confirm_action "是否禁用SSH密码认证？"; then
        echo -e "${YELLOW}━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━${NC}"
        echo -e "${RED}⚠ 警告：禁用密码认证后只能通过SSH密钥登录${NC}"
        echo -e "${YELLOW}━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━${NC}"
        
        if ! verify_ssh_keys_exist; then
            log_critical "✗ 未找到SSH密钥！"
            if confirm_action "是否现在配置SSH密钥？"; then
                manage_users
                if ! verify_ssh_keys_exist; then
                    log_error "密钥配置失败，跳过禁用密码认证"
                    return 1
                fi
            else
                log_warning "跳过禁用密码认证"
                return 0
            fi
        fi
        
        if confirm_action "确认禁用密码认证？"; then
            sed -i 's/^#*PasswordAuthentication.*/PasswordAuthentication no/' "$ssh_config"
            log_success "✓ 已禁用密码认证"
        fi
    fi
    
    # 修改SSH端口
    local new_port=""
    if [[ -t 0 ]]; then
        echo -e "${YELLOW}当前SSH端口: $(grep "^Port" "$ssh_config" 2>/dev/null | awk '{print $2}' || echo 22)${NC}"
        echo -e "${YELLOW}输入新端口号（1024-65535，留空跳过）: ${NC}\c"
        read -p "" new_port
        
        if [[ -n "$new_port" ]]; then
            if [[ "$new_port" =~ ^[0-9]+$ ]] && [[ "$new_port" -gt 1024 ]] && [[ "$new_port" -lt 65536 ]]; then
                if check_port_available "$new_port"; then
                    sed -i "s/^#*Port.*/Port $new_port/" "$ssh_config"
                    log_success "✓ SSH端口已修改为: $new_port"
                else
                    new_port=""
                fi
            fi
        fi
    fi
    
    # 其他安全配置
    cat >> "$ssh_config" << 'EOF'

# Arch Linux 安全加固配置
Protocol 2
MaxAuthTries 3
MaxSessions 2
LoginGraceTime 60
ClientAliveInterval 300
ClientAliveCountMax 2
PermitEmptyPasswords no
X11Forwarding no
UseDNS no
EOF
    
    if ! verify_ssh_config; then
        log_critical "✗ SSH配置验证失败！正在回滚..."
        local backup_file=$(ls -t "$BACKUP_DIR"/sshd_config*.bak 2>/dev/null | head -1)
        if [[ -n "$backup_file" ]]; then
            cp -p "$backup_file" "$ssh_config"
            log_success "✓ 已回滚"
        fi
        return 1
    fi
    
    log_success "✓ SSH配置验证通过"
    
    if confirm_action "是否立即重启SSH服务？"; then
        if systemctl restart sshd; then
            log_success "✓ SSH服务重启成功"
        fi
    fi
}

# 选择防火墙
choose_firewall() {
    log_section "选择防火墙方案"
    
    echo -e "${CYAN}Arch Linux 防火墙选项:${NC}"
    echo "1) nftables (现代化，推荐，原生内核支持)"
    echo "2) firewalld (易用，动态管理)"
    echo "3) ufw (简单，兼容iptables-nft)"
    echo "4) 跳过防火墙配置"
    
    if [[ -t 0 ]]; then
        echo -e "${YELLOW}选择 [1-4]: ${NC}\c"
        read -p "" fw_choice
    else
        fw_choice="1"
    fi
    
    case $fw_choice in
        1) FIREWALL_CHOICE="nftables" ;;
        2) FIREWALL_CHOICE="firewalld" ;;
        3) FIREWALL_CHOICE="ufw" ;;
        4) FIREWALL_CHOICE="none" ;;
        *) FIREWALL_CHOICE="nftables" ;;
    esac
    
    log_info "已选择防火墙: $FIREWALL_CHOICE"
}

# 配置nftables
configure_nftables() {
    log_section "配置nftables防火墙"
    
    # 安装nftables
    if ! command -v nft &> /dev/null; then
        safe_execute "pacman -S --noconfirm nftables" "安装nftables" true
    fi
    
    # 获取SSH端口
    local ssh_port=$(grep "^Port" /etc/ssh/sshd_config 2>/dev/null | awk '{print $2}')
    ssh_port=${ssh_port:-22}
    
    log_info "检测到SSH端口: $ssh_port"
    
    # 创建nftables配置
    local nft_conf="/etc/nftables.conf"
    backup_file "$nft_conf" 2>/dev/null || true
    
    cat > "$nft_conf" << EOF
#!/usr/bin/nft -f
# Arch Linux 防火墙配置

# 清空现有规则
flush ruleset

# 定义表
table inet filter {
    # 输入链
    chain input {
        type filter hook input priority 0; policy drop;
        
        # 允许环回接口
        iif lo accept
        
        # 允许已建立的连接
        ct state established,related accept
        
        # 允许SSH
        tcp dport $ssh_port accept
        
        # 允许HTTP/HTTPS (如果需要)
        # tcp dport { 80, 443 } accept
        
        # 允许ICMP (ping)
        icmp type echo-request limit rate 1/second accept
        icmpv6 type { echo-request, nd-neighbor-solicit, nd-router-advert, nd-neighbor-advert } accept
        
        # 记录并拒绝其他流量
        counter drop
    }
    
    # 转发链
    chain forward {
        type filter hook forward priority 0; policy drop;
    }
    
    # 输出链
    chain output {
        type filter hook output priority 0; policy accept;
    }
}
EOF
    
    # 询问是否开放其他端口
    if confirm_action "是否开放HTTP(80)？"; then
        safe_execute "ufw allow 80/tcp comment 'HTTP'" "开放HTTP" false
    fi
    
    if confirm_action "是否开放HTTPS(443)？"; then
        safe_execute "ufw allow 443/tcp comment 'HTTPS'" "开放HTTPS" false
    fi
    
    echo "y" | ufw enable >> "$LOG_FILE" 2>&1
    
    ufw status verbose
    
    # 启用UFW服务
    safe_execute "systemctl enable --now ufw.service" "启用UFW服务" false
    
    log_success "✓ UFW配置完成"
}

# 防火墙配置入口
configure_firewall() {
    choose_firewall
    
    case $FIREWALL_CHOICE in
        "nftables")
            configure_nftables
            ;;
        "firewalld")
            configure_firewalld
            ;;
        "ufw")
            configure_ufw
            ;;
        "none")
            log_warning "跳过防火墙配置"
            ;;
    esac
}

# 选择入侵防护系统
choose_ips() {
    log_section "选择入侵防护系统"
    
    echo -e "${CYAN}Arch Linux IPS选项:${NC}"
    echo "1) sshguard (轻量级，原生nftables/firewalld支持，推荐)"
    echo "2) crowdsec (现代化，社区驱动，支持多种防火墙)"
    echo "3) fail2ban (传统，功能全面)"
    echo "4) 跳过"
    
    if [[ -t 0 ]]; then
        echo -e "${YELLOW}选择 [1-4]: ${NC}\c"
        read -p "" ips_choice
    else
        ips_choice="1"
    fi
    
    case $ips_choice in
        1) IPS_CHOICE="sshguard" ;;
        2) IPS_CHOICE="crowdsec" ;;
        3) IPS_CHOICE="fail2ban" ;;
        4) IPS_CHOICE="none" ;;
        *) IPS_CHOICE="sshguard" ;;
    esac
    
    log_info "已选择IPS: $IPS_CHOICE"
}

# 配置sshguard
configure_sshguard() {
    log_section "配置SSHGuard"
    
    if ! command -v sshguard &> /dev/null; then
        safe_execute "pacman -S --noconfirm sshguard" "安装sshguard" false
    fi
    
    # 根据防火墙选择后端
    local backend=""
    case $FIREWALL_CHOICE in
        "nftables")
            backend="nftables"
            ;;
        "firewalld")
            backend="firewalld"
            ;;
        "ufw")
            backend="iptables"
            ;;
        *)
            backend="nftables"
            ;;
    esac
    
    log_info "使用防火墙后端: $backend"
    
    # 创建配置
    local sshguard_conf="/etc/sshguard.conf"
    backup_file "$sshguard_conf" 2>/dev/null || true
    
    cat > "$sshguard_conf" << EOF
# SSHGuard 配置
BACKEND="$backend"
LOGREADER="LANG=C /usr/bin/journalctl -afb -p info -n1 -t sshd -o cat"
THRESHOLD=30
BLOCK_TIME=3600
DETECTION_TIME=600
EOF
    
    # 启用服务
    safe_execute "systemctl enable --now sshguard.service" "启用sshguard" false
    
    # 检查状态
    systemctl status sshguard.service --no-pager || true
    
    log_success "✓ SSHGuard配置完成"
}

# 配置CrowdSec
configure_crowdsec() {
    log_section "配置CrowdSec"
    
    log_info "CrowdSec需要从AUR安装..."
    
    if command -v yay &> /dev/null; then
        safe_execute "yay -S --noconfirm crowdsec crowdsec-firewall-bouncer-nftables" "安装CrowdSec" false
    elif command -v paru &> /dev/null; then
        safe_execute "paru -S --noconfirm crowdsec crowdsec-firewall-bouncer-nftables" "安装CrowdSec" false
    else
        log_warning "未安装AUR助手，跳过CrowdSec安装"
        return 1
    fi
    
    # 启用服务
    safe_execute "systemctl enable --now crowdsec.service" "启用CrowdSec" false
    
    # 配置bouncer
    safe_execute "cscli bouncers add firewall-bouncer" "添加防火墙bouncer" false
    
    log_success "✓ CrowdSec配置完成"
}

# 配置Fail2ban
configure_fail2ban() {
    log_section "配置Fail2ban"
    
    if ! command -v fail2ban-client &> /dev/null; then
        safe_execute "pacman -S --noconfirm fail2ban" "安装fail2ban" false
    fi
    
    local ssh_port=$(grep "^Port" /etc/ssh/sshd_config 2>/dev/null | awk '{print $2}')
    ssh_port=${ssh_port:-ssh}
    
    local jail_local="/etc/fail2ban/jail.local"
    backup_file "$jail_local" 2>/dev/null || true
    
    cat > "$jail_local" << EOF
[DEFAULT]
bantime = 3600
findtime = 600
maxretry = 5

[sshd]
enabled = true
port = $ssh_port
filter = sshd
backend = systemd
maxretry = 3
bantime = 7200
EOF
    
    safe_execute "systemctl enable --now fail2ban.service" "启用fail2ban" false
    
    sleep 2
    fail2ban-client status 2>/dev/null || true
    
    log_success "✓ Fail2ban配置完成"
}

# IPS配置入口
configure_ips() {
    choose_ips
    
    case $IPS_CHOICE in
        "sshguard")
            configure_sshguard
            ;;
        "crowdsec")
            configure_crowdsec
            ;;
        "fail2ban")
            configure_fail2ban
            ;;
        "none")
            log_warning "跳过IPS配置"
            ;;
    esac
}

# 密码策略
harden_users() {
    log_section "密码策略加固"
    
    # 安装pam_pwquality
    if ! pacman -Q libpwquality &> /dev/null; then
        safe_execute "pacman -S --noconfirm libpwquality" "安装密码质量检查" false
    fi
    
    local pwquality_conf="/etc/security/pwquality.conf"
    backup_file "$pwquality_conf"
    
    cat >> "$pwquality_conf" << 'EOF'

# Arch Linux 密码策略
minlen = 12
dcredit = -1
ucredit = -1
lcredit = -1
ocredit = -1
maxrepeat = 3
EOF
    
    # 配置PAM
    local common_password="/etc/pam.d/passwd"
    if [[ -f "$common_password" ]]; then
        backup_file "$common_password"
        
        if ! grep -q "pam_pwquality.so" "$common_password"; then
            sed -i '/password.*required/a password required pam_pwquality.so retry=3' "$common_password"
        fi
    fi
    
    log_success "✓ 密码策略配置完成"
}

# 内核参数加固
harden_kernel() {
    log_section "内核参数加固"
    
    local sysctl_conf="/etc/sysctl.d/99-security.conf"
    backup_file "$sysctl_conf" 2>/dev/null || true
    
    cat > "$sysctl_conf" << 'EOF'
# Arch Linux 内核安全配置
net.ipv4.ip_forward = 0
net.ipv6.conf.all.forwarding = 0
net.ipv4.tcp_syncookies = 1
net.ipv4.conf.all.accept_redirects = 0
net.ipv6.conf.all.accept_redirects = 0
net.ipv4.conf.default.accept_redirects = 0
net.ipv6.conf.default.accept_redirects = 0
net.ipv4.conf.all.secure_redirects = 0
net.ipv4.conf.default.secure_redirects = 0
net.ipv4.conf.all.accept_source_route = 0
net.ipv6.conf.all.accept_source_route = 0
net.ipv4.conf.default.accept_source_route = 0
net.ipv6.conf.default.accept_source_route = 0
net.ipv4.conf.all.log_martians = 1
net.ipv4.conf.default.log_martians = 1
net.ipv4.icmp_echo_ignore_broadcasts = 1
net.ipv4.conf.all.rp_filter = 1
net.ipv4.conf.default.rp_filter = 1
net.ipv4.tcp_max_syn_backlog = 2048
net.ipv4.tcp_synack_retries = 2
net.ipv4.tcp_syn_retries = 5
kernel.dmesg_restrict = 1
kernel.kptr_restrict = 2
EOF
    
    safe_execute "sysctl -p '$sysctl_conf'" "应用内核参数" false
    
    log_success "✓ 内核参数配置完成"
}

# 安装安全工具
install_security_tools() {
    log_section "安装安全工具"
    
    local tools=(
        "rkhunter:Rootkit检测"
        "lynis:安全审计"
        "aide:文件完整性检查"
        "logwatch:日志分析"
    )
    
    for tool_desc in "${tools[@]}"; do
        local tool="${tool_desc%:*}"
        local desc="${tool_desc#*:}"
        
        if confirm_action "是否安装 $tool ($desc)？"; then
            if pacman -Q "$tool" &> /dev/null; then
                log_info "$tool 已安装"
                continue
            fi
            
            safe_execute "pacman -S --noconfirm '$tool'" "安装 $tool" false
            
            # 特殊配置
            case "$tool" in
                "rkhunter")
                    safe_execute "rkhunter --update" "更新rkhunter" false
                    safe_execute "rkhunter --propupd" "更新属性" false
                    ;;
            esac
        fi
    done
}

# 配置自动更新
configure_auto_updates() {
    log_section "配置自动更新"
    
    echo -e "${CYAN}Arch Linux 自动更新选项:${NC}"
    echo "1) 使用systemd timer (推荐)"
    echo "2) 跳过"
    
    if [[ -t 0 ]]; then
        echo -e "${YELLOW}选择 [1-2]: ${NC}\c"
        read -p "" update_choice
    else
        update_choice="2"
    fi
    
    if [[ "$update_choice" == "1" ]]; then
        # 创建更新脚本
        cat > /usr/local/bin/arch-auto-update.sh << 'EOF'
#!/bin/bash
/usr/bin/pacman -Syu --noconfirm >> /var/log/arch-auto-update.log 2>&1
EOF
        chmod +x /usr/local/bin/arch-auto-update.sh
        
        # 创建systemd service
        cat > /etc/systemd/system/arch-auto-update.service << 'EOF'
[Unit]
Description=Arch Linux Auto Update
After=network-online.target

[Service]
Type=oneshot
ExecStart=/usr/local/bin/arch-auto-update.sh
EOF
        
        # 创建systemd timer
        cat > /etc/systemd/system/arch-auto-update.timer << 'EOF'
[Unit]
Description=Arch Linux Auto Update Timer

[Timer]
OnCalendar=daily
Persistent=true

[Install]
WantedBy=timers.target
EOF
        
        safe_execute "systemctl enable --now arch-auto-update.timer" "启用自动更新" false
        
        log_success "✓ 自动更新配置完成（每天执行）"
    else
        log_warning "跳过自动更新配置"
    fi
}

# 生成报告
generate_report() {
    log_section "生成安全报告"
    
    local report_file="/root/arch_hardening_report_$(date +%Y%m%d_%H%M%S).txt"
    
    cat > "$report_file" << EOF
================================================================================
Arch Linux 服务器安全加固报告 v2.0
================================================================================
生成时间: $(date)
主机名: $(hostname)
内核版本: $(uname -r)
包管理器: $(pacman --version | head -1)

--------------------------------------------------------------------------------
1. SSH配置
--------------------------------------------------------------------------------
$(grep -E "^(Port|PermitRootLogin|PasswordAuthentication)" /etc/ssh/sshd_config 2>/dev/null)

--------------------------------------------------------------------------------
2. 防火墙: $FIREWALL_CHOICE
--------------------------------------------------------------------------------
$(case $FIREWALL_CHOICE in
    "nftables") nft list ruleset 2>/dev/null ;;
    "firewalld") firewall-cmd --list-all 2>/dev/null ;;
    "ufw") ufw status verbose 2>/dev/null ;;
    *) echo "未配置" ;;
esac)

--------------------------------------------------------------------------------
3. 入侵防护: $IPS_CHOICE
--------------------------------------------------------------------------------
$(case $IPS_CHOICE in
    "sshguard") systemctl status sshguard --no-pager 2>/dev/null || echo "未运行" ;;
    "crowdsec") cscli metrics 2>/dev/null || echo "未运行" ;;
    "fail2ban") fail2ban-client status 2>/dev/null || echo "未运行" ;;
    *) echo "未配置" ;;
esac)

--------------------------------------------------------------------------------
4. SSH密钥配置
--------------------------------------------------------------------------------
$(for home in /home/* /root; do
    if [[ -f "$home/.ssh/authorized_keys" ]]; then
        echo "$(basename $home): $(grep -c '^ssh-' "$home/.ssh/authorized_keys" 2>/dev/null || echo 0) 个密钥"
    fi
done)

--------------------------------------------------------------------------------
备份目录: $BACKUP_DIR
日志文件: $LOG_FILE
回滚脚本: $ROLLBACK_SCRIPT
--------------------------------------------------------------------------------
EOF
    
    log_success "✓ 报告已生成: $report_file"
    
    echo -e "${CYAN}━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━${NC}"
    cat "$report_file"
    echo -e "${CYAN}━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━${NC}"
}

################################################################################
# 主程序
################################################################################

main() {
    clear
    
    echo -e "${BLUE}"
    cat << "EOF"
╔══════════════════════════════════════════════════════════════╗
║          Arch Linux 服务器安全加固脚本 v2.0                ║
║          针对Arch Linux优化的现代化工具选择                ║
╚══════════════════════════════════════════════════════════════╝

防火墙选项:
  • nftables (现代化，推荐)
  • firewalld (易用)
  • ufw (简单)

入侵防护:
  • sshguard (轻量，原生支持nftables/firewalld)
  • crowdsec (社区驱动，现代化)
  • fail2ban (传统，功能全面)

EOF
    echo -e "${NC}"
    
    check_root
    check_arch_linux
    
    # 初始化
    mkdir -p "$BACKUP_DIR"
    echo "#!/bin/bash" > "$ROLLBACK_SCRIPT"
    echo "# Arch Linux 回滚脚本 - $(date)" >> "$ROLLBACK_SCRIPT"
    chmod +x "$ROLLBACK_SCRIPT"
    
    log "开始Arch Linux安全加固..."
    log_info "日志: $LOG_FILE"
    log_info "备份: $BACKUP_DIR"
    
    if ! confirm_action "确认继续？"; then
        exit 0
    fi
    
    # 执行加固
    collect_system_info
    install_aur_helper
    update_system
    manage_users
    harden_ssh
    configure_firewall
    configure_ips
    harden_users
    harden_kernel
    install_security_tools
    configure_auto_updates
    generate_report
    
    log_section "加固完成"
    
    echo -e "${GREEN}"
    cat << EOF
╔══════════════════════════════════════════════════════════════╗
║                 Arch Linux 加固完成！                       ║
╚══════════════════════════════════════════════════════════════╝

配置信息:
  防火墙: $FIREWALL_CHOICE
  入侵防护: $IPS_CHOICE
  备份: $BACKUP_DIR
  日志: $LOG_FILE

下一步:
1. 测试SSH登录（不要断开当前连接）
2. 检查防火墙: nft list ruleset / firewall-cmd --list-all / ufw status
3. 检查服务: systemctl status sshd $FIREWALL_CHOICE $IPS_CHOICE
4. 查看日志: journalctl -xe

EOF
    echo -e "${NC}"
    
    if confirm_action "是否现在重启？"; then
        log "5秒后重启..."
        sleep 5
        reboot
    fi
}

main "$@")端口？"; then
        sed -i 's/# tcp dport { 80, 443 } accept/tcp dport 80 accept/' "$nft_conf"
    fi
    
    if confirm_action "是否开放HTTPS(443)端口？"; then
        sed -i '/tcp dport 80 accept/a\        tcp dport 443 accept' "$nft_conf"
    fi
    
    # 测试配置
    if nft -c -f "$nft_conf" 2>&1 | tee -a "$LOG_FILE"; then
        log_success "✓ nftables配置验证通过"
        
        # 应用规则
        safe_execute "nft -f '$nft_conf'" "应用nftables规则" true
        
        # 启用服务
        safe_execute "systemctl enable nftables.service" "启用nftables服务" false
        safe_execute "systemctl start nftables.service" "启动nftables服务" false
        
        # 显示当前规则
        echo -e "${CYAN}━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━${NC}"
        nft list ruleset
        echo -e "${CYAN}━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━${NC}"
        
        log_success "✓ nftables配置完成"
    else
        log_error "✗ nftables配置验证失败"
    fi
}

# 配置firewalld
configure_firewalld() {
    log_section "配置firewalld防火墙"
    
    if ! command -v firewall-cmd &> /dev/null; then
        safe_execute "pacman -S --noconfirm firewalld" "安装firewalld" true
    fi
    
    # 启动服务
    safe_execute "systemctl enable --now firewalld" "启用firewalld" true
    
    # 获取SSH端口
    local ssh_port=$(grep "^Port" /etc/ssh/sshd_config 2>/dev/null | awk '{print $2}')
    ssh_port=${ssh_port:-22}
    
    # 配置规则
    safe_execute "firewall-cmd --permanent --add-port=$ssh_port/tcp" "开放SSH端口" true
    
    if confirm_action "是否开放HTTP(80)？"; then
        safe_execute "firewall-cmd --permanent --add-service=http" "开放HTTP" false
    fi
    
    if confirm_action "是否开放HTTPS(443)？"; then
        safe_execute "firewall-cmd --permanent --add-service=https" "开放HTTPS" false
    fi
    
    # 重载配置
    safe_execute "firewall-cmd --reload" "重载firewall配置" false
    
    # 显示状态
    firewall-cmd --list-all
    
    log_success "✓ firewalld配置完成"
}

# 配置UFW
configure_ufw() {
    log_section "配置UFW防火墙"
    
    # 安装UFW和iptables-nft
    if ! command -v ufw &> /dev/null; then
        safe_execute "pacman -S --noconfirm ufw iptables-nft" "安装UFW" true
    fi
    
    local ssh_port=$(grep "^Port" /etc/ssh/sshd_config 2>/dev/null | awk '{print $2}')
    ssh_port=${ssh_port:-22}
    
    safe_execute "ufw default deny incoming" "设置默认策略" false
    safe_execute "ufw default allow outgoing" "设置默认策略" false
    safe_execute "ufw allow $ssh_port/tcp comment 'SSH'" "开放SSH" true
    
    if confirm_action "是否开放HTTP(80