H3ConuMS-v2/backend/app/services/ssh_service.py

"""SSH 连接服务"""
import paramiko
import re
import time
from typing import Dict, List, Optional, Tuple
from dataclasses import dataclass

@dataclass
class ONUInfo:
    """ONU 设备完整信息"""
    mac_address: str
    status: str  # online/offline
    distance_m: Optional[int] = None  # 距离（米）
    distance_str: Optional[str] = None  # 距离原始字符串，如 "<1000"
    slot_number: Optional[int] = None  # 插槽号
    port_number: Optional[int] = None  # 端口号
    port_id: Optional[str] = None  # 完整端口标识，如 "1/0/1:1"
    loid: Optional[str] = None  # LOID
    model: Optional[str] = None  # 设备型号


class SSHService:
    """SSH 连接和命令执行服务"""

    def __init__(self, host: str, username: str, password: str, port: int = 22):
        self.host = host
        self.username = username
        self.password = password
        self.port = port
        self.client: Optional[paramiko.SSHClient] = None
        self.shell = None

    def connect(self) -> bool:
        """建立 SSH 连接，等待初始 banner 输出完毕"""
        try:
            self.client = paramiko.SSHClient()
            self.client.set_missing_host_key_policy(paramiko.WarningPolicy())
            self.client.connect(
                hostname=self.host,
                port=self.port,
                username=self.username,
                password=self.password,
                timeout=30
            )
            self.shell = self.client.invoke_shell(width=200, height=50)
            # 等待登录 banner 输出完毕，直到出现命令提示符 ">"
            deadline = time.time() + 10
            buf = ""
            while time.time() < deadline:
                if self.shell.recv_ready():
                    buf += self.shell.recv(4096).decode('utf-8', errors='ignore')
                    if re.search(r'<[^>]+>', buf):
                        break
                else:
                    time.sleep(0.2)
            return True
        except Exception as e:
            raise Exception(f"SSH 连接失败: {str(e)}")

    def execute_command(self, command: str) -> str:
        """执行命令并处理 More 分页，等待命令提示符出现后返回"""
        if not self.shell:
            raise Exception("SSH 未连接")

        self.shell.send(command + "\n")
        output = ""
        # 等待命令回显出现，再开始收集输出
        time.sleep(0.5)

        deadline = time.time() + 60
        while time.time() < deadline:
            if self.shell.recv_ready():
                chunk = self.shell.recv(4096).decode('utf-8', errors='ignore')
                output += chunk
                if "---- More ----" in chunk:
                    self.shell.send(" ")
                    time.sleep(0.3)
                elif re.search(r'^<[^>]+>\s*$', chunk, re.MULTILINE):
                    # 匹配整行为 <DEVICE_NAME> 的提示符行（不匹配回显中的 ">"）
                    break
            else:
                time.sleep(0.2)

        return output

    def _send_and_wait(self, cmd: str, expect: str, timeout: int = 10) -> str:
        """发送命令并等待期望字符串出现，超时返回已收集的输出"""
        self.shell.send(cmd + "\n")
        buf = ""
        deadline = time.time() + timeout
        while time.time() < deadline:
            if self.shell.recv_ready():
                buf += self.shell.recv(4096).decode('utf-8', errors='ignore')
                if expect in buf:
                    return buf
            else:
                time.sleep(0.2)
        return buf

    def clear_onu_port(self, port_id: str) -> bool:
        """清除指定端口的 ONU 配置（恢复默认）
        流程: system-view -> interface Onu{port_id} -> default -> Y
        """
        if not self.shell:
            raise Exception("SSH 未连接")

        # 进入系统视图
        out = self._send_and_wait("system-view", "]")
        if "]" not in out:
            raise Exception("进入 system-view 失败")

        # 进入端口
        out = self._send_and_wait(f"interface Onu{port_id}", "]")
        if "]" not in out:
            raise Exception(f"进入端口 Onu{port_id} 失败")

        # 执行 default，等待确认提示
        self.shell.send("default\n")
        buf = ""
        deadline = time.time() + 10
        while time.time() < deadline:
            if self.shell.recv_ready():
                buf += self.shell.recv(4096).decode('utf-8', errors='ignore')
                if "[Y/N]" in buf or "[y/n]" in buf:
                    break
            else:
                time.sleep(0.2)

        if "[Y/N]" not in buf and "[y/n]" not in buf:
            raise Exception("未收到确认提示")

        # 确认
        self.shell.send("Y\n")
        time.sleep(1)
        # 排空缓冲区
        if self.shell.recv_ready():
            self.shell.recv(4096)

        # 退出到用户视图
        self._send_and_wait("quit", "]", timeout=5)
        self._send_and_wait("quit", ">", timeout=5)

        return True

    def detect_loopback(self) -> dict:
        """执行环路检测，返回 {has_loop: bool, interfaces: [str]}"""
        output = self.execute_command("display loopback-detection")
        has_loop = "Loop is detected on following interfaces" in output
        interfaces = []
        if has_loop:
            for line in output.splitlines():
                m = re.match(r'\s+(Onu\S+)', line)
                if m:
                    interfaces.append(m.group(1))
        return {"has_loop": has_loop, "interfaces": interfaces, "raw": output}

    def parse_onu_status(self, output: str) -> Dict[str, str]:
        """解析 ONU 状态输出"""
        devices = {}
        lines = output.split('\n')

        for line in lines:
            # 匹配包含 MAC 地址的行
            mac_match = re.search(r'([0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4})', line, re.IGNORECASE)
            if mac_match:
                mac = mac_match.group(1).lower()
                if re.search(r'\b(up|online)\b', line.lower()):
                    devices[mac] = 'online'
                elif re.search(r'\b(offline|down)\b', line.lower()):
                    devices[mac] = 'offline'

        return devices

    def _clean_output(self, output: str) -> str:
        """清理终端控制字符和 More 分页标记，避免污染解析"""
        # 移除 ANSI 转义序列
        output = re.sub(r'\x1b\[[0-9;]*[a-zA-Z]', '', output)
        # 移除 ---- More ---- 标记（仅标记本身，保留同行后续设备数据）
        output = re.sub(r'---- More ----', '', output)
        # 将独立的 \r（不跟 \n）替换为空，避免覆盖行内容
        output = re.sub(r'\r(?!\n)', '', output)
        return output

    def parse_onu_info(self, output: str) -> Tuple[Dict[str, ONUInfo], Dict[str, List[ONUInfo]]]:
        """增强解析：提取完整 ONU 信息
        返回: (unique_devices, duplicate_devices)
          - unique_devices: MAC -> ONUInfo（每个 MAC 只保留最新端口）
          - duplicate_devices: MAC -> [ONUInfo, ...] （出现在多个端口的 MAC）
        """
        all_records: Dict[str, List[ONUInfo]] = {}
        output = self._clean_output(output)
        lines = output.split('\n')

        current_slot = None

        for line in lines:
            # 检测新的插槽区域: Olt1/0/1
            slot_match = re.search(r'Olt(\d+)/(\d+)/(\d+)', line)
            if slot_match:
                current_slot = int(slot_match.group(3))
                continue

            # 跳过表头行和空行
            if 'MAC' in line and 'LOID' in line:
                continue
            if not line.strip():
                continue

            # 解析设备行
            device = self._parse_device_line(line, current_slot)
            if device:
                if device.mac_address not in all_records:
                    all_records[device.mac_address] = []
                all_records[device.mac_address].append(device)

        unique_devices: Dict[str, ONUInfo] = {}
        duplicate_devices: Dict[str, List[ONUInfo]] = {}

        for mac, records in all_records.items():
            if len(records) == 1:
                unique_devices[mac] = records[0]
            else:
                duplicate_devices[mac] = records
                # unique 中保留在线的，若都离线则保留最后一条
                online = [r for r in records if r.status == 'online']
                unique_devices[mac] = online[0] if online else records[-1]

        return unique_devices, duplicate_devices

    def _parse_device_line(self, line: str, slot: Optional[int]) -> Optional[ONUInfo]:
        """解析单行设备信息"""
        # 匹配 MAC 地址
        mac_match = re.search(r'([0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4})', line, re.IGNORECASE)
        if not mac_match:
            return None

        mac = mac_match.group(1).lower()

        # 提取状态：H3C OLT 不同固件版本可能输出 Up/UP/up/Online/online
        status = 'offline'
        line_lower = line.lower()
        # 检查行末状态字段（避免误匹配 "Onu" 中的字母）
        if re.search(r'\b(up|online)\b', line_lower):
            status = 'online'

        # 提取端口信息: Onu1/0/2:1 -> slot=2, port=1, port_id="1/0/2:1"
        slot_num, port_num, port_id = None, None, None
        port_match = re.search(r'Onu(\d+)/(\d+)/(\d+):(\d+)', line)
        if port_match:
            slot_num = int(port_match.group(3))  # 第三段数字为槽位
            port_num = int(port_match.group(4))  # 冒号后为端口号
            port_id = f"{port_match.group(1)}/{port_match.group(2)}/{port_match.group(3)}:{port_match.group(4)}"

        # 提取距离 - Port 列前的字段，如 "<1000" 或 "N/A"
        distance_str = None
        dist_match = re.search(r'(\S+)\s+Onu\d+/\d+/\d+:\d+', line)
        if dist_match:
            val = dist_match.group(1)
            if val != 'N/A':
                distance_str = val

        # 提取 LOID - MAC 后第一个非空字段
        loid = None
        loid_match = re.search(r'[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}\s+(\S+)', line, re.IGNORECASE)
        if loid_match:
            loid_val = loid_match.group(1)
            if loid_val != 'N/A' and loid_val.isdigit():
                loid = loid_val

        # 提取设备型号 - Port 列之后的第一个字段，如 "WA6520H-EGPON/A"
        model = None
        model_match = re.search(r'Onu\d+/\d+/\d+:\d+\s+(\S+)', line)
        if model_match:
            potential_model = model_match.group(1)
            if potential_model != 'N/A':
                model = potential_model

        return ONUInfo(
            mac_address=mac,
            status=status,
            distance_m=None,
            distance_str=distance_str,
            slot_number=slot_num or slot,
            port_number=port_num,
            port_id=port_id,
            loid=loid,
            model=model
        )

    def get_olt_ports(self) -> list:
        """获取所有 Olt 端口状态，返回 [{'name': 'Olt1/0/1', 'status': 'up'}, ...]"""
        output = self.execute_command("display interface brief")
        ports = []
        for line in output.splitlines():
            line = line.strip()
            if not line.startswith("Olt"):
                continue
            parts = line.split()
            if len(parts) < 2:
                continue
            name = parts[0]
            link = parts[1].upper()
            # ADM 表示 administratively down（手动关闭）
            if link == "ADM":
                status = "adm-down"
            elif link == "UP":
                status = "up"
            else:
                status = "down"
            ports.append({"name": name, "status": status})
        return ports

    def toggle_olt_port(self, port_name: str, action: str) -> bool:
        """开启或关闭 OLT 端口
        action: 'shutdown' 或 'undo shutdown'
        流程: system-view -> interface {port_name} -> shutdown/undo shutdown -> quit -> quit
        """
        if not self.shell:
            raise Exception("SSH 未连接")

        out = self._send_and_wait("system-view", "]")
        if "]" not in out:
            raise Exception("进入 system-view 失败")

        out = self._send_and_wait(f"interface {port_name}", "]")
        if "]" not in out:
            raise Exception(f"进入端口 {port_name} 失败")

        out = self._send_and_wait(action, "]")
        if "]" not in out:
            raise Exception(f"执行 {action} 失败")

        self._send_and_wait("quit", "]", timeout=5)
        self._send_and_wait("quit", ">", timeout=5)
        return True

    def sync_ntp(self, old_server: str, new_server: str) -> bool:
        """同步 NTP 时间服务器配置
        流程: system-view -> undo ntp old -> ntp new -> clock timezone -> quit -> save force
        old_server 若不存在会报错，直接忽略继续执行。
        """
        if not self.shell:
            raise Exception("SSH 未连接")

        out = self._send_and_wait("system-view", "]")
        if "]" not in out:
            raise Exception("进入 system-view 失败")

        # 删除旧 NTP 服务器（若不存在会报错，忽略即可）
        self._send_and_wait(f"undo ntp-service unicast-server {old_server}", "]")

        # 添加新 NTP 服务器
        out = self._send_and_wait(f"ntp-service unicast-server {new_server}", "]")
        if "]" not in out:
            raise Exception(f"配置 NTP 服务器 {new_server} 失败")

        # 设置时区为北京时间
        out = self._send_and_wait("clock timezone Beijing add 08:00:00", "]")
        if "]" not in out:
            raise Exception("配置时区失败")

        # 退出系统视图
        self._send_and_wait("quit", ">")

        # 强制保存配置
        self._send_and_wait("save force", ">", timeout=30)

        return True

    def get_onu_events(self, port_id: str) -> list:
        """
        查询 ONU 上下线事件记录
        命令: display epon onu-event interface Onu{port_id}
        返回: [{'date', 'time', 'event', 'status', 'datetime_str'}, ...]
        时间按倒序（最新在前）返回
        """
        output = self.execute_command(
            f"display epon onu-event interface Onu{port_id}"
        )
        output = self._clean_output(output)
        events = []
        for line in output.splitlines():
            line = line.strip()
            m = re.match(r'(\d{4}/\d{2}/\d{2})\s+(\d{2}:\d{2}:\d{2})\s+(.+)', line)
            if not m:
                continue
            date_str, time_str, rest = m.group(1), m.group(2), m.group(3).strip()
            # 最后一个单词是 ONU Status（Up/Offline），前面整体是 Event 名称
            parts = rest.rsplit(None, 1)
            if len(parts) == 2:
                event, status = parts[0].strip(), parts[1].strip()
            else:
                event, status = rest, ''
            events.append({
                'date': date_str,
                'time': time_str,
                'event': event,
                'status': status,
                'datetime_str': f"{date_str} {time_str}",
            })
        return events

    def __enter__(self):
        """上下文管理器入口，自动连接"""
        self.connect()
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        """上下文管理器出口，自动关闭连接"""
        self.close()
        return False

    def close(self):
        """关闭 SSH 连接"""
        if self.client:
            self.client.close()