scm-harmonic-cmts-admin/routers/terminal.py

97 lines
3.4 KiB
Python
Raw Normal View History

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
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