""" CTF 方向配置管理 API 提供方向配置、知识库、Prompt 模板的管理接口。 """ import json import os from flask import Blueprint, request, jsonify, current_app from functools import wraps from app.models.database.models import ( db, User, CategoryConfig, CategoryAdmin, Role ) from app.services.category import CategoryService from app.services.prompt import PromptCompiler from app.services.auth import AuthService from app.services.ai.tools import ToolExecutor, SandboxConfig def get_current_user(): """获取当前用户""" return AuthService.get_current_user() category_api_bp = Blueprint('category_api', __name__, url_prefix='/api/admin/categories') def admin_required(f): """管理员权限装饰器""" @wraps(f) def decorated_function(*args, **kwargs): user = get_current_user() if not user or user.role != Role.ADMIN: return jsonify({'success': False, 'error': '需要管理员权限'}), 403 return f(*args, **kwargs) return decorated_function def category_permission_required(min_role='viewer'): """ 方向权限装饰器 - 支持三级权限 权限级别: - viewer: 只读权限，可查看配置 - editor: 编辑权限，可修改配置 - owner: 管理权限，可管理权限和删除 """ def decorator(f): @wraps(f) def decorated_function(*args, **kwargs): user = get_current_user() if not user: return jsonify({'success': False, 'error': '请先登录'}), 401 # 系统管理员拥有所有权限 if user.role == Role.ADMIN: return f(*args, **kwargs) category_id = kwargs.get('category_id') or request.view_args.get('category_id') if not category_id: return f(*args, **kwargs) # 获取用户在该方向的角色 admin_record = CategoryAdmin.query.filter_by( category_id=category_id, user_id=user.id ).first() if not admin_record: return jsonify({'success': False, 'error': '没有该方向的访问权限'}), 403 # 权限级别映射 role_levels = {'viewer': 1, 'editor': 2, 'owner': 3} user_level = role_levels.get(admin_record.role, 0) required_level = role_levels.get(min_role, 1) if user_level < required_level: role_names = {'viewer': '查看', 'editor': '编辑', 'owner': '管理'} return jsonify({ 'success': False, 'error': f'需要{role_names.get(min_role, min_role)}权限' }), 403 return f(*args, **kwargs) return decorated_function return decorator def category_admin_required(f): """方向管理员权限装饰器（编辑级别）- 兼容旧代码""" return category_permission_required('editor')(f) # ==================== 方向配置 API ==================== @category_api_bp.route('', methods=['GET']) def list_categories(): """获取所有方向列表""" user = get_current_user() if not user: return jsonify({'success': False, 'error': '请先登录'}), 401 # 管理员可以看到所有方向，普通用户只能看到自己管理的 if user.role == Role.ADMIN: categories = CategoryService.get_all_categories() else: categories = CategoryAdmin.get_user_categories(user.id) return jsonify({ 'success': True, 'categories': [c.to_dict(include_config=False) for c in categories] }) @category_api_bp.route('/', methods=['GET']) @category_admin_required def get_category(category_id): """获取方向详情""" category = CategoryService.get_category(category_id) if not category: return jsonify({'success': False, 'error': '方向不存在'}), 404 return jsonify({ 'success': True, 'category': category.to_dict(include_config=True) }) @category_api_bp.route('', methods=['POST']) @admin_required def create_category(): """创建新方向""" data = request.get_json() if not data: return jsonify({'success': False, 'error': '无效的请求数据'}), 400 category_id = data.get('id') if not category_id: return jsonify({'success': False, 'error': '方向 ID 不能为空'}), 400 # 检查是否已存在 existing = CategoryConfig.query.get(category_id) if existing: return jsonify({'success': False, 'error': '方向 ID 已存在'}), 400 try: category = CategoryService.create_category(data) return jsonify({ 'success': True, 'category': category.to_dict(), 'message': '方向创建成功' }) except Exception as e: return jsonify({'success': False, 'error': str(e)}), 500 @category_api_bp.route('/', methods=['PUT']) @category_admin_required def update_category(category_id): """更新方向配置""" data = request.get_json() if not data: return jsonify({'success': False, 'error': '无效的请求数据'}), 400 try: category = CategoryService.update_category(category_id, data) if not category: return jsonify({'success': False, 'error': '方向不存在'}), 404 return jsonify({ 'success': True, 'category': category.to_dict(), 'message': '方向配置更新成功' }) except Exception as e: return jsonify({'success': False, 'error': str(e)}), 500 @category_api_bp.route('/', methods=['DELETE']) @admin_required def delete_category(category_id): """删除方向""" try: success = CategoryService.delete_category(category_id) if not success: return jsonify({'success': False, 'error': '方向不存在'}), 404 return jsonify({ 'success': True, 'message': '方向删除成功' }) except Exception as e: return jsonify({'success': False, 'error': str(e)}), 500 @category_api_bp.route('//toggle', methods=['POST']) @admin_required def toggle_category(category_id): """启用/禁用方向""" category = CategoryConfig.query.get(category_id) if not category: return jsonify({'success': False, 'error': '方向不存在'}), 404 category.enabled = not category.enabled db.session.commit() return jsonify({ 'success': True, 'enabled': category.enabled, 'message': f"方向已{'启用' if category.enabled else '禁用'}" }) # ==================== 表单配置 API ==================== @category_api_bp.route('//form-fields', methods=['GET']) @category_admin_required def get_form_fields(category_id): """获取表单字段配置""" category = CategoryConfig.query.get(category_id) if not category: return jsonify({'success': False, 'error': '方向不存在'}), 404 return jsonify({ 'success': True, 'form_fields': category.get_form_fields(), 'form_layout': category.get_form_layout() }) @category_api_bp.route('//form-fields', methods=['PUT']) @category_admin_required def update_form_fields(category_id): """更新表单字段配置""" data = request.get_json() if not data: return jsonify({'success': False, 'error': '无效的请求数据'}), 400 category = CategoryConfig.query.get(category_id) if not category: return jsonify({'success': False, 'error': '方向不存在'}), 404 if 'form_fields' in data: category.set_form_fields(data['form_fields']) if 'form_layout' in data: category.set_form_layout(data['form_layout']) db.session.commit() return jsonify({ 'success': True, 'message': '表单配置更新成功' }) # ==================== 阶段配置 API ==================== @category_api_bp.route('//stages', methods=['GET']) @category_admin_required def get_stages(category_id): """获取阶段配置""" category = CategoryConfig.query.get(category_id) if not category: return jsonify({'success': False, 'error': '方向不存在'}), 404 return jsonify({ 'success': True, 'stages': category.get_stages() }) @category_api_bp.route('//stages', methods=['PUT']) @category_admin_required def update_stages(category_id): """更新阶段配置""" data = request.get_json() if not data or 'stages' not in data: return jsonify({'success': False, 'error': '无效的请求数据'}), 400 category = CategoryConfig.query.get(category_id) if not category: return jsonify({'success': False, 'error': '方向不存在'}), 404 category.set_stages(data['stages']) # 自动编译并存储 Prompt 模板 from app.services.prompt.compiler_service import PromptCompilerService try: PromptCompilerService.compile_and_save_prompts(category) except Exception as e: # 编译失败不影响保存配置，只记录日志 import logging logging.warning(f'自动编译 Prompt 失败: {str(e)}') db.session.commit() return jsonify({ 'success': True, 'message': '阶段配置更新成功' }) # ==================== Prompt 模板 API ==================== @category_api_bp.route('/import', methods=['POST']) @admin_required def import_category(): """导入方向配置（JSON）- 支持单个配置""" data = request.get_json() if not data: return jsonify({'success': False, 'error': '无效的请求数据'}), 400 config = data.get('config') mode = data.get('mode', 'create') # create, overwrite, merge if not config: return jsonify({'success': False, 'error': '缺少配置数据'}), 400 # 验证必要字段 if not config.get('id'): return jsonify({'success': False, 'error': '配置必须包含 id 字段'}), 400 if not config.get('name'): return jsonify({'success': False, 'error': '配置必须包含 name 字段'}), 400 category_id = config['id'] existing = CategoryConfig.query.get(category_id) try: if existing: if mode == 'create': return jsonify({'success': False, 'error': f'方向 {category_id} 已存在，请选择覆盖或合并模式'}), 400 elif mode == 'overwrite': # 覆盖模式：删除旧的，创建新的 db.session.delete(existing) db.session.commit() category = CategoryService.create_category(config) else: # merge # 合并模式：只更新提供的字段 category = CategoryService.update_category(category_id, config) else: # 新建 category = CategoryService.create_category(config) return jsonify({ 'success': True, 'category': category.to_dict(), 'message': f'方向 {category.name} 导入成功' }) except Exception as e: db.session.rollback() return jsonify({'success': False, 'error': f'导入失败: {str(e)}'}), 500 @category_api_bp.route('/import-batch', methods=['POST']) @admin_required def import_categories_batch(): """批量导入方向配置（JSON数组）""" data = request.get_json() if not data: return jsonify({'success': False, 'error': '无效的请求数据'}), 400 configs = data.get('configs') mode = data.get('mode', 'create') # create, overwrite, merge if not configs: return jsonify({'success': False, 'error': '缺少配置数据'}), 400 # 确保是数组 if not isinstance(configs, list): return jsonify({'success': False, 'error': '配置数据必须是数组格式'}), 400 if len(configs) == 0: return jsonify({'success': False, 'error': '配置数组不能为空'}), 400 results = { 'success': [], 'failed': [], 'total': len(configs) } # 批量导入 for idx, config in enumerate(configs): try: # 验证必要字段 if not config.get('id'): results['failed'].append({ 'index': idx + 1, 'id': None, 'error': '配置必须包含 id 字段' }) continue if not config.get('name'): results['failed'].append({ 'index': idx + 1, 'id': config.get('id'), 'error': '配置必须包含 name 字段' }) continue category_id = config['id'] existing = CategoryConfig.query.get(category_id) if existing: if mode == 'create': results['failed'].append({ 'index': idx + 1, 'id': category_id, 'name': existing.name, 'error': f'方向已存在，请选择覆盖或合并模式' }) continue elif mode == 'overwrite': # 覆盖模式：删除旧的，创建新的 db.session.delete(existing) db.session.commit() category = CategoryService.create_category(config) else: # merge # 合并模式：只更新提供的字段 category = CategoryService.update_category(category_id, config) else: # 新建 category = CategoryService.create_category(config) results['success'].append({ 'index': idx + 1, 'id': category_id, 'name': category.name, 'message': f'方向 {category.name} 导入成功' }) except Exception as e: db.session.rollback() results['failed'].append({ 'index': idx + 1, 'id': config.get('id', '未知'), 'error': f'导入失败: {str(e)}' }) # 返回批量导入结果 success_count = len(results['success']) failed_count = len(results['failed']) return jsonify({ 'success': failed_count == 0, # 全部成功才返回 success: true 'results': results, 'message': f'批量导入完成：成功 {success_count} 个，失败 {failed_count} 个' }) @category_api_bp.route('/export', methods=['GET']) @admin_required def export_all_categories(): """导出所有方向配置（JSON数组）""" try: categories = CategoryConfig.query.order_by(CategoryConfig.sort_order).all() configs = [] for category in categories: config = { 'id': category.id, 'name': category.name, 'icon': category.icon, 'description': category.description, 'enabled': category.enabled, 'sort_order': category.sort_order, 'form_fields': category.get_form_fields(), 'form_layout': category.get_form_layout(), 'output_config': category.get_output_config(), 'ui_config': category.get_ui_config(), 'advanced_config': category.get_advanced_config(), 'prompt_template_path': category.prompt_template_path, 'knowledge_base_path': category.knowledge_base_path, 'knowledge_db_path': category.knowledge_db_path, 'choice_script_path': category.choice_script_path, 'output_dir': category.output_dir } # 优先导出新格式（difficulties），如果不存在则导出旧格式 difficulties = category.get_difficulties() if difficulties and any(difficulties.values()): # 检查是否有非空配置 config['difficulties'] = difficulties else: # 降级到旧格式 config['stages'] = category.get_stages() config['difficulty_rules'] = category.get_difficulty_rules() configs.append(config) return jsonify({ 'success': True, 'configs': configs, 'count': len(configs) }) except Exception as e: return jsonify({'success': False, 'error': str(e)}), 500 @category_api_bp.route('//export', methods=['GET']) @category_permission_required('viewer') def export_category(category_id): """导出方向配置（直接下载 JSON 文件）""" from flask import Response import json from datetime import datetime category = CategoryConfig.query.get(category_id) if not category: return jsonify({'success': False, 'error': '方向不存在'}), 404 # 构建完整配置 config = { 'id': category.id, 'name': category.name, 'icon': category.icon, 'description': category.description, 'enabled': category.enabled, 'sort_order': category.sort_order, 'form_fields': category.get_form_fields(), 'form_layout': category.get_form_layout(), 'output_config': category.get_output_config(), 'ui_config': category.get_ui_config(), 'advanced_config': category.get_advanced_config(), 'prompt_template_path': category.prompt_template_path, 'knowledge_base_path': category.knowledge_base_path, 'knowledge_db_path': category.knowledge_db_path, 'choice_script_path': category.choice_script_path, 'output_dir': category.output_dir } # 优先导出新格式（difficulties），如果不存在则导出旧格式 difficulties = category.get_difficulties() if difficulties and any(difficulties.values()): # 检查是否有非空配置 config['difficulties'] = difficulties else: # 降级到旧格式 config['stages'] = category.get_stages() config['difficulty_rules'] = category.get_difficulty_rules() # 转换为 JSON 字符串（格式化，缩进2个空格） json_str = json.dumps(config, ensure_ascii=False, indent=2) # 生成文件名（使用方向ID和时间戳） timestamp = datetime.now().strftime('%Y%m%d_%H%M%S') filename = f"{category_id}_category_config_{timestamp}.json" # 创建响应，直接返回 JSON 文件下载 response = Response( json_str, mimetype='application/json', headers={ 'Content-Disposition': f'attachment; filename="{filename}"', 'Content-Type': 'application/json; charset=utf-8' } ) return response # 已废弃：旧的 PromptTemplate CRUD API 已移除，现在使用 stage-extensions 和 compile-prompt # @category_api_bp.route('//prompts', methods=['GET', 'POST']) # @category_api_bp.route('//prompts/', methods=['PUT', 'DELETE']) # @category_api_bp.route('//prompts//history', methods=['GET']) # @category_api_bp.route('//prompts//history/', methods=['GET']) # @category_api_bp.route('//prompts//restore/', methods=['POST']) @category_api_bp.route('//prompts/compile', methods=['POST']) @category_admin_required def compile_prompt(category_id): """编译 Prompt 模板""" data = request.get_json() or {} try: compiler = PromptCompiler() # 如果提供了模板内容，使用提供的内容 if 'template_content' in data: compiled = compiler.compile_from_template( data['template_content'], category_id, data.get('user_input') ) else: compiled = compiler.compile(category_id, data.get('user_input')) return jsonify({ 'success': True, 'compiled_prompt': compiled, 'stats': { 'total_lines': len(compiled.split('\n')), 'total_chars': len(compiled) } }) except Exception as e: return jsonify({'success': False, 'error': str(e)}), 500 @category_api_bp.route('//prompts/validate', methods=['POST']) @category_admin_required def validate_prompt(category_id): """验证 Prompt 模板与配置的一致性""" data = request.get_json() if not data or 'template_content' not in data: return jsonify({'success': False, 'error': '请提供模板内容'}), 400 category = CategoryConfig.query.get(category_id) if not category: return jsonify({'success': False, 'error': '方向不存在'}), 404 try: compiler = PromptCompiler() result = compiler.validate_template( data['template_content'], category.to_dict(include_config=True) ) return jsonify({ 'success': True, 'validation': result }) except Exception as e: return jsonify({'success': False, 'error': str(e)}), 500 # ==================== 知识库 API ==================== # 注意：知识库知识点获取API已移至 pages/wizard.py 中的 /api/categories//knowledge # ==================== 方向管理员 API ==================== # 注意：方向管理员管理已通过用户权限管理API实现，见 admin/views.py 中的 /api/user//permissions # ==================== 初始化 API ==================== @category_api_bp.route('/init', methods=['POST']) @admin_required def init_categories(): """初始化默认方向配置""" try: CategoryService.init_default_categories() return jsonify({ 'success': True, 'message': '默认方向配置初始化成功' }) except Exception as e: return jsonify({'success': False, 'error': str(e)}), 500 # ==================== 阶段 Prompt API ==================== @category_api_bp.route('//stage-prompts', methods=['GET']) @category_admin_required def get_stage_prompts(category_id): """获取阶段 Prompt 配置（包含系统固定部分和用户扩展）""" from app.services.prompt.generator import StagePromptGenerator category = CategoryConfig.query.get(category_id) if not category: return jsonify({'success': False, 'error': '方向不存在'}), 404 stages = category.get_stages() or [] advanced = category.get_advanced_config() or {} # 为每个阶段生成系统 Prompt result_stages = StagePromptGenerator.generate_all_stages(category_id, stages) return jsonify({ 'success': True, 'stages': result_stages, 'global_rules': advanced.get('global_rules', '') }) @category_api_bp.route('//stage-extensions', methods=['POST']) @category_admin_required def save_stage_extensions(category_id): """保存阶段扩展内容和全局规则""" category = CategoryConfig.query.get(category_id) if not category: return jsonify({'success': False, 'error': '方向不存在'}), 404 data = request.get_json() extensions = data.get('extensions', {}) global_rules = data.get('global_rules', '') try: stages = category.get_stages() or [] # 更新每个阶段的 user_extension for stage in stages: stage_id = stage.get('id', '') if stage_id in extensions: stage['user_extension'] = extensions[stage_id] # 保存到数据库 category.stages = json.dumps(stages) # 保存全局规则到 advanced_config advanced = category.get_advanced_config() or {} advanced['global_rules'] = global_rules category.advanced_config = json.dumps(advanced) # 自动编译并存储 Prompt 模板 from app.services.prompt.compiler_service import PromptCompilerService try: PromptCompilerService.compile_and_save_prompts(category) except Exception as e: # 编译失败不影响保存配置，只记录日志 import logging logging.warning(f'自动编译 Prompt 失败: {str(e)}') db.session.commit() # 不再自动生成 Prompt 文件，所有配置存储在数据库中 return jsonify({ 'success': True, 'message': 'Prompt 配置保存成功' }) except Exception as e: db.session.rollback() return jsonify({'success': False, 'error': str(e)}), 500 @category_api_bp.route('//compile-prompt', methods=['GET', 'POST']) @category_admin_required def compile_full_prompt(category_id): """编译完整的 Prompt（预览用）""" from app.services.prompt.generator import StagePromptGenerator category = CategoryConfig.query.get(category_id) if not category: return jsonify({'success': False, 'error': '方向不存在'}), 404 try: data = request.get_json() or {} if request.method == 'POST' else {} use_current_edits = data.get('use_current_edits', False) # 如果使用当前编辑的内容 if use_current_edits: current_system_prompt = data.get('system_prompt', '') current_stages = data.get('stages', []) # 使用当前编辑的 stages 数据 stages = current_stages system_prompt = current_system_prompt else: # 从数据库读取 stages = category.get_stages() or [] advanced = category.get_advanced_config() or {} system_prompt = advanced.get('system_prompt', '') advanced = category.get_advanced_config() or {} # 生成带系统 Prompt 的阶段列表 # 如果使用当前编辑的内容，需要手动构建 full_stages if use_current_edits: # 检查是否使用新格式 is_new_format = data.get('is_new_format', False) full_stages = [] for stage_data in stages: stage_info = { 'id': stage_data.get('id', ''), 'name': stage_data.get('name', ''), 'description': stage_data.get('description', ''), 'output': stage_data.get('output_format', ''), 'skip_forbidden': stage_data.get('skip_forbidden', False), 'system_prompt': stage_data.get('system_prompt', ''), 'user_extension': stage_data.get('user_extension', '') } if is_new_format: # 新格式：prompt 直接存储在 stage.prompt 字段 prompt_content = stage_data.get('prompt', '') if prompt_content: # 将 prompt 内容添加到 user_extension if stage_info['user_extension']: stage_info['user_extension'] += '\n\n' + prompt_content else: stage_info['user_extension'] = prompt_content else: # 旧格式：从 prompts 字典中提取 prompts = stage_data.get('prompts', {}) if prompts and isinstance(prompts, dict): # 优先使用中文键名查找 prompt = None for diff in ['中等', '简单', '入门', '困难', 'medium', 'easy', 'beginner', 'hard']: if diff in prompts: prompt = prompts[diff] break if prompt: prompt_content = prompt.get('content', '') if prompt_content: if stage_info['user_extension']: stage_info['user_extension'] += '\n\n' + prompt_content else: stage_info['user_extension'] = prompt_content full_stages.append(stage_info) else: full_stages = StagePromptGenerator.generate_all_stages(category_id, stages) # 构建预览上下文 difficulty_rules = category.get_difficulty_rules() or [] output_config = category.get_output_config() or {} # 生成难度表格 difficulty_table = '' for rule in difficulty_rules: difficulty_table += f"| {rule.get('name', '')} | {rule.get('max_count', 1)} 个 | {rule.get('writeup_count', 5)} 篇 | {rule.get('depth_range', [1, 5])} |\n" context = { 'category': category.name, 'category_id': category_id, 'difficulty': '中等', 'language': 'PHP', 'knowledge_points': 'SQL注入, XSS', 'scene': '博客系统', 'writeup_count': 8, 'max_knowledge': 3, 'depth_range': '[4.0, 7.0]', 'diff_rate': 0.4, 'difficulty_table': difficulty_table, 'available_languages': 'PHP, Python, Node.js, Go, Java', 'port_range': output_config.get('port_range', [42500, 42600]), 'flag_format': output_config.get('flag_format', 'DASCTF{...}'), 'directory_pattern': output_config.get('directory_pattern', '{timestamp}_{name}') } # 检查是否有知识库获取阶段，如果有则执行脚本获取知识库内容 knowledge_stage = None for stage in full_stages: if stage.get('id') == 'knowledge_acquisition': knowledge_stage = stage break # 如果存在知识库获取阶段，将管理员配置的脚本运行说明添加到 Prompt 中 if knowledge_stage: script_instruction = knowledge_stage.get('knowledge_script_instruction', '').strip() # 只有管理员填写了脚本运行说明时才添加到 Prompt 中 if script_instruction: # 将管理员填写的脚本运行说明添加到阶段的 user_extension if knowledge_stage.get('user_extension'): knowledge_stage['user_extension'] = script_instruction + "\n\n" + knowledge_stage['user_extension'] else: knowledge_stage['user_extension'] = script_instruction # 编译完整 Prompt（使用当前编辑的系统提示） prompt = StagePromptGenerator.compile_full_prompt( category_id, full_stages, context, system_prompt if system_prompt else None ) # 不再保存 Prompt 到文件，所有配置存储在数据库中 return jsonify({ 'success': True, 'prompt': prompt }) except Exception as e: import traceback traceback.print_exc() return jsonify({'success': False, 'error': str(e)}), 500 # 已废弃：Prompt 文件管理 API 已移除，所有配置存储在数据库中 # @category_api_bp.route('//prompts/files', methods=['GET']) # ==================== 知识库文件管理 API ==================== # ==================== 知识库文件管理 API（统一data文件夹）==================== def get_data_dir(category_id): """获取data文件夹路径""" current_file = os.path.abspath(__file__) ctf_dir = os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(current_file)))) base_path = os.path.join(ctf_dir, 'ge10') return os.path.join(base_path, category_id, 'data') def normalize_path(path): """规范化路径，防止路径遍历攻击""" if not path: return '' # 移除 .. 和 . 等危险路径 normalized = os.path.normpath(path) # 确保路径不包含 .. if '..' in normalized or normalized.startswith('/'): return None return normalized @category_api_bp.route('//knowledge/files/tree', methods=['GET']) @category_admin_required def get_file_tree(category_id): """获取文件树结构""" try: data_dir = get_data_dir(category_id) if not os.path.exists(data_dir): os.makedirs(data_dir, exist_ok=True) return jsonify({ 'success': True, 'tree': [] }) def build_tree(path, relative_path=''): """递归构建文件树""" items = [] try: for item in sorted(os.listdir(path)): item_path = os.path.join(path, item) item_rel_path = os.path.join(relative_path, item) if relative_path else item if os.path.isdir(item_path): children = build_tree(item_path, item_rel_path) items.append({ 'name': item, 'type': 'directory', 'path': item_rel_path, 'children': children }) else: stat = os.stat(item_path) items.append({ 'name': item, 'type': 'file', 'path': item_rel_path, 'size': stat.st_size, 'modified': stat.st_mtime }) except PermissionError: pass return sorted(items, key=lambda x: (x['type'] == 'file', x['name'].lower())) tree = build_tree(data_dir) return jsonify({ 'success': True, 'tree': tree }) except Exception as e: import traceback current_app.logger.error(f"获取文件树失败: {str(e)}\n{traceback.format_exc()}") return jsonify({'success': False, 'error': str(e)}), 500 @category_api_bp.route('//knowledge/files/upload', methods=['POST']) @category_admin_required def upload_knowledge_file(category_id): """上传文件到data文件夹，支持压缩包自动解压""" from werkzeug.utils import secure_filename import zipfile import tarfile import shutil if 'file' not in request.files: return jsonify({'success': False, 'error': '请选择文件'}), 400 file = request.files['file'] relative_path = request.form.get('path', '') # 相对路径（可选，用于子目录） extract = request.form.get('extract', 'false').lower() == 'true' # 是否解压 if file.filename == '': return jsonify({'success': False, 'error': '请选择文件'}), 400 try: data_dir = get_data_dir(category_id) # 确定保存目录 if relative_path: normalized = normalize_path(relative_path) if not normalized: return jsonify({'success': False, 'error': '非法路径'}), 400 save_dir = os.path.join(data_dir, normalized) else: save_dir = data_dir os.makedirs(save_dir, exist_ok=True) safe_filename = secure_filename(file.filename) file_path = os.path.join(save_dir, safe_filename) # 如果是压缩包且需要解压 if extract: # 先保存临时文件 temp_path = file_path + '.tmp' file.save(temp_path) extracted_count = 0 error_files = [] try: # 判断压缩包类型并解压 if safe_filename.lower().endswith('.zip'): with zipfile.ZipFile(temp_path, 'r') as zip_ref: # 获取所有文件列表 file_list = zip_ref.namelist() for member in file_list: # 安全检查：防止路径遍历攻击 if '..' in member or member.startswith('/'): error_files.append(member) continue # 解压文件 target_path = os.path.join(save_dir, member) # 确保目录存在 if member.endswith('/'): os.makedirs(target_path, exist_ok=True) else: os.makedirs(os.path.dirname(target_path), exist_ok=True) with zip_ref.open(member) as source, open(target_path, 'wb') as target: target.write(source.read()) extracted_count += 1 elif safe_filename.lower().endswith(('.tar', '.tar.gz', '.tgz')): mode = 'r:gz' if safe_filename.lower().endswith(('.tar.gz', '.tgz')) else 'r' with tarfile.open(temp_path, mode) as tar_ref: for member in tar_ref.getmembers(): # 安全检查 if '..' in member.name or member.name.startswith('/'): error_files.append(member.name) continue target_path = os.path.join(save_dir, member.name) # 确保在data目录内 if not os.path.abspath(target_path).startswith(os.path.abspath(data_dir)): error_files.append(member.name) continue tar_ref.extract(member, save_dir) if member.isfile(): extracted_count += 1 else: # 不支持的解压格式 os.remove(temp_path) return jsonify({'success': False, 'error': '不支持的压缩包格式，仅支持 .zip, .tar, .tar.gz'}), 400 # 删除临时文件 os.remove(temp_path) message = f'压缩包解压成功，共解压 {extracted_count} 个文件' if error_files: message += f'，{len(error_files)} 个文件因安全原因跳过' return jsonify({ 'success': True, 'extracted_count': extracted_count, 'error_files': error_files, 'message': message }) except Exception as e: # 清理临时文件 if os.path.exists(temp_path): os.remove(temp_path) raise e else: # 普通文件上传 # 检查文件是否已存在 if os.path.exists(file_path): return jsonify({'success': False, 'error': f'文件 {safe_filename} 已存在'}), 400 file.save(file_path) stat = os.stat(file_path) file_rel_path = os.path.relpath(file_path, data_dir) return jsonify({ 'success': True, 'file': { 'name': safe_filename, 'path': file_rel_path, 'size': stat.st_size, 'modified': stat.st_mtime }, 'message': f'文件 {safe_filename} 上传成功' }) except Exception as e: import traceback current_app.logger.error(f"上传文件失败: {str(e)}\n{traceback.format_exc()}") return jsonify({'success': False, 'error': str(e)}), 500 @category_api_bp.route('//knowledge/files/download', methods=['GET']) @category_admin_required def download_knowledge_file(category_id): """下载文件""" from flask import send_file file_path = request.args.get('path') # 相对路径 if not file_path: return jsonify({'success': False, 'error': '缺少路径参数'}), 400 try: data_dir = get_data_dir(category_id) # 规范化路径 normalized = normalize_path(file_path) if not normalized: return jsonify({'success': False, 'error': '非法路径'}), 400 full_path = os.path.join(data_dir, normalized) # 安全检查：确保文件在data目录内 if not os.path.abspath(full_path).startswith(os.path.abspath(data_dir)): return jsonify({'success': False, 'error': '非法路径'}), 403 if not os.path.exists(full_path): return jsonify({'success': False, 'error': '文件不存在'}), 404 if not os.path.isfile(full_path): return jsonify({'success': False, 'error': '不是文件'}), 400 filename = os.path.basename(full_path) return send_file(full_path, as_attachment=True, download_name=filename) except Exception as e: import traceback current_app.logger.error(f"下载文件失败: {str(e)}\n{traceback.format_exc()}") return jsonify({'success': False, 'error': str(e)}), 500 @category_api_bp.route('//knowledge/files/delete', methods=['DELETE']) @category_admin_required def delete_knowledge_file(category_id): """删除文件或文件夹""" data = request.get_json() file_path = data.get('path') # 相对路径 is_directory = data.get('is_directory', False) if not file_path: return jsonify({'success': False, 'error': '缺少路径参数'}), 400 try: data_dir = get_data_dir(category_id) # 规范化路径 normalized = normalize_path(file_path) if not normalized: return jsonify({'success': False, 'error': '非法路径'}), 400 full_path = os.path.join(data_dir, normalized) # 安全检查：确保路径在data目录内 if not os.path.abspath(full_path).startswith(os.path.abspath(data_dir)): return jsonify({'success': False, 'error': '非法路径'}), 403 if not os.path.exists(full_path): return jsonify({'success': False, 'error': '文件或文件夹不存在'}), 404 # 删除文件或文件夹 if is_directory: import shutil shutil.rmtree(full_path) message = f'文件夹 {os.path.basename(full_path)} 删除成功' else: os.remove(full_path) message = f'文件 {os.path.basename(full_path)} 删除成功' return jsonify({ 'success': True, 'message': message }) except Exception as e: import traceback current_app.logger.error(f"删除失败: {str(e)}\n{traceback.format_exc()}") return jsonify({'success': False, 'error': str(e)}), 500 @category_api_bp.route('//knowledge/files/create', methods=['POST']) @category_admin_required def create_file_or_folder(category_id): """创建文件或文件夹""" data = request.get_json() name = data.get('name') path = data.get('path', '') # 父目录路径（相对路径） is_directory = data.get('is_directory', False) if not name: return jsonify({'success': False, 'error': '缺少名称参数'}), 400 try: data_dir = get_data_dir(category_id) # 确定父目录 if path: normalized = normalize_path(path) if not normalized: return jsonify({'success': False, 'error': '非法路径'}), 400 parent_dir = os.path.join(data_dir, normalized) else: parent_dir = data_dir # 安全检查 if not os.path.abspath(parent_dir).startswith(os.path.abspath(data_dir)): return jsonify({'success': False, 'error': '非法路径'}), 403 os.makedirs(parent_dir, exist_ok=True) # 创建文件或文件夹 new_path = os.path.join(parent_dir, name) if os.path.exists(new_path): return jsonify({'success': False, 'error': f'{"文件夹" if is_directory else "文件"} {name} 已存在'}), 400 if is_directory: os.makedirs(new_path, exist_ok=True) message = f'文件夹 {name} 创建成功' else: # 创建空文件 with open(new_path, 'w', encoding='utf-8') as f: f.write('') message = f'文件 {name} 创建成功' return jsonify({ 'success': True, 'message': message }) except Exception as e: import traceback current_app.logger.error(f"创建失败: {str(e)}\n{traceback.format_exc()}") return jsonify({'success': False, 'error': str(e)}), 500 @category_api_bp.route('//knowledge/files/read', methods=['GET']) @category_admin_required def read_file_content(category_id): """读取文件内容""" file_path = request.args.get('path') if not file_path: return jsonify({'success': False, 'error': '缺少路径参数'}), 400 try: data_dir = get_data_dir(category_id) normalized = normalize_path(file_path) if not normalized: return jsonify({'success': False, 'error': '非法路径'}), 400 full_path = os.path.join(data_dir, normalized) if not os.path.abspath(full_path).startswith(os.path.abspath(data_dir)): return jsonify({'success': False, 'error': '非法路径'}), 403 if not os.path.exists(full_path): return jsonify({'success': False, 'error': '文件不存在'}), 404 if not os.path.isfile(full_path): return jsonify({'success': False, 'error': '不是文件'}), 400 with open(full_path, 'r', encoding='utf-8') as f: content = f.read() return jsonify({ 'success': True, 'content': content }) except Exception as e: import traceback current_app.logger.error(f"读取文件失败: {str(e)}\n{traceback.format_exc()}") return jsonify({'success': False, 'error': str(e)}), 500 @category_api_bp.route('//knowledge/files/write', methods=['POST']) @category_admin_required def write_file_content(category_id): """写入文件内容""" data = request.get_json() file_path = data.get('path') content = data.get('content', '') if not file_path: return jsonify({'success': False, 'error': '缺少路径参数'}), 400 try: data_dir = get_data_dir(category_id) normalized = normalize_path(file_path) if not normalized: return jsonify({'success': False, 'error': '非法路径'}), 400 full_path = os.path.join(data_dir, normalized) if not os.path.abspath(full_path).startswith(os.path.abspath(data_dir)): return jsonify({'success': False, 'error': '非法路径'}), 403 # 确保父目录存在 os.makedirs(os.path.dirname(full_path), exist_ok=True) with open(full_path, 'w', encoding='utf-8') as f: f.write(content) return jsonify({ 'success': True, 'message': '文件保存成功' }) except Exception as e: import traceback current_app.logger.error(f"写入文件失败: {str(e)}\n{traceback.format_exc()}") return jsonify({'success': False, 'error': str(e)}), 500 @category_api_bp.route('//knowledge/shell', methods=['POST']) @category_permission_required('owner') def execute_knowledge_shell(category_id): """在知识库目录下执行命令（管理员/方向所有者）- 流式输出版本""" from flask import Response, stream_with_context import subprocess import threading import queue payload = request.get_json() or {} command = (payload.get('command') or '').strip() cwd = (payload.get('cwd') or '').strip() timeout = payload.get('timeout') if not command: return jsonify({'success': False, 'error': '命令不能为空'}), 400 try: data_dir = get_data_dir(category_id) os.makedirs(data_dir, exist_ok=True) working_dir = data_dir if cwd: normalized = normalize_path(cwd) if not normalized: return jsonify({'success': False, 'error': '非法工作目录'}), 400 working_dir = os.path.join(data_dir, normalized) # 安全检查：工作目录必须在 data 目录内 if not os.path.abspath(working_dir).startswith(os.path.abspath(data_dir)): return jsonify({'success': False, 'error': '工作目录不在允许范围内'}), 400 os.makedirs(working_dir, exist_ok=True) # 构建沙箱 sandbox = SandboxConfig.create_for_knowledge_base( category_id=category_id, policy=SandboxConfig.load_policy() ) # 安全检查命令 allowed, reason = sandbox.is_command_allowed(command, working_dir=working_dir) if not allowed: return jsonify({'success': False, 'error': f'命令被拒绝: {reason}'}), 400 # 检查工作目录 if working_dir and not sandbox.is_path_allowed(working_dir): return jsonify({'success': False, 'error': '工作目录不在允许范围内'}), 400 # 获取超时时间 actual_timeout = sandbox.get_timeout(command) if timeout: actual_timeout = min(timeout, actual_timeout) def generate(): """生成流式输出 - 真正的实时输出""" try: # 设置环境变量 env = os.environ.copy() env['PAGER'] = 'cat' # 禁用分页 env['PYTHONUNBUFFERED'] = '1' # Python 无缓冲输出 # 使用 Popen 以便实时读取输出 process = subprocess.Popen( command, shell=True, cwd=working_dir, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, # 合并 stderr 到 stdout text=True, bufsize=0, # 无缓冲，立即输出 env=env ) # 发送开始标记 yield f"data: {json.dumps({'type': 'start', 'command': command})}\n\n" # 实时读取并立即发送输出 import time start_time = time.time() while True: # 检查超时 if actual_timeout and (time.time() - start_time) > actual_timeout: process.kill() process.wait() yield f"data: {json.dumps({'type': 'error', 'message': f'命令执行超时 ({actual_timeout}秒)'})}\n\n" break # 检查进程是否结束 if process.poll() is not None: # 读取剩余输出 remaining = process.stdout.read() if remaining: yield f"data: {json.dumps({'type': 'output', 'data': remaining})}\n\n" # 发送完成标记 yield f"data: {json.dumps({'type': 'done', 'return_code': process.returncode})}\n\n" break # 实时读取一行（非阻塞） line = process.stdout.readline() if line: # 立即发送这一行 yield f"data: {json.dumps({'type': 'output', 'data': line})}\n\n" else: # 没有输出，短暂休眠避免CPU占用过高 time.sleep(0.01) except Exception as e: import traceback current_app.logger.error(f"流式执行命令失败: {str(e)}\n{traceback.format_exc()}") yield f"data: {json.dumps({'type': 'error', 'message': f'执行错误: {str(e)}'})}\n\n" return Response( stream_with_context(generate()), mimetype='text/event-stream', headers={ 'Cache-Control': 'no-cache', 'X-Accel-Buffering': 'no' # 禁用 nginx 缓冲 } ) except Exception as e: import traceback current_app.logger.error(f"知识库命令执行失败: {str(e)}\n{traceback.format_exc()}") return jsonify({'success': False, 'error': str(e)}), 500