#!/bin/bash

# ==============================================================================
# 脚本名称: sftp_setup.sh
# 功能描述: 自动化配置一个安全的、chrooted的SFTP用户，用于文件备份。
#           用户将被限制在特定目录，无法SSH登录，也无法访问系统其他部分。
# 调用方式: sudo ./sftp_setup.sh
# 作者:     AI Assistant
# 版本:     1.4 (修复Subsystem重复定义问题)
# ==============================================================================

# --- 颜色定义，用于美化输出 ---
COLOR_GREEN='\033[0;32m'
COLOR_RED='\033[0;31m'
COLOR_YELLOW='\033[1;33m'
NC='\033[0m' # No Color

# --- 辅助函数 ---
log_info() {
    echo -e "${COLOR_GREEN}[INFO]${NC} $1"
}

log_warn() {
    echo -e "${COLOR_YELLOW}[WARN]${NC} $1"
}

log_error() {
    echo -e "${COLOR_RED}[ERROR]${NC} $1"
    exit 1
}

# 检查是否以root权限运行
check_root() {
    if [[ $EUID -ne 0 ]]; then
        log_error "此脚本必须以 root 权限运行。请使用 'sudo $0'"
    fi
}

# 清理旧的SFTP配置
cleanup_old_sftp_config() {
    local config_file="$1"
    
    log_info "正在清理旧的 SFTP 配置..."
    
    # 创建临时文件
    local temp_file=$(mktemp)
    
    # 使用更精确的模式来清理配置
    # 删除所有Subsystem sftp相关的行（包括注释和空白行后的注释）
    awk '
    /^[[:space:]]*#?[[:space:]]*Subsystem[[:space:]]+sftp/ {
        # 跳过这一行（删除）
        next
    }
    /^[[:space:]]*#.*sftp.*subsystem/ {
        # 删除包含sftp subsystem的注释行
        next
    }
    /^Match[[:space:]]+User/ {
        # 如果遇到新的Match User块，停止删除
        in_old_block = 0
    }
    in_old_block && /^[[:space:]]+#?[[:space:]]*(ChrootDirectory|ForceCommand|AllowTcpForwarding|X11Forwarding|PermitTunnel)/ {
        # 在旧块中删除相关配置行
        next
    }
    {
        # 打印其他所有行
        print
    }
    ' "$config_file" > "$temp_file"
    
    # 检查是否有变化
    if ! diff "$config_file" "$temp_file" &>/dev/null; then
        mv "$temp_file" "$config_file"
        log_info "已清理旧的 SFTP 配置。"
    else
        rm -f "$temp_file"
        log_info "未发现需要清理的旧配置。"
    fi
}

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

log_info "开始 SFTP 用户自动配置流程..."

# 1. 检查root权限
check_root

# 2. 交互式获取用户输入
echo "--- 请输入 SFTP 用户配置信息 ---"
read -p "1. 请输入新的 SFTP 用户名: " SFTP_USERNAME
read -s -p "2. 请输入该用户的密码: " SFTP_PASSWORD
echo "" # 换行
read -s -p "3. 请再次输入密码以确认: " SFTP_PASSWORD_CONFIRM
echo "" # 换行
read -p "4. 请输入备份文件的根目录 (例如 /sftp/backup): " SFTP_ROOT_DIR

# 3. 验证输入
if [[ -z "$SFTP_USERNAME" || -z "$SFTP_PASSWORD" || -z "$SFTP_ROOT_DIR" ]]; then
    log_error "用户名、密码和根目录不能为空！"
fi

if [[ "$SFTP_PASSWORD" != "$SFTP_PASSWORD_CONFIRM" ]]; then
    log_error "两次输入的密码不匹配！"
fi

if [[ ! "$SFTP_ROOT_DIR" = /* ]]; then
    log_error "根目录必须以 '/' 开头，为绝对路径！"
fi

# 定义SFTP组和上传目录名
SFTP_GROUP="sftpgroup"
UPLOAD_DIR_NAME="uploads"
SFTP_UPLOAD_DIR="${SFTP_ROOT_DIR}/${UPLOAD_DIR_NAME}"

log_info "配置信息确认无误，开始执行配置..."

# 4. 创建SFTP组
if ! grep -q "^${SFTP_GROUP}:" /etc/group; then
    log_info "创建 SFTP 专用组: ${SFTP_GROUP}"
    groupadd "$SFTP_GROUP"
else
    log_warn "组 '${SFTP_GROUP}' 已存在，跳过创建。"
fi

# 5. 创建SFTP用户
if id "$SFTP_USERNAME" &>/dev/null; then
    log_warn "用户 '${SFTP_USERNAME}' 已存在，将跳过用户创建，但会重置密码和目录权限。"
else
    log_info "创建 SFTP 用户: ${SFTP_USERNAME}"
    useradd -G "$SFTP_GROUP" -s /sbin/nologin -d "$SFTP_ROOT_DIR" "$SFTP_USERNAME"
fi

# 6. 设置用户密码
log_info "为用户 '${SFTP_USERNAME}' 设置密码..."
echo "${SFTP_USERNAME}:${SFTP_PASSWORD}" | chpasswd

# 7. 创建目录结构
log_info "创建 SFTP 目录结构..."
mkdir -p "$SFTP_ROOT_DIR"
mkdir -p "$SFTP_UPLOAD_DIR"

# 8. 设置目录权限 (最关键的一步)
log_info "设置目录权限..."
chown root:root "$SFTP_ROOT_DIR"
chmod 755 "$SFTP_ROOT_DIR"
chown "${SFTP_USERNAME}:${SFTP_GROUP}" "$SFTP_UPLOAD_DIR"
chmod 755 "$SFTP_UPLOAD_DIR"
log_info "权限设置完成。"

# 9. 配置 SSHD (修复版)
SSHD_CONFIG="/etc/ssh/sshd_config"
SSHD_BACKUP="${SSHD_CONFIG}.bak.$(date +%F_%H-%M-%S)"

log_info "配置 SSH 守护进程..."

# 备份原始文件
cp "$SSHD_CONFIG" "$SSHD_BACKUP"
log_info "已备份原始 sshd_config 到 ${SSHD_BACKUP}"

# --- 修复的配置逻辑 ---

# 1. 首先清理所有旧的SFTP相关配置
cleanup_old_sftp_config "$SSHD_CONFIG"

# 2. 检查是否已存在Subsystem sftp配置
if grep -q "^Subsystem[[:space:]]\+sftp" "$SSHD_CONFIG"; then
    log_warn "发现已存在的 Subsystem sftp 配置，将使用现有配置。"
else
    log_info "添加新的 'Subsystem sftp internal-sftp' 配置。"
    echo "Subsystem sftp internal-sftp" >> "$SSHD_CONFIG"
fi

# 3. 清理该用户可能存在的旧配置
log_info "清理用户 '${SFTP_USERNAME}' 的旧配置..."
# 删除该用户的所有Match块
sed -i -e "/^Match User ${SFTP_USERNAME}$/,/^Match User\|^$/d" "$SSHD_CONFIG"

# 4. 添加新的用户配置
log_info "为用户 '${SFTP_USERNAME}' 添加新的 Chroot 配置..."
{
    printf "\n# --- SFTP Chroot for user: %s (由脚本自动添加) ---\n" "${SFTP_USERNAME}"
    printf "Match User %s\n" "${SFTP_USERNAME}"
    printf "    ForceCommand internal-sftp\n"
    printf "    ChrootDirectory %s\n" "${SFTP_ROOT_DIR}"
    printf "    AllowTcpForwarding no\n"
    printf "    X11Forwarding no\n"
    printf "    PermitTunnel no\n"
} >> "$SSHD_CONFIG"

log_info "SSH 配置更新完成。"

# --- 配置逻辑结束 ---

# 10. 测试SSH配置并重启服务
log_info "正在测试 SSH 配置文件语法..."
if sshd -t; then
    log_info "配置文件语法正确。"
    log_info "正在重启 SSH 服务以应用配置..."
    if command -v systemctl &> /dev/null; then
        systemctl restart sshd
        # 检查服务状态
        if systemctl is-active --quiet sshd; then
            log_info "SSH 服务重启成功。"
        else
            log_error "SSH 服务启动失败，请检查系统日志。"
        fi
    else
        service ssh restart
        log_info "SSH 服务重启完成。"
    fi
else
    log_error "SSH 配置文件语法错误！服务未重启。"
    log_error "以下是所有相关的 SFTP 配置行，请检查："
    grep -n -A5 -B2 "Subsystem sftp\|Match User ${SFTP_USERNAME}" "$SSHD_CONFIG" || true
    log_error "请检查 /etc/ssh/sshd_config 文件，或使用备份文件 ${SSHD_BACKUP} 恢复。"
    exit 1
fi

# 11. 验证配置
log_info "验证 SFTP 配置..."
if grep -q "Match User ${SFTP_USERNAME}" "$SSHD_CONFIG" && \
   grep -q "ChrootDirectory ${SFTP_ROOT_DIR}" "$SSHD_CONFIG"; then
    log_info "SFTP 配置验证成功。"
else
    log_warn "SFTP 配置可能未正确应用，请手动检查。"
fi

# 12. 完成
echo "=================================================="
log_info "SFTP 用户配置完成！"
echo "=================================================="
echo -e "用户名: ${COLOR_YELLOW}${SFTP_USERNAME}${NC}"
echo -e "密码:   ${COLOR_YELLOW}[已设置]${NC}"
echo -e "SFTP根目录: ${COLOR_YELLOW}${SFTP_ROOT_DIR}${NC}"
echo -e "文件上传目录: ${COLOR_YELLOW}${SFTP_UPLOAD_DIR}${NC}"
echo ""
echo -e "${COLOR_GREEN}现在，您可以在您的备份脚本中使用以下信息：${NC}"
echo -e "  SERVER_USER=\"${SFTP_USERNAME}\""
echo -e "  REMOTE_DIR=\"/${UPLOAD_DIR_NAME}\""
echo ""
echo -e "${COLOR_YELLOW}测试命令 (在服务器本机执行):${NC}"
echo -e "  sftp ${SFTP_USERNAME}@localhost"
echo ""
echo -e "${COLOR_YELLOW}重要提醒:${NC}"
echo -e "  1. 确保根目录 ${SFTP_ROOT_DIR} 的所有者是 root:root"
echo -e "  2. 用户只能访问 ${SFTP_UPLOAD_DIR} 目录"
echo -e "  3. 用户无法通过SSH登录系统"
echo "=================================================="

exit 0