TermControlAgent/app/api/v1/endpoints/ssh.py

from fastapi import APIRouter, HTTPException, status
from app.core.device.manager import device_manager
from app.core.dispatcher import dispatcher
from app.schemas.ssh import SSHExecRequest, SSHExecResponse
from app.utils.log import get_logger

logger = get_logger(__name__)

router = APIRouter()

@router.post("/devices/{device_id}/exec", response_model=SSHExecResponse, summary="执行SSH命令")
async def exec_ssh_command(device_id: str, exec_request: SSHExecRequest):
    """在指定设备上执行SSH命令"""
    try:
        # 检查设备是否存在
        device = await device_manager.get_device(device_id)
        if not device:
            raise HTTPException(
                status_code=status.HTTP_404_NOT_FOUND,
                detail=f"设备 {device_id} 不存在"
            )
        
        # 检查协议是否为SSH
        if device.protocol_type != 'ssh':
            raise HTTPException(
                status_code=status.HTTP_400_BAD_REQUEST,
                detail=f"设备 {device_id} 不是SSH协议，当前协议: {device.protocol_type}"
            )
        
        # 执行SSH命令
        result = await dispatcher.execute_operation(
            device_id=device_id,
            protocol_type=device.protocol_type,
            operation="exec_command",
            command=exec_request.command,
            timeout=exec_request.timeout,
            working_directory=exec_request.working_directory
        )
        
        return result
        
    except HTTPException:
        raise
    except Exception as e:
        logger.error(f"SSH命令执行失败: {device_id}, 错误: {e}")
        raise HTTPException(
            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
            detail=f"SSH命令执行失败: {str(e)}"
        )

@router.post("/devices/{device_id}/upload", summary="上传文件")
async def upload_file(device_id: str, local_path: str, remote_path: str):
    """通过SSH上传文件"""
    try:
        # 检查设备是否存在
        device = await device_manager.get_device(device_id)
        if not device:
            raise HTTPException(
                status_code=status.HTTP_404_NOT_FOUND,
                detail=f"设备 {device_id} 不存在"
            )
        
        # 检查协议是否为SSH
        if device.protocol_type != 'ssh':
            raise HTTPException(
                status_code=status.HTTP_400_BAD_REQUEST,
                detail=f"设备 {device_id} 不是SSH协议，当前协议: {device.protocol_type}"
            )
        
        # 执行文件上传
        result = await dispatcher.execute_operation(
            device_id=device_id,
            protocol_type=device.protocol_type,
            operation="upload_file",
            local_path=local_path,
            remote_path=remote_path
        )
        
        return result
        
    except HTTPException:
        raise
    except Exception as e:
        logger.error(f"SSH文件上传失败: {device_id}, 错误: {e}")
        raise HTTPException(
            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
            detail=f"SSH文件上传失败: {str(e)}"
        )

@router.post("/devices/{device_id}/download", summary="下载文件")
async def download_file(device_id: str, remote_path: str, local_path: str):
    """通过SSH下载文件"""
    try:
        # 检查设备是否存在
        device = await device_manager.get_device(device_id)
        if not device:
            raise HTTPException(
                status_code=status.HTTP_404_NOT_FOUND,
                detail=f"设备 {device_id} 不存在"
            )
        
        # 检查协议是否为SSH
        if device.protocol_type != 'ssh':
            raise HTTPException(
                status_code=status.HTTP_400_BAD_REQUEST,
                detail=f"设备 {device_id} 不是SSH协议，当前协议: {device.protocol_type}"
            )
        
        # 执行文件下载
        result = await dispatcher.execute_operation(
            device_id=device_id,
            protocol_type=device.protocol_type,
            operation="download_file",
            remote_path=remote_path,
            local_path=local_path
        )
        
        return result
        
    except HTTPException:
        raise
    except Exception as e:
        logger.error(f"SSH文件下载失败: {device_id}, 错误: {e}")
        raise HTTPException(
            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
            detail=f"SSH文件下载失败: {str(e)}"
        )

@router.get("/devices/{device_id}/test-connection", summary="测试SSH连接")
async def test_ssh_connection(device_id: str):
    """测试SSH连接状态"""
    try:
        # 检查设备是否存在
        device = await device_manager.get_device(device_id)
        if not device:
            raise HTTPException(
                status_code=status.HTTP_404_NOT_FOUND,
                detail=f"设备 {device_id} 不存在"
            )
        
        # 检查协议是否为SSH
        if device.protocol_type != 'ssh':
            raise HTTPException(
                status_code=status.HTTP_400_BAD_REQUEST,
                detail=f"设备 {device_id} 不是SSH协议，当前协议: {device.protocol_type}"
            )
        
        # 测试连接
        result = await dispatcher.execute_operation(
            device_id=device_id,
            protocol_type=device.protocol_type,
            operation="test_connection"
        )
        
        return {
            "device_id": device_id,
            "connected": result,
            "message": "连接正常" if result else "连接失败"
        }
        
    except HTTPException:
        raise
    except Exception as e:
        logger.error(f"SSH连接测试失败: {device_id}, 错误: {e}")
        raise HTTPException(
            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
            detail=f"SSH连接测试失败: {str(e)}"
        )