100 lines
3.7 KiB
Python
100 lines
3.7 KiB
Python
import asyncio
|
|
import asyncssh
|
|
from fastapi import APIRouter, WebSocket, WebSocketDisconnect
|
|
from logger import get_logger
|
|
|
|
logger = get_logger("app.ssh")
|
|
router = APIRouter()
|
|
|
|
@router.websocket("/ws/terminal")
|
|
async def websocket_terminal(websocket: WebSocket, host: str, username: str, password: str):
|
|
await websocket.accept()
|
|
conn = None
|
|
process = None
|
|
ws_task = None
|
|
ssh_task = None
|
|
try:
|
|
# 建立連線
|
|
conn = await asyncssh.connect(host, username=username, password=password, known_hosts=None)
|
|
|
|
# 💡 關鍵修復：使用 term_size 參數來指定寬高 (width, height)
|
|
process = await conn.create_process(
|
|
term_type='xterm-256color',
|
|
term_size=(80, 24), # 正確的 asyncssh 參數格式
|
|
encoding=None
|
|
)
|
|
|
|
async def forward_to_ws():
|
|
try:
|
|
while True:
|
|
data_bytes = await process.stdout.read(8192)
|
|
if not data_bytes:
|
|
logger.debug("設備端主動關閉了 stdout 通道")
|
|
break
|
|
|
|
safe_text = data_bytes.decode('utf-8', errors='replace')
|
|
await websocket.send_text(safe_text)
|
|
except asyncio.CancelledError:
|
|
pass
|
|
except Exception as e:
|
|
logger.error("❌ [WS Forward Error] 讀取設備畫面時發生錯誤", exc_info=True)
|
|
|
|
async def forward_to_ssh():
|
|
try:
|
|
while True:
|
|
data_text = await websocket.receive_text()
|
|
data_bytes = data_text.encode('utf-8')
|
|
|
|
# 攔截 xterm.js 預設的 DEL (\x7f)，轉換為 BS (\x08)
|
|
if b'\x7f' in data_bytes:
|
|
data_bytes = data_bytes.replace(b'\x7f', b'\x08')
|
|
|
|
process.stdin.write(data_bytes)
|
|
await process.stdin.drain()
|
|
except WebSocketDisconnect:
|
|
# 主動拋出，讓外層捕捉以進行資源回收
|
|
# raise
|
|
# 🌟 修正：不要再 raise 拋出去了，直接 return 結束任務，讓外層自然回收
|
|
logger.debug("前端 WebSocket 正常斷開 (使用者重整或關閉網頁)")
|
|
return
|
|
except asyncio.CancelledError:
|
|
pass
|
|
except Exception as e:
|
|
logger.error("❌ [SSH Forward Error] 寫入指令到設備時發生錯誤", exc_info=True)
|
|
|
|
ws_task = asyncio.create_task(forward_to_ws())
|
|
ssh_task = asyncio.create_task(forward_to_ssh())
|
|
|
|
done, pending = await asyncio.wait(
|
|
[ws_task, ssh_task],
|
|
return_when=asyncio.FIRST_COMPLETED
|
|
)
|
|
|
|
for task in pending:
|
|
task.cancel()
|
|
|
|
except WebSocketDisconnect:
|
|
logger.debug("WebSocket 正常斷線，正在終止背景任務...")
|
|
if ws_task and not ws_task.done():
|
|
ws_task.cancel()
|
|
if ssh_task and not ssh_task.done():
|
|
ssh_task.cancel()
|
|
except Exception as e:
|
|
error_msg = f"\r\n\x1b[31mSSH Connection Error: {str(e)}\x1b[0m\r\n"
|
|
try:
|
|
await websocket.send_text(error_msg)
|
|
except:
|
|
pass
|
|
logger.error("❌ [Connection Error] 建立連線或執行過程中發生錯誤", exc_info=True)
|
|
finally:
|
|
# 🌟 確保 process 與 conn 絕對被關閉，防止 Memory Leak
|
|
if process:
|
|
process.close()
|
|
if conn:
|
|
conn.close()
|
|
try:
|
|
await websocket.close()
|
|
except:
|
|
pass
|
|
|