# -*- coding: utf-8 -*- """ Augment CLI 服务适配器 这个模块提供了与 Augment CLI 的集成，使其能够像 OpenAI API 一样被调用。 """ import os import subprocess import tempfile import json import time import shutil import re from typing import Dict, List, Any, Optional, Union from pathlib import Path from .stage_detector import StageDetector from .instruction_builder import InstructionBuilder from .guideline_manager import GuidelineManager class AugmentService: """Augment CLI 服务适配器""" def __init__(self, node_path: Optional[str] = None, user_id: Optional[int] = None, verbose: bool = False, model: Optional[str] = None, **kwargs): """初始化 Augment 服务 Args: node_path: Node.js 可执行文件路径，如果不提供则使用默认路径 user_id: 用户ID，用于获取用户的Augment Token verbose: 是否显示详细输出（默认 False，只显示关键信息） model: 模型ID（可选），如果不提供则使用默认模型 """ # Node.js 路径配置 self.node_path = node_path or os.environ.get( "NODE_PATH", "$HOME/.nvm/versions/node/v22.13.1/bin" ) # 项目根目录（ctf目录） # 从 augment_service.py 向上查找：augment_service.py -> core -> ai_driver -> services -> app -> ctf # 路径: ctf/app/services/ai_driver/core/augment_service.py # parent(1): core, parent(2): ai_driver, parent(3): services, parent(4): app, parent(5): ctf self.project_root = Path(__file__).parent.parent.parent.parent.parent # ge10 工作目录（Augment 的工作目录），统一使用 ctf/ge10 self.ge10_dir = self.project_root / "ge10" if not self.ge10_dir.exists(): raise Exception(f"ge10 目录不存在: {self.ge10_dir}") # 用户ID（用于获取Token） self.user_id = user_id # BaseAIProvider 兼容属性 self.provider_type = 'augment' self.provider_name = 'Augment CLI' self.model = model or 'augment' # 使用传入的模型或默认模型 self.base_url = None # Augment 不使用 API，使用 CLI # 日志回调函数 self.log_callback = None # 超时恢复标记（用于标记AI超时但文件已生成的情况） self.timeout_recovered = False self.timeout_warning = None # 是否显示详细输出 self.verbose = verbose # 任务ID（用于日志文件隔离） self.task_id = None # 初始化日志记录器 import logging self.logger = logging.getLogger(__name__) # 不再在初始化时验证 Augment CLI（避免每次创建服务都检查版本） # 如果 CLI 不可用，在生成题目时会自然失败 # self._verify_augment_cli() self.logger.info("Augment 服务初始化完成") def _verify_augment_cli(self): """验证 Augment CLI 是否可用""" try: # 设置 PATH 环境变量 env = os.environ.copy() env['PATH'] = f"{self.node_path}:{env.get('PATH', '')}" # 检查 auggie 命令是否可用 result = subprocess.run( ['auggie', '--version'], capture_output=True, text=True, env=env, timeout=10 ) if result.returncode != 0: raise Exception(f"Augment CLI 不可用: {result.stderr}") print(f"✅ Augment CLI 版本: {result.stdout.strip()}") except FileNotFoundError: raise Exception("未找到 auggie 命令，请确保 Augment CLI 已安装") except subprocess.TimeoutExpired: raise Exception("Augment CLI 验证超时") except Exception as e: raise Exception(f"Augment CLI 验证失败: {str(e)}") def set_log_callback(self, callback): """设置日志回调函数 Args: callback: 日志回调函数，接受 (step_num, role, content) 参数 """ self.log_callback = callback self.logger.debug("已设置日志回调函数") def start_conversation(self, system_prompt: Optional[str] = None): """开始一个新的对话 注意：Augment 使用 guidelines 而不是 system_prompt， 所以这个方法主要是为了兼容性 Args: system_prompt: 系统提示（在 Augment 中会被忽略） """ if system_prompt and self.log_callback: self.log_callback(0, "system", "Augment 使用 .augment/rules/ 中的 guidelines，不使用 system_prompt") self.logger.info("Augment 对话开始（使用 guidelines）") def add_message(self, role: str, content: str): """添加消息到对话历史 注意：Augment 不维护对话历史，每次调用都是独立的 这个方法主要是为了兼容性 Args: role: 消息角色 content: 消息内容 """ # 仅用于日志记录 if self.log_callback: self.log_callback(0, role, content) def generate_ctf_challenge( self, language: str, vulnerabilities: List[str], scene: Union[str, Dict[str, Any]], difficulty: str, extra_requirements: str = '', category_id: str = 'web', form_data: Optional[Dict[str, Any]] = None, task_id: Optional[str] = None ) -> Dict[str, Any]: """生成 CTF 题目（Augment 专用接口） Args: language: 编程语言 vulnerabilities: 漏洞名称列表（二级分类名称） scene: 场景名称（可能是字符串或字典） difficulty: 难度级别 extra_requirements: 用户额外要求 category_id: 方向ID form_data: 完整的表单数据（可选，用于传递所有表单字段） task_id: 任务ID（可选，用于创建独立工作目录避免多任务冲突） Returns: 包含生成结果的字典 """ try: # 初始化阶段配置（从 generation_statuses 获取） if task_id: from app.routes.generator.utils import generation_statuses, generation_lock with generation_lock: task_status = generation_statuses.get(task_id, {}) stage_names = task_status.get('stage_names', {}) total_stages = task_status.get('total_stages', 0) if stage_names: StageDetector.set_stage_names(stage_names) self.logger.info(f"已设置动态阶段配置: {len(stage_names)} 个阶段") print(f"📊 已加载 {len(stage_names)} 个阶段配置: {list(stage_names.values())}", flush=True) # 确保 total_stages 已设置 if total_stages == 0 and stage_names: total_stages = len(stage_names) task_status['total_stages'] = total_stages # 初始化所有阶段为 waiting 状态 if 'step_statuses' not in task_status: task_status['step_statuses'] = {i: 'waiting' for i in range(total_stages)} self.logger.info(f"已初始化 {total_stages} 个阶段的状态") else: self.logger.warning("未找到阶段配置，使用默认阶段") # 尝试从数据库加载 Prompt 并设置为规则文件 db_prompt_loaded = self._setup_prompt_from_database(category_id, difficulty, language, vulnerabilities, scene) if not db_prompt_loaded: # 如果数据库没有配置，回退到文件系统的 guideline print("⚠️ 数据库中未找到 Prompt 配置，使用默认 guideline 文件", flush=True) # 注意：GuidelineManager 会尝试复制文件到 .augment/rules/，但如果目录不存在会失败 # 这里只检查文件是否存在，不创建目录（避免多任务冲突） rules_dir = self.ge10_dir / '.augment' / 'rules' default_rules_file = rules_dir / 'ctf-generation-guide.md' if default_rules_file.exists(): # 如果文件已存在，使用它（可能是之前创建的） self._current_rules_file = str(default_rules_file.resolve()) print(f"📜 使用已存在的规则文件: {self._current_rules_file}", flush=True) else: # 如果文件不存在，尝试通过 GuidelineManager 设置（但不会创建目录） # 如果目录不存在，GuidelineManager 会失败，这是预期的行为 self._setup_guideline_for_difficulty(difficulty) if default_rules_file.exists(): self._current_rules_file = str(default_rules_file.resolve()) print(f"📜 已设置默认规则文件: {self._current_rules_file}", flush=True) else: print(f"⚠️ 无法设置默认规则文件（目录可能不存在）", flush=True) # 构建指令（传递 form_data 以支持动态字段） instruction = self._build_instruction(language, vulnerabilities, scene, difficulty, extra_requirements, category_id, form_data) # 打印实际传递给 Augment 的指令 print("\n" + "="*60) print("📝 传递给 Augment AI 的指令:") print("="*60) print(instruction) print("="*60 + "\n") # 存储 task_id 用于日志文件隔离 self.task_id = task_id # 调用 Augment CLI（传入 category_id 和 task_id 以设置正确的工作目录） response = self._call_augment(instruction, category_id=category_id, task_id=task_id) # 直接使用 output/ 目录，不再需要移动临时目录 if task_id: # 工作目录已经在 output/ 下，直接使用 output_base_dir = self.ge10_dir / category_id / "output" # 查找包含 task_id 的目录（格式：YYYYMMDD_HHMMSS_taskid） if output_base_dir.exists(): matching_dirs = [d for d in output_base_dir.iterdir() if d.is_dir() and task_id in d.name] if matching_dirs: # 找到最新的匹配目录 latest_dir = max(matching_dirs, key=lambda d: d.stat().st_mtime) output_dir = str(latest_dir) else: # 如果没有找到，尝试查找最新的目录 challenge_dirs = [d for d in output_base_dir.iterdir() if d.is_dir()] if challenge_dirs: latest_challenge_dir = max(challenge_dirs, key=lambda d: d.stat().st_mtime) output_dir = str(latest_challenge_dir) else: output_dir = self._find_latest_output_dir(category_id) else: output_dir = self._find_latest_output_dir(category_id) else: # 查找生成的输出目录（根据 category_id 查找对应的目录） output_dir = self._find_latest_output_dir(category_id) # 返回结果 return { "status": "success", "response": response, "output_dir": output_dir } except Exception as e: # 捕获所有异常，返回失败状态 error_msg = str(e) print(f"❌ 生成CTF题目失败: {error_msg}") return { "status": "error", "message": error_msg, "response": None, "output_dir": None } def chat_completion( self, messages: List[Dict[str, str]], model: str = "augment", temperature: float = 0.7, max_tokens: Optional[int] = None, stream: bool = False ) -> Dict[str, Any]: """执行聊天补全（兼容 OpenAI API 接口） 注意：这个方法主要用于兼容性，实际使用时建议使用 generate_ctf_challenge Args: messages: 消息列表 model: 模型名称（在 Augment 中被忽略） temperature: 温度参数（在 Augment 中被忽略） max_tokens: 最大 token 数（在 Augment 中被忽略） stream: 是否流式输出（在 Augment 中被忽略） Returns: 包含 AI 响应的字典 """ # 提取用户消息（最后一条非系统消息） user_message = None for msg in reversed(messages): if msg['role'] == 'user': user_message = msg['content'] break if not user_message: raise ValueError("未找到用户消息") # 调用 Augment CLI response = self._call_augment(user_message) # 返回兼容 OpenAI API 的响应格式 return { "choices": [{ "message": { "role": "assistant", "content": response }, "finish_reason": "stop" }], "usage": { "prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0 } } def _setup_guideline_for_difficulty(self, difficulty: str): """根据难度选择对应的guideline文件（委托给 GuidelineManager）""" GuidelineManager.setup_for_augment(self.ge10_dir, difficulty, verbose=True) def _setup_prompt_from_database( self, category_id: str, difficulty: str, language: str, vulnerabilities: List[str], scene: str ) -> bool: """从数据库加载 Prompt 并设置为 Augment 规则文件 Args: category_id: 方向ID difficulty: 难度级别 language: 编程语言 vulnerabilities: 漏洞列表 scene: 场景 Returns: 是否成功加载 """ try: from flask import has_app_context from app.models.database.models import CategoryConfig from app.services.prompt.compiler_service import PromptCompilerService if not has_app_context(): print("⚠️ 不在 Flask 应用上下文中，无法访问数据库", flush=True) return False category = CategoryConfig.query.get(category_id) if not category: print(f"⚠️ 方向配置不存在: {category_id}", flush=True) return False # 难度映射 difficulty_map = { '入门': 'beginner', '简单': 'easy', '中等': 'medium', '困难': 'hard' } difficulty_key = difficulty_map.get(difficulty, 'beginner') # 尝试获取已编译的 Prompt compiled_prompts = category.get_compiled_prompts() compiled_prompt = None if compiled_prompts and difficulty_key in compiled_prompts: old_prompt = compiled_prompts[difficulty_key] # 检查是否包含旧的格式（"阶段X末"），如果包含则重新编译 # 使用正则表达式检测 "阶段X末" 或 "阶段 X 末" 格式 import re old_format_pattern = r'阶段\s*\d+\s*末' # 最多重试3次，确保清理干净 max_retries = 3 retry_count = 0 while re.search(old_format_pattern, old_prompt) and retry_count < max_retries: retry_count += 1 print(f"⚠️ 检测到旧的 Prompt 格式（包含'阶段X末'），第 {retry_count} 次重新编译...", flush=True) # 重新编译并保存到数据库 if PromptCompilerService.compile_and_save_prompts(category): # 刷新数据库会话，确保获取最新数据 from app.models.database.models import db db.session.refresh(category) # 重新获取编译后的 prompt compiled_prompts = category.get_compiled_prompts() if compiled_prompts and difficulty_key in compiled_prompts: old_prompt = compiled_prompts[difficulty_key] # 验证是否还有旧格式 if not re.search(old_format_pattern, old_prompt): print(f"✅ 重新编译成功，已清理旧格式并保存到数据库", flush=True) break else: print(f"⚠️ 重新编译后仍包含旧格式，继续重试...", flush=True) else: print(f"⚠️ 重新编译后仍未找到 Prompt", flush=True) break else: print(f"⚠️ 重新编译失败", flush=True) break if retry_count >= max_retries and re.search(old_format_pattern, old_prompt): print(f"⚠️ 警告：重新编译 {max_retries} 次后仍包含旧格式，可能 stage 配置中包含'阶段X末'内容", flush=True) # 使用已编译的模板，替换占位符 compiled_prompt = PromptCompilerService.replace_placeholders( old_prompt, language, vulnerabilities, scene ) print(f"✅ 已从数据库加载 Prompt 模板（方向: {category_id}, 难度: {difficulty}）", flush=True) if not compiled_prompt: print(f"⚠️ 数据库中未找到方向 {category_id} 难度 {difficulty} 的 Prompt，尝试重新编译...", flush=True) # 尝试重新编译 if PromptCompilerService.compile_and_save_prompts(category): compiled_prompts = category.get_compiled_prompts() if compiled_prompts and difficulty_key in compiled_prompts: compiled_prompt = PromptCompilerService.replace_placeholders( compiled_prompts[difficulty_key], language, vulnerabilities, scene ) print(f"✅ 重新编译成功，已加载 Prompt 模板", flush=True) if not compiled_prompt: print(f"⚠️ 数据库中未找到方向 {category_id} 难度 {difficulty} 的 Prompt", flush=True) return False # 不再创建 .augment/rules/ 目录，规则文件通过 -if 参数传递给 Augment CLI # 规则文件会在 _call_augment 中作为临时文件创建并传递给 Augment CLI print(f"✅ 已加载 Prompt 模板，长度: {len(compiled_prompt)} 字符", flush=True) # 保存编译后的 prompt 内容，供后续使用 self._compiled_prompt = compiled_prompt return True except Exception as e: print(f"❌ 从数据库加载 Prompt 失败: {str(e)}", flush=True) import traceback traceback.print_exc() return False def _build_instruction( self, language: str, vulnerabilities: List[str], scene: str, difficulty: str, extra_requirements: str = '', category_id: str = 'web', form_data: Optional[Dict[str, Any]] = None ) -> str: """构建 Augment 指令（委托给 InstructionBuilder）""" # 如果有 form_data，使用动态指令构建 if form_data: return InstructionBuilder.build_dynamic_instruction( category_id, form_data, language, vulnerabilities, scene, difficulty, extra_requirements ) else: # 回退到简单指令构建（向后兼容） return InstructionBuilder.build_simple_instruction( language, vulnerabilities, scene, difficulty, extra_requirements ) def _detect_stage(self, line: str) -> Optional[int]: """检测输出行中的阶段信息（委托给 StageDetector）""" return StageDetector.detect_stage(line) def _should_print_line(self, line: str) -> bool: """判断是否应该打印这一行（非详细模式下） Args: line: 输出行 Returns: 是否应该打印 """ import re # 如果是详细模式，打印所有内容 if self.verbose: return True # 先过滤掉明确不需要的内容 # 1. 进度计数 if re.search(r'\[已输出 \d+ 行\]', line) or re.search(r'\[已处理 \d+ 行\]', line): return False # 2. 空行或只有空格的行 if not line.strip(): return False # 3. ANSI 颜色代码（灰色调试信息） if re.search(r'\[90m', line): return False # 4. 深度缩进的代码片段（4个或更多空格） if re.search(r'^\s{4,}[^\s]', line): return False # 5. 工具调用的详细参数（通常很长） if re.search(r'^\s+(command|cwd|wait|max_wait_seconds|path|type):', line): return False # 6. 代码行号（cat -n 的输出） if re.search(r'^\s+\d+\s+', line): return False # 关键信息模式：打印重要的行 # 1. 阶段标记（最重要）- 支持有无#号的格式 if re.search(r'(?:##\s*)?阶段\s*\d+[：:]', line): return True # 2. 工具调用（只显示工具名称） if '🔧 Tool call:' in line: return True # 3. 工具结果（只显示状态） if '📋 Tool result:' in line: return True # 4. 命令完成状态 if '✅ Command completed' in line or '✅ Successfully' in line: return True # 5. 错误和警告 if any(marker in line for marker in ['❌', '⚠️', 'Error', 'error', 'ERROR', 'Warning', 'WARNING']): return True # 6. 文件保存通知 if 'File saved' in line or 'Saved file' in line: return True # 7. Docker 相关的关键操作 docker_keywords = [ 'Building', 'Successfully built', 'Successfully tagged', 'Creating network', 'Creating container', 'Starting container', 'Container started', 'Stopping', 'Removing' ] if any(keyword in line for keyword in docker_keywords): return True # 8. AI 的思考过程（🤖 标记） if '🤖' in line: # 只显示简短的思考，过滤掉长篇大论 if len(line.strip()) < 100: return True return False # 默认不显示 return False def _extract_stage_info(self, stage: int, output_lines: list) -> Optional[Dict]: """从输出中提取阶段关键信息（委托给 StageDetector，支持动态阶段配置）""" output_text = ''.join(output_lines) # 尝试获取阶段配置（从 generation_statuses 或 stage_names） stage_config = None if self.task_id: try: from app.routes.generator.utils import generation_statuses, generation_lock with generation_lock: task_status = generation_statuses.get(self.task_id, {}) stage_names = task_status.get('stage_names', {}) if stage_names and stage in stage_names: # 构建阶段配置字典 stage_config = { 'id': stage, 'name': stage_names[stage] } except Exception as e: self.logger.debug(f"获取阶段配置失败: {e}") # 调用 StageDetector，传递阶段配置 return StageDetector.extract_stage_info(stage, output_text, stage_config=stage_config) def _move_temp_to_output(self, category_id: str, task_id: str) -> Optional[str]: """将临时工作目录移动到 output/ 目录 Args: category_id: 方向ID task_id: 任务ID Returns: 移动后的输出目录路径，如果失败则返回 None """ from pathlib import Path import shutil temp_dir = self.ge10_dir / category_id / f"temp_{task_id}" output_base_dir = self.ge10_dir / category_id / "output" output_base_dir.mkdir(parents=True, exist_ok=True) if not temp_dir.exists(): self.logger.warning(f"临时目录不存在: {temp_dir}") return None # 查找临时目录中的 output/ 子目录（AI 可能在临时目录下创建了 output/） temp_output_dir = temp_dir / "output" if temp_output_dir.exists(): # 查找 output/ 下的题目目录 challenge_dirs = [d for d in temp_output_dir.iterdir() if d.is_dir()] if challenge_dirs: # 找到最新的题目目录 latest_challenge_dir = max(challenge_dirs, key=lambda d: d.stat().st_mtime) challenge_dir_name = latest_challenge_dir.name # 移动到 output/ 目录 target_dir = output_base_dir / challenge_dir_name if target_dir.exists(): # 如果目标已存在，使用时间戳后缀 import datetime timestamp = datetime.datetime.now().strftime('%Y%m%d_%H%M%S') challenge_dir_name = f"{challenge_dir_name}_{timestamp}" target_dir = output_base_dir / challenge_dir_name shutil.move(str(latest_challenge_dir), str(target_dir)) self.logger.info(f"已将题目目录从临时目录移动到: {target_dir}") # 清理临时目录 try: shutil.rmtree(str(temp_dir)) self.logger.info(f"已清理临时目录: {temp_dir}") except Exception as e: self.logger.warning(f"清理临时目录失败: {e}") return str(target_dir) # 如果没有 output/ 子目录，检查临时目录本身是否包含题目文件 # 检查是否有 writeup.md 等关键文件 has_writeup = (temp_dir / "writeup.md").exists() has_docker = (temp_dir / "docker").exists() if has_writeup or has_docker: # 临时目录本身就是题目目录，需要移动到 output/ # 从目录名或时间戳生成题目目录名 import datetime timestamp = datetime.datetime.now().strftime('%Y%m%d_%H%M%S') challenge_dir_name = f"{timestamp}_challenge" # 尝试从 writeup 中提取题目名 if has_writeup: try: writeup_content = (temp_dir / "writeup.md").read_text(encoding='utf-8') import re title_match = re.search(r'^#\s+([^#\n]+?)(?:\s*-\s*[Ww]riteup)?\s*$', writeup_content, re.MULTILINE) if title_match: challenge_name = title_match.group(1).strip() # 清理文件名不合法字符 challenge_name = re.sub(r'[^\w\s-]', '', challenge_name) challenge_name = re.sub(r'\s+', '_', challenge_name) challenge_dir_name = f"{timestamp}_{challenge_name}" except Exception as e: self.logger.warning(f"从 writeup 提取题目名失败: {e}") target_dir = output_base_dir / challenge_dir_name if target_dir.exists(): # 如果目标已存在，添加时间戳后缀 target_dir = output_base_dir / f"{challenge_dir_name}_{int(datetime.datetime.now().timestamp())}" shutil.move(str(temp_dir), str(target_dir)) self.logger.info(f"已将临时目录移动到: {target_dir}") return str(target_dir) # 如果都没有，尝试查找最新的输出目录 self.logger.warning(f"临时目录中未找到题目文件，尝试查找最新输出目录") return self._find_latest_output_dir(category_id) def _find_latest_output_dir(self, category_id: str = 'web') -> Optional[str]: """查找最新生成的输出目录（根据 category_id 查找对应的目录） Args: category_id: 方向ID，默认为 'web' """ from pathlib import Path # 根据 category_id 查找对应的输出目录：ge10/{category_id}/output category_output_dir = self.ge10_dir / category_id / "output" if category_output_dir.exists(): # 直接查找该目录下的子目录（题目目录） subdirs = [d for d in category_output_dir.iterdir() if d.is_dir()] if subdirs: latest_dir = max(subdirs, key=lambda d: d.stat().st_mtime) return str(latest_dir) return None def _merge_directory(self, source: Path, target: Path) -> None: """递归合并目录内容（移动源目录内容到目标目录） Args: source: 源目录路径 target: 目标目录路径 """ import shutil if not source.exists(): return if not target.exists(): target.mkdir(parents=True, exist_ok=True) for item in source.iterdir(): target_item = target / item.name if item.is_dir(): # 如果是目录，递归合并 if target_item.exists(): self._merge_directory(item, target_item) else: # 目标不存在，直接移动 shutil.move(str(item), str(target_item)) else: # 如果是文件 if target_item.exists(): # 目标文件已存在，先删除再移动 if target_item.is_file(): target_item.unlink() else: # 如果是目录，递归删除 shutil.rmtree(target_item) shutil.move(str(item), str(target_item)) # 源目录已空，尝试删除（如果失败也不影响） try: if source.exists() and not any(source.iterdir()): source.rmdir() except OSError: pass def _move_session_to_challenge_dir(self, output_dir: str) -> None: """将会话历史从全局缓存目录移动到题目目录 Args: output_dir: 题目输出目录 """ try: # 源目录：全局缓存目录 source_cache = self.project_root / ".augment_cache" / f"user_{self.user_id}" # 目标目录：题目的 .augment_session/ 目录 target_cache = Path(output_dir) / ".augment_session" if not source_cache.exists(): print(f"⚠️ 源缓存目录不存在: {source_cache}", flush=True) return # 创建目标目录 target_cache.mkdir(parents=True, exist_ok=True) # 移动 sessions 目录 source_sessions = source_cache / "sessions" target_sessions = target_cache / "sessions" if source_sessions.exists(): import shutil if target_sessions.exists(): # 如果目标已存在，合并内容（递归处理） self._merge_directory(source_sessions, target_sessions) else: # 直接移动整个目录 shutil.move(str(source_sessions), str(target_sessions)) print(f"✅ 已将会话历史移动到: {target_cache}", flush=True) # 移动其他缓存目录（如果存在） for subdir in ['binaries', 'checkpoint-documents', 'task-storage']: source_subdir = source_cache / subdir target_subdir = target_cache / subdir if source_subdir.exists(): import shutil if target_subdir.exists(): # 如果目标已存在，合并内容（递归处理） self._merge_directory(source_subdir, target_subdir) else: # 直接移动整个目录 shutil.move(str(source_subdir), str(target_subdir)) except Exception as e: print(f"⚠️ 移动会话历史失败: {e}", flush=True) def _call_augment(self, instruction: str, category_id: str = 'web', task_id: Optional[str] = None) -> str: """调用 Augment CLI 执行指令 Args: instruction: 用户指令 category_id: 方向ID，用于设置工作目录 task_id: 任务ID（可选，用于创建独立工作目录避免多任务冲突） Returns: Augment 的响应文本 """ # 确定工作目录： # 为了让 Augment 的行为和 API 模式保持一致，直接以 ge10/{category_id} 作为工作根目录， # 让 AI 在其中创建 output/YYYYMMDD_HHMMSS_题目名称/ 结构。 workspace_dir = self.ge10_dir / category_id if task_id: # 仅用于日志标记当前任务，目录本身不附加 task_id，方便与 API 模式复用相同的 output 结构 self.logger.info(f"使用共享工作目录（与 API 模式一致）: {workspace_dir}，task_id={task_id}") else: self.logger.warning("未提供 task_id，使用共享工作目录（与 API 模式一致）") workspace_dir.mkdir(parents=True, exist_ok=True) # 不再在工作目录中创建配置文件，避免多任务冲突 # 规则文件通过 -if 参数传递给 Augment CLI，不需要在工作目录中创建 # 临时处理工作目录下的 CLAUDE.md 文件，避免被 Augment 自动加载造成冲突 # Augment 会自动扫描工作目录下的 .md 文件作为规则文件，这可能导致冲突 claude_md_in_workspace = workspace_dir / 'CLAUDE.md' claude_md_backup = None if claude_md_in_workspace.exists(): # 临时重命名，避免被 Augment 自动加载 import time claude_md_backup = workspace_dir / f'CLAUDE.md.backup_{int(time.time())}' try: claude_md_in_workspace.rename(claude_md_backup) print(f"📌 已临时重命名工作目录下的 CLAUDE.md，避免规则文件冲突", flush=True) except Exception as e: print(f"⚠️ 重命名 CLAUDE.md 失败: {e}，可能会造成规则文件冲突", flush=True) # 创建合并的临时文件（编译后的 prompt + 用户指令） # 使用进程ID+时间戳生成唯一文件名，避免多进程冲突 import time unique_suffix = f"{os.getpid()}_{int(time.time() * 1000)}" temp_rules_file = workspace_dir / f".augment_rules_{unique_suffix}.md" try: # 使用编译后的 prompt（如果存在） if hasattr(self, '_compiled_prompt') and self._compiled_prompt: rules_content = self._compiled_prompt # 将规则文件内容和用户指令合并 combined_content = f"{rules_content}\n\n---\n\n{instruction}" temp_rules_file.write_text(combined_content, encoding='utf-8') print(f"📜 已将 Prompt 和用户指令合并到临时文件: {temp_rules_file}", flush=True) else: # 如果没有编译后的 prompt，只使用指令 temp_rules_file.write_text(instruction, encoding='utf-8') print(f"⚠️ 未找到编译后的 Prompt，仅使用用户指令: {temp_rules_file}", flush=True) except Exception as e: print(f"⚠️ 创建合并文件失败: {e}，回退到仅使用指令文件", flush=True) # 如果合并失败，创建仅包含指令的临时文件 temp_rules_file.write_text(instruction, encoding='utf-8') try: # 设置环境变量 env = os.environ.copy() env['PATH'] = f"{self.node_path}:{env.get('PATH', '')}" # 获取用户Token并设置到环境变量 user_cache_dir = None if self.user_id: from app.services.ai.helpers.augment_helper import get_user_augment_token, convert_token_to_augment_format token_str = get_user_augment_token(self.user_id) if token_str: # 转换 Token 格式为 Augment CLI 期望的格式 converted_token = convert_token_to_augment_format(token_str) if converted_token: env['AUGMENT_SESSION_AUTH'] = converted_token print(f"✅ 已设置用户 {self.user_id} 的Augment Token", flush=True) print(f"🔑 Token (前50字符): {converted_token[:50]}...", flush=True) # 为每个用户创建独立的缓存目录 user_cache_dir = self.project_root / ".augment_cache" / f"user_{self.user_id}" user_cache_dir.mkdir(parents=True, exist_ok=True) print(f"📂 用户缓存目录: {user_cache_dir}", flush=True) else: print(f"❌ Token 转换失败，原始 Token: {token_str[:50]}...", flush=True) else: print(f"⚠️ 警告: 用户 {self.user_id} 未设置Token，将使用系统默认Token", flush=True) else: print("⚠️ 警告: 未指定用户ID，将使用系统默认Token", flush=True) # 检查环境变量是否设置成功 if 'AUGMENT_SESSION_AUTH' in env: print(f"✅ 环境变量 AUGMENT_SESSION_AUTH 已设置 (长度: {len(env['AUGMENT_SESSION_AUTH'])})", flush=True) else: print(f"❌ 环境变量 AUGMENT_SESSION_AUTH 未设置!", flush=True) # 构建 Augment CLI 命令 # 使用合并后的临时文件作为 -if 参数（指导文件），工作目录设置为 ge10/{category_id} cmd = [ 'auggie', '-if', str(temp_rules_file.resolve()), # 使用合并后的临时文件（规则+指令） '-p', # 打印输出 '--workspace-root', str(workspace_dir.resolve()) # 工作目录设置为 ge10/{category_id} ] # 如果指定了模型，添加 --model 参数 if self.model and self.model != 'augment': cmd.extend(['--model', self.model]) print(f"🤖 使用模型: {self.model}", flush=True) print(f"📂 工作目录: {workspace_dir.resolve()}", flush=True) # 如果有用户缓存目录，添加到命令参数 if user_cache_dir: cmd.extend(['--augment-cache-dir', str(user_cache_dir)]) import logging logger = logging.getLogger(__name__) logger.info(f"执行 Augment CLI: {' '.join(cmd)}") logger.info(f"工作目录: {workspace_dir}") logger.info(f"合并的规则文件: {temp_rules_file.resolve()}") # 创建日志文件（保存到 log/ 目录） import datetime from pathlib import Path # 确保 logs 目录存在 log_dir = Path("logs") log_dir.mkdir(exist_ok=True) # 日志文件名使用完整的 task_id 以确保唯一性 if task_id: # 使用完整的 task_id 作为文件名（task_id 本身已包含时间戳和 UUID） log_file = log_dir / f"augment_{task_id}.txt" else: # 如果没有 task_id，使用时间戳和进程ID确保唯一性 import uuid unique_id = f"{datetime.datetime.now().strftime('%Y%m%d%H%M%S')}_{os.getpid()}_{uuid.uuid4().hex[:8]}" log_file = log_dir / f"augment_{unique_id}.txt" logger.info(f"日志文件: {log_file} (task_id: {task_id})") # 设置日志文件路径到 generation_statuses（按 task_id 索引） from app.routes.generator.utils import generation_statuses, generation_lock if task_id: with generation_lock: if task_id not in generation_statuses: generation_statuses[task_id] = {} generation_statuses[task_id]["log_file"] = str(log_file) generation_statuses[task_id]["log_position"] = 0 else: # 向后兼容：如果没有 task_id，使用全局 generation_status from app.routes.generator.utils import generation_status generation_status["log_file"] = str(log_file) generation_status["log_position"] = 0 # 记录开始时间 start_time = time.time() # 执行命令（实时输出到控制台） import sys process = subprocess.Popen( cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, # 合并 stderr 到 stdout text=True, env=env, cwd=str(workspace_dir), bufsize=0, # 无缓冲 universal_newlines=True ) # 实时读取并打印输出 output_lines = [] line_count = 0 current_stage = -1 # 初始为 -1，表示还未开始任何阶段 # 打开日志文件 # 使用统一的日志格式器 from app.services.ai.core.log_formatter import UnifiedLogWriter, LogLevel with UnifiedLogWriter(str(log_file), "Augment") as log_writer: try: generation_complete = False # 标记是否检测到完成标记 current_stage = -1 # 初始阶段为 -1（未开始） while True: line = process.stdout.readline() if not line: # 检查进程是否结束 if process.poll() is not None: break continue # 写入原始行到日志（保持原始输出） log_writer.write_raw(line, flush=True) output_lines.append(line) # 控制台输出（根据 verbose 模式过滤） if self._should_print_line(line): print(line, end='', flush=True) sys.stdout.flush() # 检测生成完成标记（优先检测，因为这是最可靠的完成信号） output_text_so_far = ''.join(output_lines) if StageDetector.detect_completion(output_text_so_far): print(f"✅ 检测到生成完成标记 [CTF_GENERATION_COMPLETE]", flush=True) log_writer.write_success("检测到生成完成标记，所有阶段已完成") # 获取总阶段数，标记所有阶段为完成 total_stages = StageDetector.get_total_stages() if total_stages > 0: # 标记所有阶段为完成（仅记录到日志，不再上报进度） for stage_num in range(total_stages): if stage_num > current_stage: output_text = ''.join(output_lines) stage_info = self._extract_stage_info(stage_num, output_lines) if hasattr(self, '_extract_stage_info') else None stage_name = StageDetector.get_stage_name(stage_num) log_writer.write_stage(stage_num, stage_name, "完成") # 标记当前阶段为完成（仅记录到日志，不再上报进度） if current_stage >= 0: output_text = ''.join(output_lines) stage_info = self._extract_stage_info(current_stage, output_lines) if hasattr(self, '_extract_stage_info') else None stage_name = StageDetector.get_stage_name(current_stage) log_writer.write_stage(current_stage, stage_name, "完成") # 设置完成标志，后续不再处理阶段切换 generation_complete = True break # 检测阶段变化 stage_detected = self._detect_stage(line) if stage_detected is not None: print(f"🔍 检测到阶段标记: {stage_detected}, 当前阶段: {current_stage}", flush=True) if stage_detected is not None and stage_detected != current_stage: print(f"🔄 阶段切换: {current_stage} -> {stage_detected}", flush=True) # 确保 StageDetector 已设置阶段名称（可能在多任务环境下被重置） if task_id: from app.routes.generator.utils import generation_statuses, generation_lock with generation_lock: task_status = generation_statuses.get(task_id, {}) stage_names = task_status.get('stage_names', {}) if stage_names: StageDetector.set_stage_names(stage_names) # 使用统一的阶段转换处理逻辑（不再上报进度，仅用于阶段日志/目录检测等） output_text = ''.join(output_lines) current_stage = StageDetector.handle_stage_transition( current_stage=current_stage, new_stage=stage_detected, output_text=output_text, extract_info_func=self._extract_stage_info, logger=self.logger ) # 记录阶段信息（使用统一格式） stage_msg = StageDetector.get_stage_name(current_stage) log_writer.write_stage(current_stage, stage_msg, "开始") # 每 100 行显示一次进度（仅在详细模式下） line_count += 1 if self.verbose and line_count % 100 == 0: progress_msg = f"[已处理 {line_count} 行]" print(progress_msg, flush=True) log_writer.write_info(progress_msg) elif not self.verbose and line_count % 100 == 0: # 非详细模式：只写入日志，不打印 log_writer.write_info(f"[已处理 {line_count} 行]") # 等待进程完成 try: process.wait(timeout=600) # 10 分钟超时 except subprocess.TimeoutExpired: process.kill() raise Exception("Augment CLI 执行超时（超过 10 分钟）") # 如果没有检测到完成标记，将最后一个阶段标记为完成（使用统一方法） if not generation_complete: output_text = ''.join(output_lines) StageDetector.finalize_stage( current_stage=current_stage, output_text=output_text, extract_info_func=self._extract_stage_info, logger=self.logger ) # 记录执行时间（使用统一格式） elapsed_time = time.time() - start_time log_writer.write_success(f"Augment 执行完成，耗时: {elapsed_time:.2f} 秒") self.logger.info(f"Augment 执行时间: {elapsed_time:.2f} 秒") except Exception as e: # 记录错误到日志 log_writer.write_error(f"Augment 执行异常: {str(e)}", e) self.logger.error(f"Augment 执行异常: {str(e)}", exc_info=True) raise # 合并输出 output = ''.join(output_lines) # 检查执行结果 if process.returncode != 0: error_msg = f"Augment CLI 执行失败，返回码: {process.returncode}" self.logger.error(error_msg) # 检查是否有输出目录生成（即使AI超时，文件可能已经生成） # 注意：这里需要根据 category_id 查找对应方向的 output 目录 output_dir_found = False # 从 workspace_dir 中提取 category_id（如果可能） category_id_from_workspace = None try: workspace_parts = Path(workspace_dir).parts if 'ge10' in workspace_parts: ge10_index = workspace_parts.index('ge10') if len(workspace_parts) > ge10_index + 1: category_id_from_workspace = workspace_parts[ge10_index + 1] except: pass # 如果无法从工作目录提取，尝试从输出中检测 for line in output_lines: if 'output/' in line: # 尝试检测输出目录（使用新的目录结构） detected_dir = StageDetector.detect_output_dir(line, str(self.ge10_dir)) if detected_dir and os.path.exists(detected_dir): # 检查是否有关键文件 has_writeup = os.path.exists(os.path.join(detected_dir, 'writeup.md')) has_exp = os.path.exists(os.path.join(detected_dir, 'exp.py')) has_docker = os.path.exists(os.path.join(detected_dir, 'docker')) if has_writeup or has_docker: # 只要有 writeup 或 docker 就认为生成成功 output_dir_found = True dir_name = os.path.basename(detected_dir) print(f"⚠️ 虽然AI超时，但检测到题目文件已生成: {dir_name}") print(f" - writeup.md: {'✅' if has_writeup else '❌'}") print(f" - exp.py: {'✅' if has_exp else '❌'}") print(f" - docker/: {'✅' if has_docker else '❌'}") break if not output_dir_found: # 记录错误到日志 if self.log_callback: self.log_callback(0, "error", error_msg) raise Exception(error_msg) else: # 文件已生成，视为成功（带警告） warning_msg = f"AI执行超时（返回码: {process.returncode}），但题目文件已完整生成" print(f"⚠️ {warning_msg}") # 设置超时恢复标记 self.timeout_recovered = True self.timeout_warning = warning_msg if self.log_callback: self.log_callback(0, "warning", f"{error_msg}（但文件已生成）") # 记录输出到日志 if self.log_callback: self.log_callback(0, "assistant", output) print(f"✅ Augment 执行完成，输出长度: {len(output)} 字符") return output except subprocess.TimeoutExpired: error_msg = "Augment CLI 执行超时（超过 10 分钟）" print(f"❌ {error_msg}") if self.log_callback: self.log_callback(0, "error", error_msg) raise Exception(error_msg) except Exception as e: error_msg = f"Augment CLI 执行异常: {str(e)}" print(f"❌ {error_msg}") if self.log_callback: self.log_callback(0, "error", error_msg) raise finally: # 清理临时规则文件 try: if 'temp_rules_file' in locals() and temp_rules_file.exists(): temp_rules_file.unlink() print(f"✅ 已清理临时规则文件: {temp_rules_file}", flush=True) except Exception as e: print(f"⚠️ 清理临时文件失败: {e}", flush=True) # 恢复工作目录下的 CLAUDE.md 文件（如果之前被重命名了） try: if 'claude_md_backup' in locals() and claude_md_backup and claude_md_backup.exists(): original_name = workspace_dir / 'CLAUDE.md' claude_md_backup.rename(original_name) print(f"✅ 已恢复 CLAUDE.md 文件", flush=True) except Exception as e: print(f"⚠️ 恢复 CLAUDE.md 文件失败: {e}", flush=True) # ============================================================ # 工作空间隔离与会话管理 # ============================================================ def create_isolated_workspace(self, challenge_id: int, output_dir: str) -> Dict[str, str]: """为题目创建独立的工作空间（直接使用题目的output_dir） Args: challenge_id: 题目ID output_dir: 原始输出目录路径(ge10/{category_id}/output/xxx) Returns: 包含工作空间信息的字典: { 'workspace_dir': 工作空间目录路径, 'session_dir': 会话目录路径 } """ if not output_dir or not os.path.exists(output_dir): raise ValueError(f"题目输出目录不存在: {output_dir}") # 直接使用题目的output_dir作为工作空间 workspace_dir = Path(output_dir) # 创建会话目录（如果不存在） session_dir = workspace_dir / ".augment_session" session_dir.mkdir(exist_ok=True) # 不再创建配置文件，避免污染题目目录 # 规则文件通过 -if 参数传递给 Augment CLI，不需要在工作目录中创建 print(f"✅ 使用题目目录作为工作空间: {workspace_dir}") return { 'workspace_dir': str(workspace_dir), 'session_dir': str(session_dir) } def _create_workspace_security_config(self, workspace_dir: Path, category_id: str = 'web'): """在工作空间中创建安全配置文件和规则文件 Args: workspace_dir: 工作空间目录 category_id: 方向ID，用于限制访问范围 """ # 创建 .augmentignore 文件,限制AI访问范围（如果不存在或需要更新） augmentignore = workspace_dir / ".augmentignore" # 计算相对于 ge10 的路径深度，用于限制访问 try: workspace_relative = workspace_dir.relative_to(self.ge10_dir) depth = len(workspace_relative.parts) except ValueError: # 如果无法计算相对路径，使用默认深度 depth = 2 # 构建更严格的限制规则 ignore_rules = """# 安全限制:禁止访问工作空间外的文件 # 严格限制：只能访问当前工作目录及其子目录 # 禁止访问父目录（多层防护），但允许访问共享资源（data/） ../* ../../* ../../../* ../../../../* ../../../../../* # 允许访问父目录的共享资源（同一 category 下的 data/ 目录） !../data/ # 禁止访问其他方向的目录 """ # 添加对其他方向的限制 if category_id: for other_dir in ['web', 'pwn', 'reverse', 'crypto', 'misc']: if other_dir != category_id: ignore_rules += f"../{other_dir}/*\n" ignore_rules += f"../../{other_dir}/*\n" ignore_rules += """# 禁止访问系统敏感目录 /etc/* /root/* /home/* ~/* /usr/* /var/* /bin/* /sbin/* /opt/* # 禁止访问项目根目录外的任何内容 ../../../* ../../../../* # 允许访问当前目录和子目录（output/、data/ 等） !output/ !data/ !prompts/ """ # 始终更新 .augmentignore 文件，确保安全规则最新 augmentignore.write_text(ignore_rules) self.logger.info(f"已创建/更新安全配置文件: {augmentignore}") # 创建 .augment/rules/ 目录并复制 CTF 生成规则（如果不存在） rules_dir = workspace_dir / ".augment" / "rules" rules_dir.mkdir(parents=True, exist_ok=True) # 复制 CTF 生成规则文件 source_rule = self.ge10_dir / ".augment" / "rules" / "ctf-generation-guide.md" target_rule = rules_dir / "ctf-generation-guide.md" if source_rule.exists() and not target_rule.exists(): import shutil shutil.copy2(source_rule, target_rule) self.logger.info(f"已复制 CTF 生成规则到工作空间: {target_rule}") # 创建 README_SECURITY.md 提示AI注意安全边界（如果不存在） readme = workspace_dir / "README_SECURITY.md" if not readme.exists(): readme.write_text("""# 工作空间安全说明 ⚠️ **重要提示** 这是一个CTF题目工作空间。请注意: 1. **只能修改当前目录内的文件** 2. **禁止访问父目录或其他题目的文件** 3. **禁止执行危险的系统命令** 4. **所有修改应该围绕完善当前题目进行** ## 允许的操作 - 修改源代码文件 - 修改Docker配置 - 修改writeup.md - 添加新的测试文件 - 修复构建错误 - 添加提示文件 ## 禁止的操作 - 访问 `../` 父目录 - 删除整个工作空间 - 修改其他题目的文件 - 访问系统敏感目录 请遵守这些规则,专注于完善当前题目。 """) print(f"✅ 已创建工作空间安全配置") def validate_user_instruction(self, instruction: str) -> tuple[bool, Optional[str]]: """验证用户指令的安全性 Args: instruction: 用户输入的指令 Returns: (is_valid, error_message): 是否有效和错误信息 """ # 危险模式列表 FORBIDDEN_PATTERNS = [ (r'\.\./+', "禁止访问父目录"), (r'/etc/', "禁止访问系统配置目录"), (r'/root/', "禁止访问root目录"), (r'/home/(?!.*workspaces/challenge_)', "禁止访问其他用户目录"), (r'rm\s+-rf\s+/', "禁止执行危险的删除命令"), (r'workspaces/challenge_\d+(?!/)', "禁止访问其他题目的工作空间"), (r'sudo\s+', "禁止使用sudo命令"), (r'chmod\s+777', "禁止修改文件权限为777"), # 移除对 /dev/ 重定向的限制，允许正常的重定向操作 # (r'>\s*/dev/', "禁止写入设备文件"), ] for pattern, error_msg in FORBIDDEN_PATTERNS: if re.search(pattern, instruction, re.IGNORECASE): return False, f"安全检查失败: {error_msg}" return True, None def continue_conversation( self, challenge_id: int, workspace_dir: str, user_instruction: str, session_id: Optional[str] = None ) -> Dict[str, Any]: """继续之前的对话来完善题目 Args: challenge_id: 题目ID workspace_dir: 工作空间目录 user_instruction: 用户新的指令 session_id: 会话ID（可选，如果提供则恢复特定会话） Returns: 执行结果字典，包含 response, session_id, workspace_dir """ # 验证指令安全性 is_valid, error_msg = self.validate_user_instruction(user_instruction) if not is_valid: raise ValueError(error_msg) # 验证工作空间路径 workspace_path = Path(workspace_dir) # 如果是相对路径，转换为绝对路径 if not workspace_path.is_absolute(): workspace_path = self.project_root / workspace_path if not workspace_path.exists(): raise ValueError(f"工作空间不存在: {workspace_dir}") # 确保工作空间在允许的范围内（ge10/{category_id}目录） # 不再限制在 ge10/output 目录，而是允许在 ge10/{category_id} 目录下 ge10_dir = self.project_root / "ge10" try: # 使用 resolve() 获取规范化的绝对路径，然后检查是否在 ge10 目录下 workspace_resolved = workspace_path.resolve() ge10_resolved = ge10_dir.resolve() workspace_resolved.relative_to(ge10_resolved) except (ValueError, RuntimeError): raise ValueError(f"工作空间路径不合法，必须在 {ge10_dir} 目录下") print(f"\n{'='*60}") print(f"💬 继续对话 - 题目 #{challenge_id}") print(f"📂 工作空间: {workspace_dir}") print(f"📝 用户指令: {user_instruction}") if session_id: print(f"🔗 会话ID: {session_id}") print(f"{'='*60}\n") # 调用 Augment CLI 继续对话 result = self._call_augment_continue( instruction=user_instruction, workspace_dir=workspace_dir, session_id=session_id ) return { "status": "success", "response": result['response'], "session_id": result['session_id'], "workspace_dir": workspace_dir } def continue_conversation_stream( self, challenge_id: int, workspace_dir: str, user_instruction: str, session_id: Optional[str] = None ): """继续对话 - 流式输出版本 Args: challenge_id: 题目ID workspace_dir: 工作空间目录 user_instruction: 用户新的指令 session_id: 会话ID（可选） Yields: 输出的文本块 """ # 验证指令安全性 is_valid, error_msg = self.validate_user_instruction(user_instruction) if not is_valid: raise ValueError(error_msg) # 验证工作空间路径 workspace_path = Path(workspace_dir) if not workspace_path.is_absolute(): workspace_path = self.project_root / workspace_path if not workspace_path.exists(): raise ValueError(f"工作空间不存在: {workspace_path}") # 验证路径安全性 try: workspace_path.resolve().relative_to(self.project_root.resolve()) except ValueError: raise ValueError(f"工作空间路径不合法，必须在 {self.project_root} 目录下") print(f"\n{'='*60}") print(f"💬 继续对话（流式） - 题目 #{challenge_id}") print(f"📂 工作空间: {workspace_dir}") print(f"📝 用户指令: {user_instruction}") if session_id: print(f"🔗 会话ID: {session_id}") print(f"{'='*60}\n") # 调用流式版本 for chunk in self._call_augment_continue_stream( instruction=user_instruction, workspace_dir=workspace_dir, session_id=session_id ): yield chunk def _clean_stream_chunk(self, chunk: str) -> str: """对流式输出的文本块进行基本清理，保留换行符""" import re # 仅移除ANSI颜色代码 ansi_escape = re.compile(r'\x1B(?:[@-Z\\-_]|\[[0-?]*[ -/]*[@-~])') cleaned = ansi_escape.sub('', chunk) # 移除会话恢复信息 cleaned = re.sub(r'📂\s*Resumed session from[^🤖\n]*🤖\s*', '', cleaned) cleaned = re.sub(r'📂\s*Resumed session from[^\n]*\n', '', cleaned) cleaned = re.sub(r'Resumed session from[^\n]*\n', '', cleaned) # 移除开头的emoji cleaned = re.sub(r'^[🤖🧑👤]\s*', '', cleaned) return cleaned def _clean_augment_output(self, output: str) -> str: """清理Augment CLI的输出，移除ANSI颜色代码和不必要的调试信息 Args: output: 原始输出 Returns: 清理后的输出 """ import re # 移除ANSI颜色代码 ansi_escape = re.compile(r'\x1B(?:[@-Z\\-_]|\[[0-?]*[ -/]*[@-~])') cleaned = ansi_escape.sub('', output) # 移除会话恢复信息（多种格式） # 格式1: "📂 Resumed session from ... 🤖" (同一行) cleaned = re.sub(r'📂\s*Resumed session from[^🤖\n]*🤖\s*', '', cleaned) # 格式2: "📂 Resumed session from ...\n" (单独一行) cleaned = re.sub(r'📂\s*Resumed session from[^\n]*\n', '', cleaned) cleaned = re.sub(r'Resumed session from[^\n]*\n', '', cleaned) # 移除行首的单独emoji（🤖, 🧑, 👤） cleaned = re.sub(r'^[🤖🧑👤]\s*\n', '', cleaned, flags=re.MULTILINE) cleaned = re.sub(r'\n[🤖🧑👤]\s*\n', '\n', cleaned) # 移除开头的emoji cleaned = re.sub(r'^[🤖🧑👤]\s*', '', cleaned) # 移除多余的空行（保留单个空行） cleaned = re.sub(r'\n{3,}', '\n\n', cleaned) # 移除行尾空格 lines = cleaned.split('\n') cleaned_lines = [line.rstrip() for line in lines] cleaned = '\n'.join(cleaned_lines) # 移除开头和结尾的空白 cleaned = cleaned.strip() return cleaned def _extract_session_id_from_output(self, output: str) -> Optional[str]: """从Augment CLI输出中提取会话ID Args: output: Augment CLI的输出 Returns: 会话ID，如果未找到则返回None """ import re # 查找类似 "Resumed session from 2025/11/16 11:56:53: (52 exchanges)" 的行 # 或者从日志中提取 session ID match = re.search(r'session[:\s]+([a-f0-9-]{36})', output, re.IGNORECASE) if match: return match.group(1) return None def _save_conversation_to_session( self, session_cache_dir: Path, session_id: str, user_message: str, assistant_message: str ) -> None: """手动保存对话到会话文件 因为使用 -p (print mode) 时，Augment CLI 不会自动保存会话历史， 所以我们需要手动更新会话文件。 Args: session_cache_dir: 会话缓存目录 session_id: 会话ID user_message: 用户消息 assistant_message: AI回复 """ import json from datetime import datetime, timezone session_file = session_cache_dir / "sessions" / f"{session_id}.json" if not session_file.exists(): import logging logger = logging.getLogger(__name__) logger.warning(f"会话文件不存在，无法保存: {session_file}") return try: # 读取现有会话 with open(session_file, 'r', encoding='utf-8') as f: session_data = json.load(f) # 添加新的对话记录 new_exchange = { "exchange": { "request_message": user_message, "response_text": assistant_message, "request_id": f"manual-{datetime.now().timestamp()}", "request_nodes": [ { "id": 1, "type": 0, "text_node": { "content": user_message } } ] } } session_data['chatHistory'].append(new_exchange) # 更新修改时间 session_data['modified'] = datetime.now(timezone.utc).isoformat().replace('+00:00', 'Z') # 保存回文件 with open(session_file, 'w', encoding='utf-8') as f: json.dump(session_data, f, ensure_ascii=False, indent=2) print(f"✅ 已保存对话到会话文件（共 {len(session_data['chatHistory'])} 条）", flush=True) except Exception as e: print(f"⚠️ 保存会话失败: {e}", flush=True) def _call_augment_continue( self, instruction: str, workspace_dir: str, session_id: Optional[str] = None ) -> Dict[str, str]: """调用 Augment CLI 继续对话 Args: instruction: 用户指令 workspace_dir: 工作空间目录 session_id: 会话ID（可选，如果提供则恢复特定会话） Returns: 包含响应文本和会话ID的字典 """ # 创建临时指令文件 with tempfile.NamedTemporaryFile( mode='w', suffix='.md', delete=False, encoding='utf-8' ) as f: f.write(instruction) instruction_file = f.name try: # 转换工作空间路径为绝对路径 workspace_path = Path(workspace_dir) if not workspace_path.is_absolute(): workspace_path = self.project_root / workspace_path workspace_abs = str(workspace_path.resolve()) # 设置环境变量 env = os.environ.copy() env['PATH'] = f"{self.node_path}:{env.get('PATH', '')}" # 获取用户Token if self.user_id: from app.services.ai.helpers.augment_helper import get_user_augment_token, convert_token_to_augment_format token_str = get_user_augment_token(self.user_id) if token_str: converted_token = convert_token_to_augment_format(token_str) env['AUGMENT_SESSION_AUTH'] = converted_token print(f"✅ 已设置用户 {self.user_id} 的Augment Token", flush=True) # 使用题目目录下的 .augment_session/ 作为会话存储目录 session_cache_dir = workspace_path / ".augment_session" session_cache_dir.mkdir(parents=True, exist_ok=True) print(f"📂 会话目录: {session_cache_dir}", flush=True) # 构建命令 - 使用 -p 参数执行，然后手动保存会话历史 # 注意：-p (print mode) 是 one-shot 模式，不会自动保存会话历史 # 我们会在执行完后手动更新会话文件 if session_id: # 如果有会话ID，使用 --resume 恢复特定会话 cmd = [ 'auggie', '--resume', session_id, # 恢复特定会话 '-if', instruction_file, '-p', # 打印模式（one-shot） '--workspace-root', workspace_abs # 限制工作目录（使用绝对路径） ] print(f"📌 恢复会话: {session_id}", flush=True) else: # 否则继续最近的会话 cmd = [ 'auggie', '--continue', # 继续最近的对话 '-if', instruction_file, '-p', # 打印模式（one-shot） '--workspace-root', workspace_abs # 限制工作目录（使用绝对路径） ] print("📌 继续最近的会话", flush=True) # 添加会话缓存目录（使用题目目录下的 .augment_session/） cmd.extend(['--augment-cache-dir', str(session_cache_dir)]) print(f"🚀 执行 Augment CLI (继续对话): {' '.join(cmd)}", flush=True) print(f"📂 工作空间: {workspace_abs}", flush=True) print("=" * 60, flush=True) # 创建日志文件（包含 task_id 以确保唯一性） import datetime log_dir = Path("logs") log_dir.mkdir(exist_ok=True) # 日志文件名使用完整的 task_id 以确保唯一性 if self.task_id: # 使用完整的 task_id 作为文件名 log_file = log_dir / f"augment_continue_{self.task_id}.txt" else: # 如果没有 task_id，使用时间戳和进程ID确保唯一性 import uuid unique_id = f"{datetime.datetime.now().strftime('%Y%m%d%H%M%S')}_{os.getpid()}_{uuid.uuid4().hex[:8]}" log_file = log_dir / f"augment_continue_{unique_id}.txt" # 执行命令 start_time = time.time() process = subprocess.Popen( cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True, env=env, cwd=workspace_abs, # 在工作空间中执行（使用绝对路径） bufsize=0, universal_newlines=True ) # 实时读取输出 output_lines = [] with open(log_file, 'w', encoding='utf-8') as log_f: for line in process.stdout: output_lines.append(line) # 打印到控制台 if self._should_print_line(line): print(line, end='', flush=True) # 写入日志 log_f.write(line) log_f.flush() # 等待进程完成 process.wait(timeout=600) # 10分钟超时 elapsed_time = time.time() - start_time print("\n" + "=" * 60) print(f"⏱️ 执行时间: {elapsed_time:.2f} 秒") print("=" * 60) output = ''.join(output_lines) if process.returncode != 0: raise Exception(f"Augment CLI 执行失败，返回码: {process.returncode}") # 清理输出内容，移除ANSI颜色代码和调试信息 cleaned_output = self._clean_augment_output(output) # 尝试从输出中提取会话ID（如果没有提供） extracted_session_id = session_id if not extracted_session_id: extracted_session_id = self._extract_session_id_from_output(output) # 如果还是没有会话ID，尝试从会话目录中查找最新的会话文件 if not extracted_session_id and session_cache_dir: sessions_dir = session_cache_dir / "sessions" if sessions_dir.exists(): import glob session_files = glob.glob(str(sessions_dir / "*.json")) if session_files: # 获取最新的会话文件 latest_session = max(session_files, key=os.path.getmtime) extracted_session_id = Path(latest_session).stem print(f"📌 检测到会话ID: {extracted_session_id}", flush=True) # 手动保存会话历史（因为 -p 模式不会自动保存） if extracted_session_id and session_cache_dir: self._save_conversation_to_session( session_cache_dir=session_cache_dir, session_id=extracted_session_id, user_message=instruction, assistant_message=cleaned_output ) return { 'response': cleaned_output, 'session_id': extracted_session_id } except subprocess.TimeoutExpired: process.kill() raise Exception("Augment CLI 执行超时（超过 10 分钟）") except Exception as e: raise Exception(f"Augment CLI 执行异常: {str(e)}") finally: # 清理临时文件 try: os.unlink(instruction_file) except: pass def _call_augment_continue_stream( self, instruction: str, workspace_dir: str, session_id: Optional[str] = None ): """调用 Augment CLI 继续对话 - 流式输出版本 Args: instruction: 用户指令 workspace_dir: 工作空间目录 session_id: 会话ID（可选） Yields: 输出的文本块 """ # 创建临时指令文件 with tempfile.NamedTemporaryFile( mode='w', suffix='.md', delete=False, encoding='utf-8' ) as f: f.write(instruction) instruction_file = f.name try: # 转换工作空间路径为绝对路径 workspace_path = Path(workspace_dir) if not workspace_path.is_absolute(): workspace_path = self.project_root / workspace_path workspace_abs = str(workspace_path.resolve()) # 设置环境变量 env = os.environ.copy() env['PATH'] = f"{self.node_path}:{env.get('PATH', '')}" # 获取用户Token if self.user_id: from app.services.ai.helpers.augment_helper import get_user_augment_token, convert_token_to_augment_format token_str = get_user_augment_token(self.user_id) if token_str: converted_token = convert_token_to_augment_format(token_str) env['AUGMENT_SESSION_AUTH'] = converted_token # 使用题目目录下的 .augment_session/ 作为会话存储目录 session_cache_dir = workspace_path / ".augment_session" session_cache_dir.mkdir(parents=True, exist_ok=True) # 构建命令 if session_id: cmd = [ 'auggie', '--resume', session_id, '-if', instruction_file, '-p', # 打印模式 '--workspace-root', workspace_abs ] else: cmd = [ 'auggie', '--continue', '-if', instruction_file, '-p', '--workspace-root', workspace_abs ] cmd.extend(['--augment-cache-dir', str(session_cache_dir)]) # 执行命令 process = subprocess.Popen( cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True, env=env, cwd=workspace_abs, bufsize=1, # 行缓冲 universal_newlines=True ) # 实时读取并yield输出 full_output = [] for line in process.stdout: full_output.append(line) # 对每个块进行基本清理并yield，保留换行符 cleaned_chunk = self._clean_stream_chunk(line) yield cleaned_chunk # 等待进程完成 process.wait(timeout=600) if process.returncode != 0: raise Exception(f"Augment CLI 执行失败，返回码: {process.returncode}") # 保存会话历史 output = ''.join(full_output) cleaned_output = self._clean_augment_output(output) # 提取会话ID extracted_session_id = session_id if not extracted_session_id: extracted_session_id = self._extract_session_id_from_output(output) if not extracted_session_id and session_cache_dir: sessions_dir = session_cache_dir / "sessions" if sessions_dir.exists(): import glob session_files = glob.glob(str(sessions_dir / "*.json")) if session_files: latest_session = max(session_files, key=os.path.getmtime) extracted_session_id = Path(latest_session).stem # 手动保存会话历史 if extracted_session_id and session_cache_dir: self._save_conversation_to_session( session_cache_dir=session_cache_dir, session_id=extracted_session_id, user_message=instruction, assistant_message=cleaned_output ) except subprocess.TimeoutExpired: process.kill() raise Exception("Augment CLI 执行超时（超过 10 分钟）") except Exception as e: raise Exception(f"Augment CLI 执行异常: {str(e)}") finally: # 清理临时文件 try: os.unlink(instruction_file) except: pass