mirror of
https://github.com/Jerryplusy/AI-powered-switches.git
synced 2025-07-04 21:29:18 +00:00
283 lines
9.6 KiB
Python
283 lines
9.6 KiB
Python
import asyncio
|
||
import logging
|
||
import telnetlib3
|
||
from datetime import datetime
|
||
from pathlib import Path
|
||
from typing import Dict, List, Optional, Union
|
||
from pydantic import BaseModel
|
||
from tenacity import retry, stop_after_attempt, wait_exponential
|
||
import aiofiles
|
||
import asyncssh
|
||
|
||
|
||
# ----------------------
|
||
# 数据模型
|
||
# ----------------------
|
||
class SwitchConfig(BaseModel):
|
||
type: str # vlan/interface/acl/route
|
||
vlan_id: Optional[int] = None
|
||
interface: Optional[str] = None
|
||
name: Optional[str] = None
|
||
ip_address: Optional[str] = None
|
||
vlan: Optional[int] = None # 兼容eNSP模式
|
||
|
||
|
||
# ----------------------
|
||
# 异常类
|
||
# ----------------------
|
||
class SwitchConfigException(Exception):
|
||
pass
|
||
|
||
|
||
class EnspConnectionException(SwitchConfigException):
|
||
pass
|
||
|
||
|
||
class SSHConnectionException(SwitchConfigException):
|
||
pass
|
||
|
||
|
||
# ----------------------
|
||
# 核心配置器
|
||
# ----------------------
|
||
class SwitchConfigurator:
|
||
def __init__(
|
||
self,
|
||
username: str = "admin",
|
||
password: str = "admin",
|
||
timeout: int = 10,
|
||
max_workers: int = 5,
|
||
ensp_mode: bool = False,
|
||
ensp_port: int = 2000,
|
||
ensp_command_delay: float = 0.5,
|
||
**ssh_options
|
||
):
|
||
self.username = username
|
||
self.password = password
|
||
self.timeout = timeout
|
||
self.semaphore = asyncio.Semaphore(max_workers)
|
||
self.backup_dir = Path("config_backups")
|
||
self.backup_dir.mkdir(exist_ok=True)
|
||
self.ensp_mode = ensp_mode
|
||
self.ensp_port = ensp_port
|
||
self.ensp_delay = ensp_command_delay
|
||
self.ssh_options = ssh_options
|
||
self._connection_pool = {} # SSH连接池
|
||
|
||
# ====================
|
||
# 公开API方法
|
||
# ====================
|
||
async def apply_config(self, ip: str, config: Union[Dict, SwitchConfig]) -> Dict:
|
||
"""
|
||
应用配置到交换机（主入口）
|
||
返回格式:
|
||
{
|
||
"status": "success"|"failed",
|
||
"output": str,
|
||
"backup_path": str,
|
||
"error": Optional[str],
|
||
"timestamp": str
|
||
}
|
||
"""
|
||
if isinstance(config, dict):
|
||
config = SwitchConfig(**config)
|
||
|
||
result = await self.safe_apply(ip, config)
|
||
result["timestamp"] = datetime.now().isoformat()
|
||
return result
|
||
|
||
# ====================
|
||
# 内部实现方法
|
||
# ====================
|
||
async def _apply_config(self, ip: str, config: SwitchConfig) -> str:
|
||
"""实际配置逻辑"""
|
||
commands = (
|
||
self._generate_ensp_commands(config)
|
||
if self.ensp_mode
|
||
else self._generate_standard_commands(config)
|
||
)
|
||
return await self._send_commands(ip, commands)
|
||
|
||
async def _send_commands(self, ip: str, commands: List[str]) -> str:
|
||
"""双模式命令发送"""
|
||
return (
|
||
await self._send_ensp_commands(ip, commands)
|
||
if self.ensp_mode
|
||
else await self._send_ssh_commands(ip, commands)
|
||
)
|
||
|
||
async def _send_ensp_commands(self, ip: str, commands: List[str]) -> str:
|
||
"""Telnet协议执行（eNSP）"""
|
||
try:
|
||
reader, writer = await telnetlib3.open_connection(
|
||
host=ip,
|
||
port=self.ensp_port,
|
||
connect_minwait=self.timeout,
|
||
connect_maxwait=self.timeout
|
||
)
|
||
|
||
# 登录流程
|
||
await reader.readuntil(b"Username:")
|
||
writer.write(f"{self.username}\n")
|
||
await reader.readuntil(b"Password:")
|
||
writer.write(f"{self.password}\n")
|
||
await asyncio.sleep(1)
|
||
|
||
# 执行命令
|
||
output = ""
|
||
for cmd in commands:
|
||
writer.write(f"{cmd}\n")
|
||
await asyncio.sleep(self.ensp_delay)
|
||
while True:
|
||
try:
|
||
data = await asyncio.wait_for(reader.read(1024), timeout=1)
|
||
if not data:
|
||
break
|
||
output += data
|
||
except asyncio.TimeoutError:
|
||
break
|
||
|
||
writer.close()
|
||
return output
|
||
except Exception as e:
|
||
raise EnspConnectionException(f"eNSP连接失败: {str(e)}")
|
||
|
||
async def _send_ssh_commands(self, ip: str, commands: List[str]) -> str:
|
||
"""SSH协议执行"""
|
||
async with self.semaphore:
|
||
try:
|
||
if ip not in self._connection_pool:
|
||
self._connection_pool[ip] = await asyncssh.connect(
|
||
host=ip,
|
||
username=self.username,
|
||
password=self.password,
|
||
connect_timeout=self.timeout,
|
||
**self.ssh_options
|
||
)
|
||
|
||
results = []
|
||
for cmd in commands:
|
||
result = await self._connection_pool[ip].run(cmd)
|
||
results.append(result.stdout)
|
||
return "\n".join(results)
|
||
except asyncssh.Error as e:
|
||
if ip in self._connection_pool:
|
||
self._connection_pool[ip].close()
|
||
del self._connection_pool[ip]
|
||
raise SSHConnectionException(f"SSH操作失败: {str(e)}")
|
||
|
||
@staticmethod
|
||
def _generate_ensp_commands(config: SwitchConfig) -> List[str]:
|
||
"""生成eNSP命令序列"""
|
||
commands = ["system-view"]
|
||
if config.type == "vlan":
|
||
commands.extend([
|
||
f"vlan {config.vlan_id}",
|
||
f"description {config.name or ''}"
|
||
])
|
||
elif config.type == "interface":
|
||
commands.extend([
|
||
f"interface {config.interface}",
|
||
"port link-type access",
|
||
f"port default vlan {config.vlan}" if config.vlan else "",
|
||
f"ip address {config.ip_address}" if config.ip_address else ""
|
||
])
|
||
commands.append("return")
|
||
return [c for c in commands if c.strip()]
|
||
|
||
@staticmethod
|
||
def _generate_standard_commands(config: SwitchConfig) -> List[str]:
|
||
"""生成标准CLI命令"""
|
||
commands = []
|
||
if config.type == "vlan":
|
||
commands.extend([
|
||
f"vlan {config.vlan_id}",
|
||
f"name {config.name or ''}"
|
||
])
|
||
elif config.type == "interface":
|
||
commands.extend([
|
||
f"interface {config.interface}",
|
||
f"switchport access vlan {config.vlan}" if config.vlan else "",
|
||
f"ip address {config.ip_address}" if config.ip_address else ""
|
||
])
|
||
return commands
|
||
|
||
async def _get_current_config(self, ip: str) -> str:
|
||
"""获取当前配置"""
|
||
commands = (
|
||
["display current-configuration"]
|
||
if self.ensp_mode
|
||
else ["show running-config"]
|
||
)
|
||
try:
|
||
return await self._send_commands(ip, commands)
|
||
except (EnspConnectionException, SSHConnectionException) as e:
|
||
raise SwitchConfigException(f"配置获取失败: {str(e)}")
|
||
|
||
async def _backup_config(self, ip: str) -> Path:
|
||
"""备份配置到文件"""
|
||
backup_path = self.backup_dir / f"{ip}_{datetime.now().isoformat()}.cfg"
|
||
config = await self._get_current_config(ip)
|
||
async with aiofiles.open(backup_path, "w") as f:
|
||
await f.write(config)
|
||
return backup_path
|
||
|
||
async def _restore_config(self, ip: str, backup_path: Path) -> bool:
|
||
"""从备份恢复配置"""
|
||
try:
|
||
async with aiofiles.open(backup_path) as f:
|
||
config = await f.read()
|
||
commands = (
|
||
["system-view", config, "return"]
|
||
if self.ensp_mode
|
||
else [f"configure terminal\n{config}\nend"]
|
||
)
|
||
await self._send_commands(ip, commands)
|
||
return True
|
||
except Exception as e:
|
||
logging.error(f"恢复失败: {str(e)}")
|
||
return False
|
||
|
||
@retry(
|
||
stop=stop_after_attempt(2),
|
||
wait=wait_exponential(multiplier=1, min=4, max=10)
|
||
)
|
||
async def safe_apply(
|
||
self,
|
||
ip: str,
|
||
config: Union[Dict, SwitchConfig]
|
||
) -> Dict[str, Union[str, bool, Path]]:
|
||
"""安全配置应用（自动回滚）"""
|
||
backup_path = await self._backup_config(ip)
|
||
try:
|
||
result = await self._apply_config(ip, config)
|
||
if not await self._validate_config(ip, config):
|
||
raise SwitchConfigException("配置验证失败")
|
||
return {
|
||
"status": "success",
|
||
"output": result,
|
||
"backup_path": str(backup_path)
|
||
}
|
||
except (EnspConnectionException, SSHConnectionException, SwitchConfigException) as e:
|
||
restore_status = await self._restore_config(ip, backup_path)
|
||
return {
|
||
"status": "failed",
|
||
"error": str(e),
|
||
"backup_path": str(backup_path),
|
||
"restore_success": restore_status
|
||
}
|
||
|
||
async def _validate_config(self, ip: str, config: SwitchConfig) -> bool:
|
||
"""验证配置是否生效"""
|
||
current = await self._get_current_config(ip)
|
||
if config.type == "vlan":
|
||
return f"vlan {config.vlan_id}" in current
|
||
elif config.type == "interface" and config.vlan:
|
||
return f"switchport access vlan {config.vlan}" in current
|
||
return True
|
||
|
||
async def close(self):
|
||
"""清理所有连接"""
|
||
for conn in self._connection_pool.values():
|
||
conn.close()
|
||
self._connection_pool.clear() |