diff --git a/app.py b/app.py index a6b038cb9d62e6e4891985525d27ad564441bc27..fa9454a01467b0c08649318c1622ba4f27eff348 100644 --- a/app.py +++ b/app.py @@ -23,7 +23,30 @@ ANNOTATIONS_FOLDER = os.path.join(STATIC_FOLDER, 'annotations') app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER app.config['STATIC_FOLDER'] = STATIC_FOLDER app.config['ANNOTATIONS_FOLDER'] = ANNOTATIONS_FOLDER -app.config['MAX_CONTENT_LENGTH'] = None # 不限制上传大小 +# 限制单次上传大小(默认 4GB,可用 XIABIE_MAX_UPLOAD_MB 调),避免超大文件打满磁盘(DoS) +app.config['MAX_CONTENT_LENGTH'] = int(os.environ.get('XIABIE_MAX_UPLOAD_MB', '4096')) * 1024 * 1024 + + +def _safe_name(filename): + """清洗上传/操作文件名,防路径穿越:只取 basename(去掉任何目录成分),保留中文。 + `../../etc/passwd` -> `passwd`(join 后仍落在目标目录内);空/'.'/'..' 返回 None。""" + name = os.path.basename((filename or '').replace('\\', '/').strip()) + if not name or name in ('.', '..'): + return None + return name + + +def _confined_path(base, name): + """把 name 限制在 base 目录内,防 ../ 穿越。返回绝对路径;越界/非法返回 None。""" + safe = _safe_name(name) + if safe is None: + return None + base_abs = os.path.abspath(base) + p = os.path.abspath(os.path.join(base_abs, safe)) + if p == base_abs or p.startswith(base_abs + os.sep): + return p + return None + # 创建必要的目录 os.makedirs(UPLOAD_FOLDER, exist_ok=True) @@ -138,8 +161,11 @@ def delete_images(): for image_name in image_names: try: - # 删除图片文件 - image_path = os.path.join(app.config['UPLOAD_FOLDER'], image_name) + # 删除图片文件(限制在 uploads 目录内,防 ../ 删任意文件) + image_path = _confined_path(app.config['UPLOAD_FOLDER'], image_name) + if image_path is None: + errors.append(f"非法图片名: {image_name}") + continue if os.path.exists(image_path): os.remove(image_path) deleted_count += 1 @@ -188,10 +214,11 @@ def upload_folder(): uploaded_files = [] for file in files: - if file.filename != '': - filepath = os.path.join(app.config['UPLOAD_FOLDER'], file.filename or '') + safe = _safe_name(file.filename) + if safe: + filepath = os.path.join(app.config['UPLOAD_FOLDER'], safe) file.save(filepath) - uploaded_files.append(file.filename or '') + uploaded_files.append(safe) return jsonify({'message': 'Files uploaded successfully', 'files': uploaded_files}) @@ -235,7 +262,10 @@ def upload_labelme_dataset(): # 处理图像文件 for image_filename, image_file in image_files.items(): - # 保存图像文件 + # 保存图像文件(清洗文件名防路径穿越) + image_filename = _safe_name(image_filename) + if not image_filename: + continue image_path = os.path.join(app.config['UPLOAD_FOLDER'], image_filename) image_file.save(image_path) uploaded_files.append(image_filename) @@ -380,7 +410,9 @@ def upload_video(): def extract_frames(video_path, frame_interval, original_name=None): """从视频中抽帧并保存为图片""" - + # 清洗原视频名(会拼进抽帧文件名),防路径穿越;去掉扩展名 + original_name = os.path.splitext(_safe_name(original_name) or 'video')[0] or 'video' + # 打开视频文件 cap = cv2.VideoCapture(video_path) if not cap.isOpened(): @@ -1024,7 +1056,12 @@ def download_models(): # 获取模型列表和安装路径 models_str = request.args.get('models', '') - models = models_str.split(',') if models_str else [] + # 模型名会被拼进子进程的 python -c 源码,必须白名单校验,防代码注入。 + # 只允许 yolo 预训练权重的命名(字母数字/点/连字符/下划线),滤掉危险字符。 + import re as _re + models = [m for m in (models_str.split(',') if models_str else []) + if _re.fullmatch(r'[A-Za-z0-9._-]+', m.strip())] + models = [m.strip() for m in models] install_path = request.args.get('install_path', 'plugins/yolo11') # 确保安装路径是相对于项目根目录的 @@ -1142,13 +1179,14 @@ def upload_model(): uploaded_files = [] files = request.files.getlist('files[]') for file in files: - if file.filename: - ext = os.path.splitext(file.filename)[1].lower() + safe = _safe_name(file.filename) + if safe: + ext = os.path.splitext(safe)[1].lower() if ext in ('.pt', '.onnx'): - # 保存文件到models目录 - file_path = os.path.join(models_dir, file.filename) + # 保存文件到models目录(清洗文件名防路径穿越) + file_path = os.path.join(models_dir, safe) file.save(file_path) - uploaded_files.append(file.filename) + uploaded_files.append(safe) return jsonify({'success': True, 'uploaded_files': uploaded_files}) @@ -1176,10 +1214,12 @@ def delete_model(): if not model_name: return jsonify({'success': False, 'error': '模型名称不能为空'}) - # 构建模型文件路径 + # 构建模型文件路径(限制在 models 目录内,防 ../ 删任意文件) models_dir = os.path.join(install_path, 'models') - model_path = os.path.join(models_dir, model_name) - + model_path = _confined_path(models_dir, model_name) + if model_path is None: + return jsonify({'success': False, 'error': '非法模型名'}) + # 检查模型文件是否存在 if not os.path.exists(model_path): return jsonify({'success': False, 'error': '模型文件不存在'}) @@ -1260,47 +1300,24 @@ def export_dataset(): # 如果是'all'则不进行过滤，使用所有图片 # 分割数据集 - np.random.shuffle(images) - + # 固定随机种子,划分可复现(原来无种子,每次导出划分都变) + seed = int(os.environ.get('XIABIE_SPLIT_SEED', '42')) + np.random.RandomState(seed).shuffle(images) + total_images = len(images) - - # 彻底重写数据集分割逻辑，确保严格按照比例分割 - # 0比例的数据集绝对为空，多余的数据直接扔掉 - train_images = [] - val_images = [] - test_images = [] - - # 只处理比例大于0的数据集 - if train_ratio > 0: - # 计算训练集数量 - train_count = int(total_images * train_ratio) - # 只分配计算出的数量的图片 - train_images = images[:train_count] - - # 验证集只在train_ratio > 0时才处理，否则从0开始 - val_start = len(train_images) if train_ratio > 0 else 0 - if val_ratio > 0: - # 计算验证集数量 - val_count = int(total_images * val_ratio) - # 只分配计算出的数量的图片 - val_images = images[val_start:val_start + val_count] - - # 测试集只在train_ratio > 0或val_ratio > 0时才处理，否则从0开始 - test_start = (len(train_images) + len(val_images)) if (train_ratio > 0 or val_ratio > 0) else 0 - if test_ratio > 0: - # 计算测试集数量 - test_count = int(total_images * test_ratio) - # 只分配计算出的数量的图片 - test_images = images[test_start:test_start + test_count] - - # 确保0比例的数据集绝对为空 - if train_ratio == 0: - train_images = [] - if val_ratio == 0: - val_images = [] + + # 划分:train/val 按比例 int 取,余下(含 int 截断的零头)全归 test,**绝不丢图** + # (原逻辑用 int 截断后把零头直接扔掉,会静默丢训练数据)。 + n_train = min(int(total_images * train_ratio), total_images) + n_val = min(int(total_images * val_ratio), total_images - n_train) + train_images = list(images[:n_train]) + val_images = list(images[n_train:n_train + n_val]) + test_images = list(images[n_train + n_val:]) # 余数(零头)兜到 test,不丢图 + # test_ratio=0 时余数补回 train(而非丢弃)。约定三比例之和为 1。 if test_ratio == 0: + train_images += test_images test_images = [] - + # 处理每个分割的数据集 splits = [ ('train', train_images), @@ -1362,13 +1379,13 @@ names: {selected_classes} # 只导出选中的类别 if ann['class'] in selected_classes: # 转换为YOLO格式: class_id center_x center_y width height (归一化) - # 修改这里，使用全局类别列表中的索引而不是选中类别列表中的索引 - class_id = None - # 从全局类别列表中查找类别ID - for i, cls in enumerate(classes): - if cls['name'] == ann['class']: - class_id = i - break + # 关键修复:class_id 必须用「导出的 selected_classes」里的下标, + # 才能和 data.yaml 的 names(也是 selected_classes)对齐。 + # 原来用全局 classes 下标,导出子集时 id 与 names 错位 → 训练标签全错。 + try: + class_id = selected_classes.index(ann['class']) + except ValueError: + class_id = None # 如果在全局类别中找到了该类别，则写入标签文件 if class_id is not None: @@ -1454,7 +1471,14 @@ names: {selected_classes} if __name__ == '__main__': - app.run(debug=True, host='0.0.0.0', port=5000) + # 安全默认:只监听本机 127.0.0.1、关 debug(Werkzeug 调试器=RCE 入口)。 + # 确需对外/调试时显式用环境变量打开,并自负风险(本服务无鉴权)。 + host = os.environ.get('XIABIE_HOST', '127.0.0.1') + port = int(os.environ.get('XIABIE_PORT', '5000')) + debug = os.environ.get('XIABIE_DEBUG', '0') == '1' + if host != '127.0.0.1': + print('[WARN] 监听在 %s:%d —— 本服务无鉴权,暴露到网络等于把机器交出去!' % (host, port)) + app.run(debug=debug, host=host, port=port) def process_content_data(content_data, annotations):