#!/bin/bash

# ==============================================================================
# 脚本名称: fail2ban_setup.sh
# 脚本功能: 自动化安装与配置 Fail2ban 动态防暴力破解（加强版）
# 支持系统: Debian, Ubuntu, CentOS, RHEL, Rocky Linux, AlmaLinux
# ==============================================================================

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

# 检查是否为 root 用户
if [ "$(id -u)" != "0" ]; then
    echo -e "${RED}错误：请使用 root 用户或通过 sudo 运行此脚本！${PLAIN}"
    exit 1
fi

# 检测系统架构与发行版
if [ -f /etc/os-release ]; then
    . /etc/os-release
    OS=$ID
    VERSION_ID=$VERSION_ID
else
    echo -e "${RED}错误：无法识别当前系统系统！${PLAIN}"
    exit 1
fi

echo -e "${BLUE}==================================================${PLAIN}"
echo -e "${GREEN}          Fail2ban 动态防暴力破解一键配置脚本${PLAIN}"
echo -e "${BLUE}==================================================${PLAIN}"

# 1. 安装 Fail2ban
install_fail2ban() {
    echo -e "${YELLOW}[1/3] 正在安装 Fail2ban...${PLAIN}"
    if [[ "$OS" == "ubuntu" || "$OS" == "debian" ]]; then
        apt-get update -y
        apt-get install -y fail2ban rsyslog
    elif [[ "$OS" == "centos" || "$OS" == "rhel" || "$OS" == "rocky" || "$OS" == "alma" ]]; then
        yum install -y epel-release
        yum makecache
        yum install -y fail2ban rsyslog
    else
        echo -e "${RED}暂不支持此系统发行版：$OS${PLAIN}"
        exit 1
    fi
    
    if [ $? -eq 0 ]; then
        echo -e "${GREEN}Fail2ban 安装成功！${PLAIN}"
    else
        echo -e "${RED}Fail2ban 安装失败，请检查网络或软件源。${PLAIN}"
        exit 1
    fi
}

# 2. 配置 Fail2ban
configure_fail2ban() {
    echo -e "${YELLOW}[2/3] 正在配置防护规则...${PLAIN}"
    
    # 确保防火墙后端适配（有些新系统默认使用 nftables 或 firewalld，这里做基础兼容）
    BANACTION="iptables-multiport"
    if command -v ufw >/dev/null 2>&1; then
        BANACTION="ufw"
    elif command -v firewall-cmd >/dev/null 2>&1; then
        BANACTION="firewalld-multiport"
    fi

    # 创建自定义局部配置文件 jail.local
    cat > /etc/fail2ban/jail.local <<EOF
[DEFAULT]
# 忽略的白名单 IP，可以是单个 IP 或 CIDR 网段（空格隔开）
ignoreip = 127.0.0.1/8 ::1

# 封禁判定窗口时间（5分钟内）
findtime = 300

# 窗口时间内最大失败尝试次数
maxretry = 5

# 触发后的基础封禁时间（24小时）
bantime = 86400

# 封禁动作后端
banaction = $BANACTION

# 开启【累犯永久封禁】机制（Fail2ban 自带的恶意增量惩罚）
bantime.increment = true
bantime.factor = 2
bantime.formula = ban.time * (1 << ban.count)
bantime.maxtime = 315360000 # 最大封禁10年（相当于永久）

[sshd]
enabled = true
port    = ssh
logpath = %(sshd_log)s
backend = %(sshd_backend)s
maxretry = 5
EOF

    echo -e "${GREEN}防护规则配置完成！(默认 5 分钟内错误 5 次，封禁 24 小时；累犯自动翻倍至永久封禁)${PLAIN}"
}

# 3. 启动并验证服务
start_services() {
    echo -e "${YELLOW}[3/3] 正在启动服务...${PLAIN}"
    
    # 确保系统日志服务在运行（Fail2ban 需要读取 auth.log / secure）
    systemctl enable rsyslog --now >/dev/null 2>&1
    
    # 启动 Fail2ban
    systemctl daemon-reload
    systemctl enable fail2ban
    systemctl restart fail2ban
    
    if systemctl is-active --quiet fail2ban; then
        echo -e "${GREEN}===============================================================================${PLAIN}"
        echo -e "${GREEN}🎉 Fail2ban 动态防暴破部署成功并已开始运行！查看状态：fail2ban-client status ${PLAIN}"
        echo -e "${GREEN}===============================================================================${PLAIN}"
    else
        echo -e "${RED}警告：Fail2ban 服务启动失败，请运行 'journalctl -u fail2ban' 查看错误日志。${PLAIN}"
    fi
}

# 简易管理菜单
show_status() {
    echo -e "\n${BLUE}--- Fail2ban 当前防护状态 ---${PLAIN}"
    if systemctl is-active --quiet fail2ban; then
        # 获取被封禁的 SSH IP
        INFO=$(fail2ban-client status sshd)
        echo "$INFO"
    else
        echo -e "${RED}Fail2ban 未在运行！${PLAIN}"
    fi
}

# 执行流程
install_fail2ban
configure_fail2ban
start_services