remote_uploader.py

import asyncio
import websockets
import os
import time
import base64
import logging
import argparse
import json
import urllib.parse

# Configure logging
logging.basicConfig(level=logging.DEBUG, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

async def upload_directory(space_id, api_key, upload_dir):
    # 使用查询参数传递 api_key
    encoded_secret = urllib.parse.quote(api_key)
    uri = f"wss://remote-terminal-worker.nianxi4563.workers.dev/terminal/{space_id}?secret={encoded_secret}"
    logger.info(f"Attempting to connect to {uri}")

    try:
        async with websockets.connect(uri, ping_interval=20, ping_timeout=60) as websocket:
            logger.info("Connected to WebSocket")

            # 发送机器注册信息（尽管去除 machine_id，但保持兼容）
            machine_info = {"type": "machine", "card_id": space_id}
            await websocket.send(json.dumps(machine_info))
            logger.debug(f"Sent machine registration: {machine_info}")

            # 不等待 machine_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("No files found in directory")
                return 0, 0

            logger.info(f"Found {len(all_files)} files, starting upload...")

            success_count = 0
            failed_count = 0

            for file_path in all_files:
                try:
                    with open(file_path, 'rb') as f:
                        file_content = f.read()
                    file_b64 = base64.b64encode(file_content).decode('utf-8')
                    file_msg = {
                        "type": "file_upload",
                        "filename": os.path.basename(file_path),
                        "content": file_b64
                    }
                    await websocket.send(json.dumps(file_msg))
                    logger.info(f"Uploaded file: {file_path}")
                    success_count += 1
                except Exception as e:
                    logger.error(f"Error uploading file {file_path}: {e}")
                    failed_count += 1
                # 稍微延迟，避免服务器压力过大
                await asyncio.sleep(0.5)

            logger.info(f"Upload completed! Success: {success_count}, Failed: {failed_count}")
            return success_count, failed_count

    except Exception as e:
        logger.error(f"Failed to connect or upload: {e}")
        return 0, len(os.listdir(upload_dir)) if os.path.exists(upload_dir) else 0

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="一次性文件上传器 - 使用 WebSocket 扫描并上传指定文件夹中的所有文件")
    parser.add_argument("api_key", help="您的 API 密钥 (作为用户 token)")
    parser.add_argument("space_id", help="Space ID (作为 card_id)")
    parser.add_argument("--upload-dir", default="output", help="要上传的目录 (默认: output)")

    args = parser.parse_args()

    if not os.path.exists(args.upload_dir):
        print(f"目录不存在: {args.upload_dir}")
        exit(1)

    print(f"🔍 开始扫描目录: {args.upload_dir}")
    print(f"🔑 Space ID (card_id): {args.space_id}")
    print("-" * 50)

    success, failed = asyncio.run(upload_directory(args.space_id, args.api_key, args.upload_dir))

    print("-" * 50)
    print(f"📊 上传完成! 成功: {success}, 失败: {failed}")

    if success > 0:
        print("🎉 文件已成功上传到您的网盘!")