#!/bin/bash# ==============================================
# MongoDB 安全数据替换脚本 (执行顺序：备份→校验→确认→清空→还原指定数据→失败回滚到备份)
# 依赖命令: mongoexport, mongorestore, mongosh/mongo
# ==============================================# Configuration parameters (配置参数)
DEFAULT_HOST="localhost"
DEFAULT_PORT="27017"
DEFAULT_AUTH_SOURCE="admin"
BACKUP_DIR="./mongo_backups"
LOG_FILE="./mongo_restore_$(date +%Y%m%d_%H%M%S).log"# Color definitions (颜色定义)
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
NC='\033[0m' # No Color (无颜色)# Global variable to store the backup file path (全局变量，用于存储备份文件路径)
BACKUP_FILE=""# Log messages (日志消息函数)
log() {local level=$1local message=$2local timestamp=$(date '+%Y-%m-%d %H:%M:%S')echo "[${timestamp}] [${level}] ${message}" >> "$LOG_FILE"case "$level" in"ERROR") echo -e "${RED}[ERROR]${NC} ${message}" >&2 ;;"WARN") echo -e "${YELLOW}[WARN]${NC} ${message}" ;;"SUCCESS") echo -e "${GREEN}[SUCCESS]${NC} ${message}" ;;*) echo "[INFO] ${message}" ;;esac
}# Check for required commands (检查所需命令)
check_required_commands() {local missing=()# Check for mongosh or mongo (检查 mongosh 或 mongo)if ! command -v mongosh &> /dev/null && ! command -v mongo &> /dev/null; thenmissing+=("mongosh/mongo")fi# Check for other required commands (检查其他所需命令)for cmd in mongoexport mongorestore mongoimport; doif ! command -v "$cmd" &> /dev/null; thenmissing+=("$cmd")fidoneif [ ${#missing[@]} -gt 0 ]; thenlog "ERROR" "缺少必要的命令: ${missing[*]}"log "ERROR" "请确保以下命令已安装并配置在您的系统 PATH 中:"log "ERROR" "1. mongoexport"log "ERROR" "2. mongorestore"log "ERROR" "3. mongoimport"log "ERROR" "4. mongosh 或 mongo"exit 1fi
}# Find MongoDB client command (mongosh or mongo) (查找 MongoDB 客户端命令)
find_mongo_client() {if command -v mongosh &> /dev/null; thenMONGO_CLIENT_CMD="mongosh"elif command -v mongo &> /dev/null; thenMONGO_CLIENT_CMD="mongo"elselog "ERROR" "严重错误: 未找到 MongoDB Shell 客户端"exit 1filog "INFO" "正在使用 MongoDB 客户端: ${MONGO_CLIENT_CMD}"
}# Validate file existence and readability (验证文件存在性和可读性)
validate_file() {local file=$1local description=$2if [ ! -e "$file" ]; thenlog "ERROR" "${description} 文件不存在: $file"return 1 # Indicate failure (指示失败)fiif [ ! -f "$file" ]; thenlog "ERROR" "${description} 路径不是一个文件: $file"return 1 # Indicate failure (指示失败)fiif [ ! -r "$file" ]; thenlog "ERROR" "无法读取 ${description} 文件: $file"return 1 # Indicate failure (指示失败)fireturn 0 # Indicate success (指示成功)
}# Display help information (显示帮助信息)
show_help() {echo -e "${YELLOW}MongoDB 安全数据替换脚本 - 执行顺序: 备份 → 校验 → 确认 → 清空 → 还原指定数据 → 失败回滚到备份${NC}"echo ""echo "用法:"echo " $0 -d 数据库名 -c 集合名 [-f 导入文件] [选项]"echo ""echo "必填参数:"echo " -d, --database 数据库名称"echo " -c, --collection 集合名称"echo ""echo "可选参数 (如果未提供 -f，则必须使用 --clear-only 或 --no-clear):"echo " -f, --file 要导入的数据文件路径 (必须是 JSON/BSON 格式)"echo ""echo "选项:"echo " -h, --host MongoDB 主机 (默认: ${DEFAULT_HOST})"echo " -p, --port MongoDB 端口 (默认: ${DEFAULT_PORT})"echo " -u, --username 认证用户名"echo " -w, --password 认证密码"echo " -a, --auth-source 认证数据库 (默认: ${DEFAULT_AUTH_SOURCE})"echo " -y, --assume-yes 跳过确认并直接执行 (危险操作!)"echo " --clear-only 仅清空集合并备份，不导入新数据 (此模式下无需 -f)"echo " --no-clear 跳过清空集合步骤 (导入时保留现有数据)"echo " --help 显示此帮助信息"echo ""echo -e "${RED}注意: 此脚本已修改。所有操作现在都将强制进行备份。${NC}"echo -e "${RED}重要: 清空和导入操作只有在备份成功 并且 (如果提供了导入文件) 导入文件有效时才会继续。${NC}"echo -e "${RED}如果导入失败，将尝试自动回滚到备份数据。${NC}"echo "脚本需要 'mongoexport', 'mongorestore', 'mongoimport' 和 'mongosh' (或 'mongo') 命令在您的系统 PATH 中。"exit 0
}# Initialize environment (初始化环境)
init_env() {# First, check for required commands (首先，检查所需命令)check_required_commands# Then, find the client (然后，查找客户端)find_mongo_client# Create backup directory (创建备份目录)mkdir -p "$BACKUP_DIR" || {log "ERROR" "创建备份目录失败: $BACKUP_DIR"exit 1}# Initialize log file (初始化日志文件)echo "=== MongoDB 安全替换脚本日志 ===" > "$LOG_FILE"echo "执行时间: $(date '+%Y-%m-%d %H:%M:%S')" >> "$LOG_FILE"echo "依赖检查:" >> "$LOG_FILE"echo " mongosh/mongo: $(command -v ${MONGO_CLIENT_CMD})" >> "$LOG_FILE"echo " mongoexport: $(command -v mongoexport)" >> "$LOG_FILE"echo " mongorestore: $(command -v mongorestore)" >> "$LOG_FILE"echo " mongoimport: $(command -v mongoimport)" >> "$LOG_FILE"echo "==============================" >> "$LOG_FILE"
}# Parse command-line arguments (解析命令行参数)
parse_arguments() {while [[ $# -gt 0 ]]; docase "$1" in-d|--database)DB_NAME="$2"shift 2;;-c|--collection)COLLECTION_NAME="$2"shift 2;;-f|--file)IMPORT_FILE="$2"shift 2;;-h|--host)HOST="$2"shift 2;;-p|--port)PORT="$2"if ! [[ "$PORT" =~ ^[0-9]+$ ]]; thenlog "ERROR" "端口号必须是数字"exit 1fishift 2;;-u|--username)USERNAME="$2"shift 2;;-w|--password)PASSWORD="$2"shift 2;;-a|--auth-source)AUTH_SOURCE="$2"shift 2;;-y|--assume-yes)ASSUME_YES=trueshift;;--clear-only)CLEAR_ONLY=trueshift;;--no-clear)NO_CLEAR=trueshift;;--help)show_help;;*)log "ERROR" "未知参数: $1"show_helpexit 1;;esacdone# Validate required parameters (验证必填参数)if [ -z "$DB_NAME" ] || [ -z "$COLLECTION_NAME" ]; thenlog "ERROR" "缺少必填参数: 数据库名或集合名!"show_helpexit 1fi# Validate --clear-only and --no-clear are mutually exclusive (验证 --clear-only 和 --no-clear 互斥)if [ "$CLEAR_ONLY" = true ] && [ "$NO_CLEAR" = true ]; thenlog "ERROR" "不能同时使用 --clear-only 和 --no-clear 选项。"show_helpexit 1fi# Validate import file (-f) usage logic (验证导入文件 (-f) 的使用逻辑)if [ -n "$IMPORT_FILE" ]; thenif [ "$CLEAR_ONLY" = true ]; thenlog "ERROR" "不能同时指定导入文件 (-f) 和 --clear-only 选项。"show_helpexit 1fi# Validate file extension (验证文件扩展名)local ext="${IMPORT_FILE##*.}"if [[ ! "$ext" =~ ^(json|bson)$ ]]; thenlog "ERROR" "不支持的导入文件格式: $ext (仅支持 .json 或 .bson)"exit 1fielse # No import file specified (未指定导入文件)if [ "$CLEAR_ONLY" != true ] && [ "$NO_CLEAR" != true ]; thenlog "ERROR" "请提供要导入的文件 (-f)，或使用 --clear-only / --no-clear 选项。"log "ERROR" "注意: 所有操作都将强制进行备份。"show_helpexit 1fifi# If --no-clear is specified, an import file must also be specified (如果指定了 --no-clear，则必须同时指定导入文件)if [ "$NO_CLEAR" = true ] && [ -z "$IMPORT_FILE" ]; thenlog "ERROR" "使用 --no-clear 选项时，必须指定导入文件 (-f)。否则，操作无意义。"show_helpexit 1fi
}# Build MongoDB connection string and authentication options (构建 MongoDB 连接字符串和认证选项)
build_connection() {HOST="${HOST:-$DEFAULT_HOST}"PORT="${PORT:-$DEFAULT_PORT}"AUTH_SOURCE="${AUTH_SOURCE:-$DEFAULT_AUTH_SOURCE}"# Build URI connection string for mongosh/mongo shell (为 mongosh/mongo shell 构建 URI 连接字符串)MONGO_URI=""if [ -n "$USERNAME" ] && [ -n "$PASSWORD" ]; then# URL-encode special characters in password (对密码中的特殊字符进行 URL 编码)# Check if jq is available for URL encoding, otherwise fall back (检查 jq 是否可用，否则回退)if command -v jq &>/dev/null; thenENCODED_PASSWORD=$(printf '%s' "$PASSWORD" | jq -sRr @uri)elselog "WARN" "未找到 jq 命令。密码可能未进行 URL 编码。请确保您的密码不包含特殊字符，以便直接在 URI 中使用。"ENCODED_PASSWORD="$PASSWORD"fiMONGO_URI="mongodb://${USERNAME}:${ENCODED_PASSWORD}@${HOST}:${PORT}/${DB_NAME}?authSource=${AUTH_SOURCE}"elseMONGO_URI="mongodb://${HOST}:${PORT}/${DB_NAME}"fi# Build traditional command-line arguments for mongoexport/mongorestore (为 mongoexport/mongorestore 构建传统命令行参数)AUTH_STR=""if [ -n "$USERNAME" ] && [ -n "$PASSWORD" ]; thenAUTH_STR="--username ${USERNAME} --password ${PASSWORD} --authenticationDatabase ${AUTH_SOURCE}"fiFULL_OPTIONS="--host ${HOST} --port ${PORT} ${AUTH_STR}"
}# Backup existing data (备份现有数据)
# Returns 0 on successful (non-empty) backup, 1 otherwise (成功备份（非空）返回 0，否则返回 1)
backup_data() {local timestamp=$(date +"%Y%m%d_%H%M%S")BACKUP_FILE="${BACKUP_DIR}/${DB_NAME}_${COLLECTION_NAME}_${timestamp}.json"log "INFO" "正在开始备份现有数据 -> ${BACKUP_FILE}"if ! mongoexport ${FULL_OPTIONS} \--db "${DB_NAME}" \--collection="${COLLECTION_NAME}" \--out="${BACKUP_FILE}" \--jsonArray >> "$LOG_FILE" 2>&1; thenlog "ERROR" "备份失败! 请检查 MongoDB 连接或权限。后续操作将被跳过。"return 1fi# Validate backup file (验证备份文件)if [ ! -s "$BACKUP_FILE" ]; thenlog "WARN" "备份文件为空 (可能是空集合)。这通常没问题，但请注意。"return 0 # Still consider it a successful backup of an empty state (仍然认为这是空状态的成功备份)elseBACKUP_SIZE=$(du -h "$BACKUP_FILE" | cut -f1)log "SUCCESS" "备份成功 (大小: ${BACKUP_SIZE})"return 0fi
}# Confirm operation (delete or clear) (确认操作（删除或清空）)
confirm_operation() {if [ "$ASSUME_YES" = true ]; thenlog "WARN" "已启用自动确认，跳过用户确认"return 0filocal BACKUP_STATUS="${GREEN}强制备份已完成${NC}" # Always backed up now (现在总是已备份)local CLEAR_STATUS="永久删除"if [ "$NO_CLEAR" = true ]; thenCLEAR_STATUS="${YELLOW}不删除${NC} (保留现有数据)"elif [ "$CLEAR_ONLY" = true ]; thenCLEAR_STATUS="${YELLOW}仅清空${NC}" # Means deletion (意味着删除)filocal IMPORT_STATUS=""if [ "$CLEAR_ONLY" = true ]; thenIMPORT_STATUS="${YELLOW}不导入${NC}"elif [ -z "$IMPORT_FILE" ]; then # If no -f parameter, but not --clear-only (如果未提供 -f 参数，但不是 --clear-only)IMPORT_STATUS="${YELLOW}不导入${NC}"elseIMPORT_STATUS="从 ${GREEN}${IMPORT_FILE}${NC} 导入"fiecho -e "\n${YELLOW}======= 重要警告 =======${NC}"echo -e "即将执行以下操作:"echo -e "1. ${BACKUP_STATUS} ${DB_NAME}.${COLLECTION_NAME} 的现有数据"echo -e "2. ${CLEAR_STATUS} ${DB_NAME}.${COLLECTION_NAME} 中的所有数据"echo -e "3. ${IMPORT_STATUS} 新数据"# Specific warnings (特定警告)if [ "$NO_CLEAR" = true ] && [ -n "$IMPORT_FILE" ]; thenecho -e "${YELLOW}注意: 您已选择不清空集合。新数据将导入并与现有数据合并。${NC}"fiecho -e "备份文件将保存到: ${BACKUP_DIR}"echo -e "${YELLOW}=======================${NC}\n"read -p "您确定要继续吗? (yes/no): " -rechoif [[ ! $REPLY =~ ^[Yy][Ee][Ss]$ ]]; thenlog "INFO" "用户取消了操作"exit 0fi
}# Clear target collection (清空目标集合)
clear_collection() {log "INFO" "正在清空集合..."if ! "${MONGO_CLIENT_CMD}" "${MONGO_URI}" --quiet --eval "const result = db.${COLLECTION_NAME}.deleteMany({});print('Documents deleted: ' + result.deletedCount);" >> "$LOG_FILE" 2>&1; thenlog "ERROR" "清空集合失败!"return 1 # Indicate failure (指示失败)filog "SUCCESS" "集合清空成功。"return 0 # Indicate success (指示成功)
}# Import new data (导入新数据)
import_data() {local import_file_path=$1local source_desc=$2 # e.g., "new data" or "backup data" (例如，“新数据”或“备份数据”)log "INFO" "正在从 ${source_desc} (${import_file_path}) 开始导入数据..."case "${import_file_path##*.}" injson)if ! mongoimport ${FULL_OPTIONS} \--db "${DB_NAME}" \--collection="${COLLECTION_NAME}" \--file="${import_file_path}" \--jsonArray >> "$LOG_FILE" 2>&1; thenlog "ERROR" "${source_desc} JSON 导入失败!"return 1fi;;bson)# IMPORTANT: mongorestore --dir expects a directory, not a single file.# If import_file_path is a single .bson file, mongoimport --type bson is usually preferred.# This section assumes import_file_path is a directory containing BSON dump.if [ ! -d "${import_file_path}" ]; thenlog "ERROR" "对于 BSON 导入，--file 选项应指定包含 BSON 文件的目录，而不是单个文件。"log "ERROR" "如果 '${import_file_path}' 是单个 BSON 文件，请考虑使用 'mongoimport --type bson' 命令，或调整输入文件类型。"return 1fiif ! mongorestore ${FULL_OPTIONS} \--db "${DB_NAME}" \--collection="${COLLECTION_NAME}" \--dir="${import_file_path}" \--drop >> "$LOG_FILE" 2>&1; thenlog "ERROR" "${source_desc} BSON 导入失败!"return 1fi;;*)log "ERROR" "不支持 ${source_desc} 的格式: ${import_file_path##*.}"return 1;;esaclog "SUCCESS" "数据从 ${source_desc} 导入成功。"return 0
}# Restore data from backup (从备份还原数据)
restore_from_backup() {log "ERROR" "正在尝试从备份还原: ${BACKUP_FILE}"# Ensure collection is clear before restoring from backup to prevent duplicates if --no-clear was set# (在从备份还原之前确保集合已清空，以防止在设置 --no-clear 时出现重复数据)if ! clear_collection; thenlog "ERROR" "在从备份还原之前清空集合失败。需要手动干预!"return 1fiif import_data "$BACKUP_FILE" "备份数据"; thenlog "SUCCESS" "成功从备份还原数据。"return 0elselog "ERROR" "从备份还原数据失败! 绝对需要手动干预。请检查日志以获取详细信息。"return 1fi
}# Verify operation results (验证操作结果)
verify_operation() {log "INFO" "正在验证操作结果...""${MONGO_CLIENT_CMD}" "${MONGO_URI}" --quiet --eval "print('当前文档数量: ' + db.${COLLECTION_NAME}.countDocuments());print('示例文档:');printjson(db.${COLLECTION_NAME}.find().limit(1).toArray()[0] || '空集合');" | tee -a "$LOG_FILE"
}# Main flow (主流程)
main() {# Initialize environment (初始化环境)init_env# Parse command-line arguments (解析命令行参数)parse_arguments "$@"# Build connection parameters (构建连接参数)build_connection# Print operation summary (打印操作摘要)log "INFO" "操作摘要:"log "INFO" "数据库: ${DB_NAME}"log "INFO" "集合: ${COLLECTION_NAME}"log "INFO" "备份模式: 强制执行备份" # Updated message (更新消息)if [ "$NO_CLEAR" = true ]; thenlog "INFO" "清空模式: 跳过清空 (保留现有数据)"elif [ "$CLEAR_ONLY" = true ]; thenlog "INFO" "清空模式: 仅清空集合 (不导入新数据)"elselog "INFO" "清空模式: 先清空，然后导入"fiif [ -n "$IMPORT_FILE" ]; thenlog "INFO" "导入文件: ${IMPORT_FILE}"elselog "INFO" "导入文件: 未指定 (不执行导入)"fi[ -n "$USERNAME" ] && log "INFO" "认证用户: ${USERNAME}"# --- 步骤 1: 强制备份现有数据 ---# 如果备份失败，脚本将终止。if ! backup_data; thenlog "ERROR" "备份失败。中止后续操作。"log "INFO" "日志文件: ${LOG_FILE}"exit 1fi# --- 步骤 2: 校验还原文件是否存在，存在则执行后续操作 ---local proceed_with_modification=falseif [ -n "$IMPORT_FILE" ]; thenif validate_file "$IMPORT_FILE" "导入"; thenlog "INFO" "导入文件 '${IMPORT_FILE}' 校验成功。准备进行修改。"proceed_with_modification=trueelselog "ERROR" "导入文件校验失败。中止清空/导入操作，但备份已存在。"log "INFO" "日志文件: ${LOG_FILE}"exit 1 # 导入文件无效，退出fielif [ "$CLEAR_ONLY" = true ] || [ "$NO_CLEAR" = true ]; then# 如果没有导入文件，但设置了 --clear-only 或 --no-clear，我们仍将其视为“修改”意图。# --no-clear 且没有导入文件的情况已在 parse_arguments 中捕获。log "INFO" "操作为仅清空或与现有数据合并。准备进行修改。"proceed_with_modification=trueelse# 这种情况理论上应在 parse_arguments 中捕获，但作为回退。log "WARN" "未指定导入文件，且未处于仅清空/不清空模式。除备份外，不执行数据导入或清空。"fiif [ "$proceed_with_modification" = true ]; then# --- 步骤 3: 用户确认操作 ---confirm_operationlocal import_success=falseif [ "$NO_CLEAR" != true ]; then# --- 步骤 4: 清空目标集合 (如果未跳过) ---if ! clear_collection; thenlog "ERROR" "清空集合失败。正在尝试从备份还原。"if ! restore_from_backup; thenlog "CRITICAL" "清空失败后从备份还原也失败。绝对需要手动干预!"exit 1fi# 如果从备份还原成功，则认为此阶段的“修改”是成功的（即回滚成功）import_success=true fielselog "INFO" "由于 --no-clear 选项，跳过了清空集合步骤。"fi# --- 步骤 5: 导入新数据 (如果指定了文件且不是仅清空模式) ---# 只有在清空成功（或跳过清空）且没有发生回滚的情况下才尝试导入新数据if [ "$import_success" != true ] && [ -n "$IMPORT_FILE" ] && [ "$CLEAR_ONLY" != true ]; thenif ! import_data "$IMPORT_FILE" "新数据"; thenlog "ERROR" "新数据导入失败。正在尝试从备份还原。"# --- 步骤 6: 如果还原失败则用备份的数据还原 ---if ! restore_from_backup; thenlog "CRITICAL" "导入失败后从备份还原也失败。绝对需要手动干预!"exit 1fifielselog "INFO" "跳过了数据导入步骤。"fi# --- 步骤 7: 验证操作结果 ---verify_operation# 完成消息echo -e "\n${GREEN}=== 操作完成 ===${NC}"log "SUCCESS" "所有操作完成"elseecho -e "\n${YELLOW}=== 操作完成 (未执行数据修改) ===${NC}"log "INFO" "由于导入文件校验失败或未指定操作，未执行数据修改 (清空/导入)。"filog "INFO" "备份文件: ${BACKUP_FILE}" # 始终报告备份文件log "INFO" "日志文件: ${LOG_FILE}"
}main "$@"
