From 8a6f1b6f241cf9ee24d63a9ed79b1b1fff1b65b7 Mon Sep 17 00:00:00 2001 From: yyhuni Date: Sat, 10 Jan 2026 11:04:42 +0800 Subject: [PATCH] feat(engine): add --force-sub flag for selective engine config updates - Add --force-sub command flag to init_default_engine management command - Allow updating only sub-engines while preserving user-customized full scan config - Update docker/scripts/init-data.sh to always update full scan engine configuration - Change docker/server/start.sh to use --force flag for initial engine setup - Improve update.sh with better logging functions and formatted output - Add color-coded log functions (log_step, log_ok, log_info, log_warn, log_error) - Enhance update.sh UI with better visual formatting and warning messages - Refactor error messages and user prompts for improved clarity - This enables safer upgrades by preserving custom full scan configurations while updating sub-engines --- .../commands/init_default_engine.py | 30 +++-- docker/scripts/init-data.sh | 14 +-- docker/server/start.sh | 2 +- update.sh | 108 ++++++++++-------- 4 files changed, 93 insertions(+), 61 deletions(-) diff --git a/backend/apps/engine/management/commands/init_default_engine.py b/backend/apps/engine/management/commands/init_default_engine.py index 3898d7f5..35c48dd6 100644 --- a/backend/apps/engine/management/commands/init_default_engine.py +++ b/backend/apps/engine/management/commands/init_default_engine.py @@ -2,8 +2,9 @@ 初始化默认扫描引擎 用法： - python manage.py init_default_engine # 只创建不存在的引擎（不覆盖已有） - python manage.py init_default_engine --force # 强制覆盖所有引擎配置 + python manage.py init_default_engine # 只创建不存在的引擎（不覆盖已有） + python manage.py init_default_engine --force # 强制覆盖所有引擎配置 + python manage.py init_default_engine --force-sub # 只覆盖子引擎，保留 full scan cd /root/my-vulun-scan/docker docker compose exec server python backend/manage.py init_default_engine --force @@ -12,6 +13,7 @@ - 读取 engine_config_example.yaml 作为默认配置 - 创建 full scan（默认引擎）+ 各扫描类型的子引擎 - 默认不覆盖已有配置，加 --force 才会覆盖 +- 加 --force-sub 只覆盖子引擎配置，保留用户自定义的 full scan """ from django.core.management.base import BaseCommand @@ -30,11 +32,18 @@ class Command(BaseCommand): parser.add_argument( '--force', action='store_true', - help='强制覆盖已有的引擎配置', + help='强制覆盖已有的引擎配置（包括 full scan 和子引擎）', + ) + parser.add_argument( + '--force-sub', + action='store_true', + help='只覆盖子引擎配置，保留 full scan（升级时使用）', ) def handle(self, *args, **options): force = options.get('force', False) + force_sub = options.get('force_sub', False) + # 读取默认配置文件 config_path = Path(__file__).resolve().parent.parent.parent.parent / 'scan' / 'configs' / 'engine_config_example.yaml' @@ -99,15 +108,22 @@ class Command(BaseCommand): engine_name = f"{scan_type}" sub_engine = ScanEngine.objects.filter(name=engine_name).first() if sub_engine: - if force: + # force 或 force_sub 都会覆盖子引擎 + if force or force_sub: sub_engine.configuration = single_yaml sub_engine.save() - self.stdout.write(self.style.SUCCESS(f' ✓ 子引擎 {engine_name} 配置已更新 (ID: {sub_engine.id})')) + self.stdout.write(self.style.SUCCESS( + f' ✓ 子引擎 {engine_name} 配置已更新 (ID: {sub_engine.id})' + )) else: - self.stdout.write(self.style.WARNING(f' ⊘ {engine_name} 已存在，跳过（使用 --force 覆盖）')) + self.stdout.write(self.style.WARNING( + f' ⊘ {engine_name} 已存在，跳过（使用 --force 覆盖）' + )) else: sub_engine = ScanEngine.objects.create( name=engine_name, configuration=single_yaml, ) - self.stdout.write(self.style.SUCCESS(f' ✓ 子引擎 {engine_name} 已创建 (ID: {sub_engine.id})')) + self.stdout.write(self.style.SUCCESS( + f' ✓ 子引擎 {engine_name} 已创建 (ID: {sub_engine.id})' + )) diff --git a/docker/scripts/init-data.sh b/docker/scripts/init-data.sh index dec58e08..d9889b50 100755 --- a/docker/scripts/init-data.sh +++ b/docker/scripts/init-data.sh @@ -83,20 +83,20 @@ if not yaml_path.exists(): print('未找到配置文件，跳过') exit(0) +new_config = yaml_path.read_text() + # 检查是否已有 full scan 引擎 engine = ScanEngine.objects.filter(name='full scan').first() if engine: - if not engine.configuration or not engine.configuration.strip(): - engine.configuration = yaml_path.read_text() - engine.save(update_fields=['configuration']) - print(f'已初始化引擎配置: {engine.name}') - else: - print(f'引擎已有配置，跳过') + # 直接覆盖为最新配置 + engine.configuration = new_config + engine.save(update_fields=['configuration']) + print(f'已更新引擎配置: {engine.name}') else: # 创建引擎 engine = ScanEngine.objects.create( name='full scan', - configuration=yaml_path.read_text(), + configuration=new_config, ) print(f'已创建引擎: {engine.name}') " diff --git a/docker/server/start.sh b/docker/server/start.sh index 38c913b2..6bcd22f4 100644 --- a/docker/server/start.sh +++ b/docker/server/start.sh @@ -10,7 +10,7 @@ python manage.py migrate --noinput echo " ✓ 数据库迁移完成" echo " [1.1/3] 初始化默认扫描引擎..." -python manage.py init_default_engine +python manage.py init_default_engine --force echo " ✓ 默认扫描引擎已就绪" echo " [1.2/3] 初始化默认目录字典..." diff --git a/update.sh b/update.sh index c4741e1a..c43cd9b3 100755 --- a/update.sh +++ b/update.sh @@ -21,8 +21,8 @@ cd "$(dirname "$0")" # 权限检查 if [ "$EUID" -ne 0 ]; then - echo -e "\033[0;31m[错误] 请使用 sudo 运行此脚本\033[0m" - echo -e " 正确用法: \033[1msudo ./update.sh\033[0m" + printf "\033[0;31m✗ 请使用 sudo 运行此脚本\033[0m\n" + printf " 正确用法: \033[1msudo ./update.sh\033[0m\n" exit 1 fi @@ -49,9 +49,17 @@ YELLOW='\033[1;33m' RED='\033[0;31m' BLUE='\033[0;34m' CYAN='\033[0;36m' +DIM='\033[2m' BOLD='\033[1m' NC='\033[0m' +# 日志函数 +log_step() { printf "${CYAN}▶${NC} %s\n" "$1"; } +log_ok() { printf " ${GREEN}✓${NC} %s\n" "$1"; } +log_info() { printf " ${DIM}→${NC} %s\n" "$1"; } +log_warn() { printf " ${YELLOW}!${NC} %s\n" "$1"; } +log_error() { printf "${RED}✗${NC} %s\n" "$1"; } + # 合并 .env 新配置项（保留用户已有值） merge_env_config() { local example_file="docker/.env.example" @@ -70,58 +78,68 @@ merge_env_config() { if ! grep -q "^${key}=" "$env_file"; then printf '%s\n' "$line" >> "$env_file" - echo -e " ${GREEN}+${NC} 新增: $key" + log_info "新增配置: $key" ((new_keys++)) fi done < "$example_file" if [ $new_keys -gt 0 ]; then - echo -e " ${GREEN}OK${NC} 已添加 $new_keys 个新配置项" + log_ok "已添加 $new_keys 个新配置项" else - echo -e " ${GREEN}OK${NC} 配置已是最新" + log_ok "配置已是最新" fi } -echo "" -echo -e "${BOLD}${BLUE}╔════════════════════════════════════════╗${NC}" +# 显示标题 +printf "\n" +printf "${BOLD}${BLUE}┌────────────────────────────────────────┐${NC}\n" if [ "$DEV_MODE" = true ]; then - echo -e "${BOLD}${BLUE}║ 开发环境更新（本地构建） ║${NC}" + printf "${BOLD}${BLUE}│${NC} ${BOLD}XingRin 系统更新${NC} ${BOLD}${BLUE}│${NC}\n" + printf "${BOLD}${BLUE}│${NC} ${DIM}开发模式 · 本地构建${NC} ${BOLD}${BLUE}│${NC}\n" else - echo -e "${BOLD}${BLUE}║ 生产环境更新（Docker Hub） ║${NC}" + printf "${BOLD}${BLUE}│${NC} ${BOLD}XingRin 系统更新${NC} ${BOLD}${BLUE}│${NC}\n" + printf "${BOLD}${BLUE}│${NC} ${DIM}生产模式 · Docker Hub${NC} ${BOLD}${BLUE}│${NC}\n" fi -echo -e "${BOLD}${BLUE}╚════════════════════════════════════════╝${NC}" -echo "" +printf "${BOLD}${BLUE}└────────────────────────────────────────┘${NC}\n" +printf "\n" -# 测试性功能警告 -echo -e "${BOLD}${YELLOW}[!] 警告：此功能为测试性功能，可能会导致升级失败${NC}" -echo -e "${YELLOW} 建议运行 ./uninstall.sh 后重新clone最新代码进行全新安装${NC}" -echo "" -echo -n -e "${YELLOW}是否继续更新？(y/N) ${NC}" +# 警告提示 +printf "${YELLOW}┌─ 注意事项 ─────────────────────────────┐${NC}\n" +printf "${YELLOW}│${NC} • 此功能为测试性功能，可能导致升级失败 ${YELLOW}│${NC}\n" +printf "${YELLOW}│${NC} • 升级会覆盖所有默认引擎配置 ${YELLOW}│${NC}\n" +printf "${YELLOW}│${NC} • 自定义配置请先备份或创建新引擎 ${YELLOW}│${NC}\n" +printf "${YELLOW}│${NC} • 推荐：卸载后重新安装以获得最佳体验 ${YELLOW}│${NC}\n" +printf "${YELLOW}└────────────────────────────────────────┘${NC}\n" +printf "\n" + +printf "${YELLOW}是否继续更新？${NC} [y/N] " read -r ans_continue ans_continue=${ans_continue:-N} if [[ ! $ans_continue =~ ^[Yy]$ ]]; then - echo -e "${CYAN}已取消更新。${NC}" + printf "\n${DIM}已取消更新${NC}\n" exit 0 fi -echo "" +printf "\n" # Step 1: 停止服务 -echo -e "${CYAN}[1/5]${NC} 停止服务..." -./stop.sh 2>&1 | sed 's/^/ /' +log_step "停止服务..." +./stop.sh 2>&1 | sed 's/^/ /' +log_ok "服务已停止" # Step 2: 拉取代码 -echo "" -echo -e "${CYAN}[2/5]${NC} 拉取代码..." -git pull --rebase 2>&1 | sed 's/^/ /' -if [ $? -ne 0 ]; then - echo -e "${RED}[错误]${NC} git pull 失败，请手动解决冲突后重试" +printf "\n" +log_step "拉取最新代码..." +if git pull --rebase 2>&1 | sed 's/^/ /'; then + log_ok "代码已更新" +else + log_error "git pull 失败，请手动解决冲突后重试" exit 1 fi # Step 3: 检查配置更新 + 版本同步 -echo "" -echo -e "${CYAN}[3/5]${NC} 检查配置更新..." +printf "\n" +log_step "同步配置..." merge_env_config # 版本同步：从 VERSION 文件更新 IMAGE_TAG @@ -130,21 +148,20 @@ if [ -f "VERSION" ]; then if [ -n "$NEW_VERSION" ]; then if grep -q "^IMAGE_TAG=" "docker/.env"; then sed_inplace "s/^IMAGE_TAG=.*/IMAGE_TAG=$NEW_VERSION/" "docker/.env" - echo -e " ${GREEN}+${NC} 版本同步: IMAGE_TAG=$NEW_VERSION" else printf '%s\n' "IMAGE_TAG=$NEW_VERSION" >> "docker/.env" - echo -e " ${GREEN}+${NC} 新增版本: IMAGE_TAG=$NEW_VERSION" fi + log_ok "版本同步: $NEW_VERSION" fi fi # Step 4: 构建/拉取镜像 -echo "" -echo -e "${CYAN}[4/5]${NC} 更新镜像..." +printf "\n" +log_step "更新镜像..." if [ "$DEV_MODE" = true ]; then # 开发模式：本地构建所有镜像（包括 Worker） - echo -e " 构建 Worker 镜像..." + log_info "构建 Worker 镜像..." # 读取 IMAGE_TAG IMAGE_TAG=$(grep "^IMAGE_TAG=" "docker/.env" | cut -d'=' -f2) @@ -153,24 +170,23 @@ if [ "$DEV_MODE" = true ]; then fi # 构建 Worker 镜像（Worker 是临时容器，不在 compose 中，需要单独构建） - docker build -t docker-worker -f docker/worker/Dockerfile . 2>&1 | sed 's/^/ /' - docker tag docker-worker docker-worker:${IMAGE_TAG} 2>&1 | sed 's/^/ /' - echo -e " ${GREEN}OK${NC} Worker 镜像已构建: docker-worker:${IMAGE_TAG}" + docker build -t docker-worker -f docker/worker/Dockerfile . 2>&1 | sed 's/^/ /' + docker tag docker-worker docker-worker:${IMAGE_TAG} 2>&1 | sed 's/^/ /' + log_ok "Worker 镜像: docker-worker:${IMAGE_TAG}" - # 其他服务镜像由 start.sh --dev 构建 - echo -e " 其他服务镜像将在启动时构建..." + log_info "其他服务镜像将在启动时构建" else - # 生产模式：镜像由 start.sh 拉取 - echo -e " 镜像将在启动时从 Docker Hub 拉取..." + log_info "镜像将在启动时从 Docker Hub 拉取" fi # Step 5: 启动服务 -echo "" -echo -e "${CYAN}[5/5]${NC} 启动服务..." +printf "\n" +log_step "启动服务..." ./start.sh "$@" -echo "" -echo -e "${BOLD}${GREEN}════════════════════════════════════════${NC}" -echo -e "${BOLD}${GREEN} 更新完成！${NC}" -echo -e "${BOLD}${GREEN}════════════════════════════════════════${NC}" -echo "" +# 完成提示 +printf "\n" +printf "${GREEN}┌────────────────────────────────────────┐${NC}\n" +printf "${GREEN}│${NC} ${BOLD}${GREEN}✓${NC} ${BOLD}更新完成${NC} ${GREEN}│${NC}\n" +printf "${GREEN}└────────────────────────────────────────┘${NC}\n" +printf "\n"