110 lines
4.2 KiB
Python
110 lines
4.2 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 asyncio.wait_for(
|
||
asyncssh.connect(host, username=username, password=password, known_hosts=None),
|
||
timeout=10.0
|
||
)
|
||
|
||
# 💡 關鍵修復：使用 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(f"❌ [WS Forward Error] 讀取設備畫面時發生錯誤: {str(e)}")
|
||
|
||
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(f"❌ [SSH Forward Error] 寫入指令到設備時發生錯誤: {str(e)}")
|
||
|
||
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:
|
||
err_str = str(e)
|
||
error_msg = f"\r\n\x1b[31mSSH Connection Error: {err_str}\x1b[0m\r\n"
|
||
try:
|
||
# 先把錯誤印在終端機畫面上
|
||
await websocket.send_text(error_msg)
|
||
|
||
# 🌟 關鍵修復：使用自訂斷線碼 4001，並附上錯誤原因 (WebSocket 規範 reason 最長 123 bytes)
|
||
safe_reason = err_str[:120] if err_str else "Authentication or Connection Failed"
|
||
await websocket.close(code=4001, reason=safe_reason)
|
||
except:
|
||
pass
|
||
logger.error(f"❌ [Connection Error] 建立連線或執行過程中發生錯誤: {str(e)}")
|
||
return # 🌟 提早結束，避免下方的 finally 再次執行 close 導致報錯
|
||
finally:
|
||
# 🌟 確保 process 與 conn 絕對被關閉，防止 Memory Leak
|
||
if process:
|
||
process.close()
|
||
if conn:
|
||
conn.close()
|
||
try:
|
||
await websocket.close()
|
||
except:
|
||
pass
|
||
|