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
2026-01-24 08:13:10 +08:00 · 2026-01-10 11:04:42 +08:00
parent 255d505aba
commit 8a6f1b6f24
4 changed files with 93 additions and 61 deletions
--- 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})'
+                ))
--- 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}')
 "
--- 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] 初始化默认目录字典..."
--- 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"