websocket_uploader.py

import asyncio
import websockets
import json
import logging
import sys
import base64
import os
import time

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

async def upload_files_via_websocket(space_id, user_token, website_url):
    """通过WebSocket上传output目录中的所有文件"""
    
    output_dir = "output"
    if not os.path.exists(output_dir):
        logger.error(f"输出目录不存在: {output_dir}")
        return False
    
    # 获取所有文件
    files_to_upload = []
    for root, dirs, files in os.walk(output_dir):
        for file in files:
            file_path = os.path.join(root, file)
            if os.path.isfile(file_path):
                files_to_upload.append(file_path)
    
    if not files_to_upload:
        logger.info("输出目录中没有找到任何文件")
        return True
    
    logger.info(f"找到 {len(files_to_upload)} 个文件需要上传")
    
    # 构建WebSocket URL
    ws_url = website_url.replace('http://', 'ws://').replace('https://', 'wss://') + '/socket.io/?EIO=4&transport=websocket'
    
    try:
        # 连接到WebSocket服务器
        logger.info(f"正在连接到 WebSocket 服务器: {ws_url}")
        
        # 使用socket.io客户端协议
        import socketio
        
        sio = socketio.AsyncClient()
        
        @sio.event
        async def connect():
            logger.info("WebSocket 连接成功")
            # 注册客户端
            await sio.emit('register_client', {
                'space_id': space_id,
                'user_token': user_token
            })
        
        @sio.event
        async def registered(data):
            logger.info(f"客户端注册成功: {data}")
            # 开始上传文件
            await upload_files(sio, files_to_upload, space_id, user_token)
        
        @sio.event
        async def upload_success(data):
            logger.info(f"文件上传成功: {data}")
        
        @sio.event
        async def upload_error(data):
            logger.error(f"文件上传失败: {data}")
        
        @sio.event
        async def error(data):
            logger.error(f"WebSocket 错误: {data}")
        
        @sio.event
        async def disconnect():
            logger.info("WebSocket 连接断开")
        
        # 连接到服务器
        await sio.connect(website_url)
        
        # 等待所有文件上传完成
        await sio.wait()
        
    except Exception as e:
        logger.error(f"WebSocket 连接或上传过程中发生错误: {e}")
        return False

async def upload_files(sio, files_to_upload, space_id, user_token):
    """上传文件列表"""
    success_count = 0
    failed_count = 0
    
    for file_path in files_to_upload:
        try:
            filename = os.path.basename(file_path)
            logger.info(f"正在上传文件: {filename}")
            
            # 读取文件内容并编码为base64
            with open(file_path, 'rb') as f:
                file_content = f.read()
                file_b64 = base64.b64encode(file_content).decode('utf-8')
            
            # 发送文件上传请求
            await sio.emit('file_upload', {
                'space_id': space_id,
                'user_token': user_token,
                'filename': filename,
                'content': file_b64
            })
            
            success_count += 1
            # 稍微延迟避免服务器压力过大
            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}")
    
    # 上传完成后断开连接
    await sio.disconnect()

def main():
    if len(sys.argv) != 4:
        print("用法: python3 websocket_uploader.py <space_id> <user_token> <website_url>")
        sys.exit(1)
    
    space_id = sys.argv[1]
    user_token = sys.argv[2] 
    website_url = sys.argv[3]
    
    logger.info(f"开始 WebSocket 文件上传")
    logger.info(f"Space ID: {space_id}")
    logger.info(f"Website URL: {website_url}")
    
    # 运行异步上传
    asyncio.run(upload_files_via_websocket(space_id, user_token, website_url))

if __name__ == "__main__":
    main()