51 lines
1.8 KiB
Python
51 lines
1.8 KiB
Python
import asyncssh
|
||
from contextlib import asynccontextmanager
|
||
from collections.abc import AsyncIterator
|
||
from dataclasses import dataclass
|
||
from starlette.applications import Starlette
|
||
# --- 1. SSH 连接参数 ---
|
||
REMOTE_HOST = '202.121.182.208'
|
||
REMOTE_USER = 'koko125'
|
||
# 确保路径正确，Windows 路径建议使用正斜杠 '/'
|
||
PRIVATE_KEY_PATH = 'D:/tool/tool/id_rsa.txt'
|
||
INITIAL_WORKING_DIRECTORY = f'/cluster/home/{REMOTE_USER}/sandbox'
|
||
|
||
|
||
# --- 2. 定义共享上下文 ---
|
||
# 这个数据类将持有所有模块需要共享的资源
|
||
@dataclass
|
||
class SharedAppContext:
|
||
"""在整个服务器生命周期内持有的共享资源"""
|
||
ssh_connection: asyncssh.SSHClientConnection
|
||
sandbox_path: str
|
||
|
||
|
||
# --- 3. 创建生命周期管理器 ---
|
||
# 这是整个架构的核心，负责在服务器启动时连接SSH，在关闭时断开
|
||
@asynccontextmanager
|
||
async def shared_lifespan(app: Starlette) -> AsyncIterator[SharedAppContext]:
|
||
"""
|
||
为整个应用管理共享资源的生命周期。
|
||
"""
|
||
print("主应用启动，正在建立共享 SSH 连接...")
|
||
conn = None
|
||
try:
|
||
# 建立 SSH 连接
|
||
conn = await asyncssh.connect(
|
||
REMOTE_HOST,
|
||
username=REMOTE_USER,
|
||
client_keys=[PRIVATE_KEY_PATH]
|
||
)
|
||
print(f"SSH 连接到 {REMOTE_HOST} 成功！")
|
||
|
||
# 使用 yield 将创建好的共享资源上下文传递给 Starlette 应用
|
||
# 服务器会在此处暂停并开始处理请求
|
||
yield {"shared_context": SharedAppContext(ssh_connection=conn, sandbox_path=INITIAL_WORKING_DIRECTORY)}
|
||
|
||
finally:
|
||
# 当服务器关闭时，yield 之后的代码会被执行
|
||
if conn:
|
||
conn.close()
|
||
await conn.wait_closed()
|
||
print("共享 SSH 连接已关闭。")
|