Spaces:
Runtime error
Runtime error
import asyncio | |
import websockets | |
import subprocess | |
import json | |
import logging | |
import sys | |
import urllib.parse | |
import base64 | |
import os | |
import threading | |
import requests | |
from urllib.parse import urlparse | |
import argparse | |
# Configure logging | |
logging.basicConfig(level=logging.DEBUG, format='%(asctime)s - %(levelname)s - %(message)s') | |
logger = logging.getLogger(__name__) | |
# 全局变量,用于在用户输入时访问 WebSocket | |
global_websocket = None | |
async def connect(card_id, machine_secret, upload_files=None, server="wss://remote-terminal-worker.nianxi4563.workers.dev/terminal/"): | |
global global_websocket | |
# 使用查询参数传递 machine_secret | |
encoded_secret = urllib.parse.quote(machine_secret) | |
uri = f"{server}{card_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: | |
global_websocket = websocket | |
logger.info("Connected to WebSocket") | |
machine_info = {"type": "machine", "card_id": card_id, "token": machine_secret} # 添加用户token | |
await websocket.send(json.dumps(machine_info)) | |
logger.debug(f"Sent machine registration: {machine_info}") | |
# 如果指定了上传文件列表,发送文件内容 | |
if upload_files: | |
for upload_file in upload_files: | |
if os.path.exists(upload_file): | |
try: | |
with open(upload_file, '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(upload_file), | |
"content": file_b64 | |
# 去除 machine_id | |
} | |
await websocket.send(json.dumps(file_msg)) | |
logger.info(f"Uploaded file: {upload_file}") | |
except Exception as e: | |
logger.error(f"Error uploading file {upload_file}: {e}") | |
else: | |
logger.error(f"File not found: {upload_file}") | |
# 等待一段时间以接收任何响应,然后关闭 | |
try: | |
await asyncio.sleep(5) # 等待服务器响应 | |
except: | |
pass | |
except Exception as e: | |
logger.error(f"Failed to connect or maintain connection: {e}") | |
if __name__ == "__main__": | |
parser = argparse.ArgumentParser(description="WebSocket-based File Uploader for remote control file transfer") | |
parser.add_argument("api_key", help="API Key (used as machine_secret and token)") | |
parser.add_argument("space_id", help="Space ID (used as card_id)") | |
parser.add_argument("--server", default="wss://remote-terminal-worker.nianxi4563.workers.dev/terminal/", help="WebSocket server base URL (default: wss://remote-terminal-worker.nianxi4563.workers.dev/terminal/)") | |
parser.add_argument("--upload-dir", default="output", help="Directory to upload files from (default: output)") | |
args = parser.parse_args() | |
if not os.path.exists(args.upload_dir): | |
print(f"目录不存在: {args.upload_dir}") | |
sys.exit(1) | |
all_files = [] | |
for root, dirs, files in os.walk(args.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: | |
print("目录中没有找到任何文件") | |
sys.exit(0) | |
print(f"找到 {len(all_files)} 个文件,开始上传...") | |
asyncio.run(connect(args.space_id, args.api_key, upload_files=all_files, server=args.server)) | |
print("上传完成!") |