import requests import argparse import os import asyncio import websockets import json import base64 import logging import sys import time # 配置日志 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s') logger = logging.getLogger(__name__) class WebSocketFileUploader: def __init__(self, server_url, api_key, space_id): self.server_url = server_url self.api_key = api_key self.space_id = space_id async def connect_and_upload(self, upload_dir): """连接到 WebSocket 服务器并上传文件""" # --- 修复点: 统一将 URL 转换为小写，以处理类似 "HTTPS" 这样的异常情况 --- normalized_server_url = self.server_url.lower() # 使用转换后的 URL 进行协议替换 ws_url = normalized_server_url.replace('https://', 'wss://').replace('http://', 'ws://').rstrip('/') ws_url += f'/ws/upload/{self.space_id}?token={self.api_key}' logger.info(f"正在连接到 WebSocket: {ws_url}") try: # 增加连接超时以更快地失败 async with websockets.connect(ws_url, open_timeout=10, ping_interval=20, ping_timeout=60) as websocket: logger.info("WebSocket 连接成功") # 等待服务器的认证成功响应 response = await websocket.recv() auth_response = json.loads(response) if auth_response.get("type") == "auth_success": logger.info("认证成功，开始上传文件...") return await self._upload_files(websocket, upload_dir) else: logger.error(f"认证失败: {auth_response.get('message', '未知错误')}") logger.info("回退到 HTTP API 上传...") return self._fallback_http_upload(upload_dir) except Exception as e: logger.error(f"WebSocket 连接失败: {e}") logger.info("回退到 HTTP API 上传...") return self._fallback_http_upload(upload_dir) async def _upload_files(self, websocket, upload_dir): if not os.path.exists(upload_dir): logger.error(f"目录不存在: {upload_dir}") return False all_files = [] for root, dirs, files in os.walk(upload_dir): for file in files: file_path = os.path.join(root, file) if os.path.isfile(file_path): all_files.append(file_path) if not all_files: logger.info("📁 目录中没有找到任何文件") return True logger.info(f"📁 找到 {len(all_files)} 个文件，开始上传...") success_count = 0 failed_count = 0 for file_path in all_files: try: if await self._upload_single_file(websocket, file_path): success_count += 1 logger.info(f"✅ 文件 '{os.path.basename(file_path)}' 上传成功!") else: failed_count += 1 logger.error(f"❌ 文件 '{os.path.basename(file_path)}' 上传失败") await asyncio.sleep(0.5) except Exception as e: logger.error(f"上传文件 {file_path} 时发生异常: {e}") failed_count += 1 logger.info(f"📊 上传完成! 成功: {success_count}, 失败: {failed_count}") return success_count > 0 async def _upload_single_file(self, websocket, file_path): try: filename = os.path.basename(file_path) logger.info(f"正在上传文件: {filename}") with open(file_path, 'rb') as f: file_content = f.read() file_b64 = base64.b64encode(file_content).decode('utf-8') upload_msg = { "type": "file_upload", "filename": filename, "content": file_b64, "space_id": self.space_id } await websocket.send(json.dumps(upload_msg)) response = await websocket.recv() result = json.loads(response) if result.get("type") == "upload_success": return True else: logger.error(f"上传失败: {result.get('message', '未知错误')}") return False except Exception as e: logger.error(f"上传文件 {file_path} 时出错: {e}") return False def _fallback_http_upload(self, upload_dir): logger.info("使用 HTTP API 回退方案上传文件...") if not os.path.exists(upload_dir): logger.error(f"目录不存在: {upload_dir}") return False upload_url = f"{self.server_url.rstrip('/')}/api/remote_upload" headers = {'X-API-Key': self.api_key} data = {'space_id': self.space_id} all_files = [] for root, dirs, files in os.walk(upload_dir): for file in files: file_path = os.path.join(root, file) if os.path.isfile(file_path): all_files.append(file_path) if not all_files: logger.info("📁 目录中没有找到任何文件") return True logger.info(f"📁 找到 {len(all_files)} 个文件，开始HTTP上传...") success_count = 0 failed_count = 0 for file_path in all_files: try: filename = os.path.basename(file_path) logger.info(f"正在上传文件: {filename}") with open(file_path, 'rb') as f: files = {'file': (filename, f)} response = requests.post(upload_url, headers=headers, data=data, files=files, timeout=60) if response.status_code == 200: logger.info(f"✅ 文件 '{filename}' 上传成功!") success_count += 1 else: logger.error(f"❌ 上传失败: HTTP {response.status_code} - {response.text}") failed_count += 1 time.sleep(0.5) except Exception as e: logger.error(f"上传文件 {file_path} 时发生异常: {e}") failed_count += 1 logger.info(f"📊 上传完成! 成功: {success_count}, 失败: {failed_count}") return success_count > 0 async def upload_directory_websocket(upload_dir, server_url, api_key, space_id): logger.info(f"🔍 开始扫描目录: {upload_dir}") logger.info(f"📡 服务器地址: {server_url}") logger.info(f"🔑 Space ID: {space_id}") logger.info("-" * 50) uploader = WebSocketFileUploader(server_url, api_key, space_id) success = await uploader.connect_and_upload(upload_dir) if success: logger.info("🎉 文件已成功上传到您的网盘!") else: logger.error("❌ 文件上传失败") return success if __name__ == "__main__": parser = argparse.ArgumentParser(description="WebSocket 文件上传器 - 扫描并上传指定文件夹中的所有文件") parser.add_argument("api_key", help="您的 API 密钥") parser.add_argument("space_id", help="Space ID") parser.add_argument("--server", default="http://127.0.0.1:5001", help="服务器的 URL 地址 (例如: http://127.0.0.1:5001)") parser.add_argument("--upload-dir", default="output", help="要上传的目录 (默认: output)") args = parser.parse_args() asyncio.run(upload_directory_websocket(args.upload_dir, args.server, args.api_key, args.space_id))