# -*- coding: utf-8 -*- """ 指令构建器 统一的 CTF 题目生成指令构建逻辑，供各 AI 服务共用 """ from typing import Dict, List, Any, Union, Optional import logging logger = logging.getLogger(__name__) class InstructionBuilder: """CTF 题目生成指令构建器 统一处理指令构建、场景格式化等逻辑 """ @classmethod def format_scene(cls, scene: Union[str, Dict]) -> str: """格式化场景数据为可读字符串 Args: scene: 场景数据（可能是字符串或字典） Returns: 格式化后的场景字符串 """ if not scene: return '无特定场景' if isinstance(scene, dict): # 检查是否明确标记为无场景 if scene.get('no_scene'): return '无特定场景' scene_name = scene.get('name', '未知场景') # 检查是否有子场景 sub_scene = scene.get('sub_scene') if sub_scene: if isinstance(sub_scene, dict): sub_scene_name = sub_scene.get('name', '') if sub_scene_name: return f"{scene_name} - {sub_scene_name}" else: return f"{scene_name} - {sub_scene}" return scene_name else: return str(scene) @classmethod def format_vulnerabilities(cls, vulnerabilities: List[Any]) -> str: """格式化漏洞列表为字符串 Args: vulnerabilities: 漏洞列表（可能是字符串列表或字典列表） Returns: 格式化后的漏洞字符串 """ vuln_names = [] for vuln in vulnerabilities: if isinstance(vuln, dict): vuln_names.append(vuln.get('name', str(vuln))) else: vuln_names.append(str(vuln)) return '、'.join(vuln_names) @classmethod def build_simple_instruction( cls, language: str, vulnerabilities: List[Any], scene: Union[str, Dict], difficulty: str, extra_requirements: str = '' ) -> str: """构建简化指令（用于 Augment CLI） AI 已经通过 guidelines 知道流程，只需要提供核心参数 Args: language: 编程语言 vulnerabilities: 漏洞列表 scene: 场景 difficulty: 难度级别 extra_requirements: 用户额外要求 Returns: 指令文本 """ vuln_str = cls.format_vulnerabilities(vulnerabilities) scene_str = cls.format_scene(scene) instruction = f"""请严格按照规则文件中的流程，完成 CTF 题目的设计和生成。 ## 用户需求 - 难度：{difficulty} - 漏洞类型：{vuln_str} - 编程语言：{language} - 应用场景：{scene_str} """ if extra_requirements and extra_requirements.strip(): instruction += f"- 额外要求：{extra_requirements.strip()}\n" instruction += """ ## 执行要求 1. 严格按照规则文件中定义的阶段顺序执行 2. 每个阶段开始时输出阶段标题 3. 不要跳过任何阶段 4. 所有文件保存到 output/ 目录 """ return instruction @classmethod def build_claude_instruction( cls, language: str, vulnerabilities: List[Any], scene: Union[str, Dict], difficulty: str, extra_requirements: str = '' ) -> str: """构建 Claude CLI 指令 包含 CLAUDE.md 引导和完成所有阶段的警告 Args: language: 编程语言 vulnerabilities: 漏洞列表 scene: 场景 difficulty: 难度级别 extra_requirements: 用户额外要求 Returns: 指令文本 """ vuln_str = cls.format_vulnerabilities(vulnerabilities) scene_str = cls.format_scene(scene) instruction = f"""请按照 CLAUDE.md 中的指导，完成 CTF 题目设计的所有阶段。 用户需求： - 语言：{language} - 漏洞：{vuln_str} - 场景：{scene_str} - 难度：{difficulty} """ if extra_requirements and extra_requirements.strip(): instruction += f"- 额外要求：{extra_requirements.strip()}\n" instruction += "\n⚠️ 重要：必须完成所有阶段，包括代码生成、Docker环境构建和成品输出。不要在中途停止。" return instruction @classmethod def build_api_instruction( cls, language: str, vulnerabilities: List[Any], scene: Union[str, Dict], difficulty: str, extra_requirements: str = '', has_scene_field: bool = True # 新增参数：表单中是否有场景字段 ) -> str: """构建 API 模式指令（用于 AIService） Args: language: 编程语言 vulnerabilities: 漏洞列表 scene: 场景 difficulty: 难度级别 extra_requirements: 用户额外要求 has_scene_field: 表单配置中是否有场景字段（默认True保持向后兼容） Returns: 指令文本 """ vuln_str = cls.format_vulnerabilities(vulnerabilities) instruction = f"""- 难度：{difficulty} - 漏洞类型：{vuln_str} - 编程语言：{language}""" # 只在表单中有场景字段且有有效场景值时才显示 if has_scene_field: scene_name = cls.format_scene(scene) # 如果场景是"无特定场景"或标记为无场景，则不显示 if scene_name and scene_name != '无特定场景': if isinstance(scene, dict) and scene.get('no_scene'): pass # 不显示 elif scene_name != '无场景': instruction += f"\n- 应用场景：{scene_name}" instruction += "\n" if extra_requirements and extra_requirements.strip(): instruction += f"- 额外要求：{extra_requirements.strip()}\n" return instruction @classmethod def build_dynamic_instruction( cls, category_id: str, form_data: Dict[str, Any], language: Optional[str] = None, vulnerabilities: Optional[List[Any]] = None, scene: Optional[Union[str, Dict]] = None, difficulty: Optional[str] = None, extra_requirements: str = '' ) -> str: """根据表单配置动态构建指令（支持所有表单字段） 这个方法会遍历表单配置中的所有字段，动态生成指令内容， 而不是只使用固定的几个字段。这样即使表单字段变化，也能正确传递所有信息。 Args: category_id: 方向ID form_data: 完整的表单数据 language: 编程语言（可选，用于向后兼容） vulnerabilities: 漏洞列表（可选，用于向后兼容） scene: 场景（可选，用于向后兼容） difficulty: 难度（可选，用于向后兼容） extra_requirements: 额外要求（可选，用于向后兼容） Returns: 指令文本 """ try: from app.models.database.models import CategoryConfig category = CategoryConfig.query.get(category_id) if not category: logger.warning(f"方向配置不存在: {category_id}，回退到简单指令构建") # 回退到简单指令构建 if language and vulnerabilities and difficulty: return cls.build_simple_instruction( language, vulnerabilities, scene or '无特定场景', difficulty, extra_requirements ) else: return "请按照规则文件中的流程，完成 CTF 题目的设计和生成。" form_fields = category.get_form_fields() instruction_lines = [] # 创建一个字段ID到字段配置的映射，用于快速查找 field_config_map = {field.get('id'): field for field in form_fields} # 创建一个已处理的字段ID集合，避免重复处理 processed_field_ids = set() # 首先，按照表单字段配置顺序，动态生成指令（优先使用配置中的字段） for field in form_fields: field_id = field.get('id') field_label = field.get('label', field_id) # 使用字段标签或ID field_type = field.get('type') field_visible = field.get('visible', True) # 默认可见 field_hidden = field.get('hidden', False) # 是否隐藏 # 跳过不可见的字段、隐藏字段和特殊字段 if not field_visible or field_hidden or field_id in ['extra_requirements']: continue # 获取字段值 field_value = form_data.get(field_id) if field_value is None or field_value == '': continue # 跳过空值 # 根据字段类型格式化值 formatted_value = None if field_type in ['multi_select', 'multi_select_categorized']: # 多选字段：格式化为逗号分隔的字符串 if isinstance(field_value, list): formatted_value = ', '.join([str(v.get('name', v) if isinstance(v, dict) else v) for v in field_value if v]) elif isinstance(field_value, str): formatted_value = field_value else: formatted_value = str(field_value) elif field_type == 'select': # 单选字段：可能是字符串、字典或选项值 if isinstance(field_value, dict): formatted_value = field_value.get('name', str(field_value)) else: # 如果是选项值，尝试从字段配置中找到对应标签 options = field.get('options', []) for opt in options: opt_value = opt.get('value') if isinstance(opt, dict) else opt if opt_value == field_value: formatted_value = opt.get('label', opt.get('value', str(opt))) if isinstance(opt, dict) else str(opt) break if not formatted_value: formatted_value = str(field_value) elif field_type == 'select_with_sub': # 带子选项的选择：可能是字典 if isinstance(field_value, dict): # 检查是否标记为无场景 if field_value.get('no_scene'): continue # 跳过无场景的情况 scene_name = field_value.get('name', str(field_value)) # 如果场景名称为"无场景"或空，也跳过 if not scene_name or scene_name.strip() in ['无场景', '无特定场景', 'none', 'None', '']: continue # 确保场景名称不是空字符串 if not scene_name or not scene_name.strip(): continue sub_scene = field_value.get('sub_scene') if sub_scene: if isinstance(sub_scene, dict): sub_scene_name = sub_scene.get('name', '') if sub_scene_name: formatted_value = f"{scene_name} - {sub_scene_name}" else: formatted_value = scene_name else: formatted_value = f"{scene_name} - {sub_scene}" else: formatted_value = scene_name else: # 如果是字符串，检查是否为"无场景" if str(field_value).strip() in ['无场景', '无特定场景', 'none', 'None', '']: continue formatted_value = str(field_value) elif field_type == 'textarea': # 文本域：直接使用值 formatted_value = str(field_value).strip() else: # 其他类型：直接转换为字符串 formatted_value = str(field_value) # 如果有格式化后的值，添加到指令中 if formatted_value and formatted_value.strip(): # 跳过"无场景"、"无特定场景"等无意义的值 if formatted_value.strip() in ['无场景', '无特定场景', 'none', 'None', '']: continue instruction_lines.append(f"- {field_label}：{formatted_value}") processed_field_ids.add(field_id) # 标记为已处理 # 然后，处理 form_data 中存在但配置中不存在的字段（兜底处理） # 这样可以确保即使配置更新不及时，也能传递所有数据 for field_id, field_value in form_data.items(): # 跳过已处理的字段 if field_id in processed_field_ids: continue # 跳过空值 if field_value is None or field_value == '': continue # 跳过特殊字段 if field_id in ['extra_requirements']: continue # 如果字段不在配置中，使用字段ID作为标签，直接格式化值 formatted_value = None if isinstance(field_value, list): formatted_value = ', '.join([str(v.get('name', v) if isinstance(v, dict) else v) for v in field_value if v]) elif isinstance(field_value, dict): formatted_value = field_value.get('name', str(field_value)) else: formatted_value = str(field_value).strip() if formatted_value and formatted_value.strip(): # 跳过"无场景"、"无特定场景"等无意义的值 if formatted_value.strip() in ['无场景', '无特定场景', 'none', 'None', '']: continue instruction_lines.append(f"- {field_id}：{formatted_value}") # 如果有额外要求，单独添加（避免重复） if extra_requirements and extra_requirements.strip(): # 检查是否已经在指令中（避免重复） if not any('额外要求' in line or 'extra' in line.lower() for line in instruction_lines): instruction_lines.append(f"- 额外要求：{extra_requirements.strip()}") # 构建完整指令 if instruction_lines: instruction = """请严格按照规则文件中的流程，完成 CTF 题目的设计和生成。 ## 用户需求 """ instruction += '\n'.join(instruction_lines) instruction += """ ## 执行要求 1. 严格按照规则文件中定义的阶段顺序执行 2. 每个阶段开始时输出阶段标题 3. 不要跳过任何阶段 4. 所有文件保存到 output/ 目录 """ else: # 如果没有提取到任何字段，回退到简单指令 logger.warning(f"未能从表单数据中提取到任何字段，回退到简单指令构建") if language and vulnerabilities and difficulty: return cls.build_simple_instruction( language, vulnerabilities, scene or '无特定场景', difficulty, extra_requirements ) else: instruction = "请按照规则文件中的流程，完成 CTF 题目的设计和生成。" return instruction except Exception as e: logger.error(f"动态构建指令时出错: {e}", exc_info=True) # 出错时回退到简单指令构建 if language and vulnerabilities and difficulty: return cls.build_simple_instruction( language, vulnerabilities, scene or '无特定场景', difficulty, extra_requirements ) else: return "请按照规则文件中的流程，完成 CTF 题目的设计和生成。"