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

"""
iMC REST API 服务
- 使用 HTTP Digest Access Authentication (RFC 2617)
- 支持 nonce 过期自动续约（401 时自动重新握手）
- 功能：ONU 远程重启、光功率查询
"""
import hashlib
import re
import json
import time
import threading
import logging
import requests
from app.core.config import settings

logger = logging.getLogger(__name__)

# iMC 重启错误码映射
REBOOT_ERROR_CODES = {
    '103': 'ONU不存在',
    '119': 'SNMP连接超时',
    '120': '业务割接失败',
    '121': 'ONU未运行',
    '122': '重启失败',
}

# 并发控制：防止同一设备被重复重启
_reboot_locks: dict = {}
_reboot_lock = threading.Lock()


def _normalize_mac(mac: str) -> str:
    """标准化 MAC 为 iMC 要求的格式：1484-7790-4840（大写，4位分组）"""
    clean = mac.replace(':', '').replace('-', '').replace('.', '').upper()
    return f"{clean[0:4]}-{clean[4:8]}-{clean[8:12]}"


class IMCService:
    """iMC REST API 服务封装"""

    def __init__(self):
        self.base_url = settings.IMC_API_URL.rstrip('/')
        self.username = settings.IMC_API_USERNAME
        self.password = settings.IMC_API_PASSWORD
        self.verify_ssl = settings.IMC_API_VERIFY_SSL
        self.connect_timeout = settings.IMC_CONNECT_TIMEOUT
        self.read_timeout = settings.IMC_READ_TIMEOUT
        self.session = requests.Session()
        self.realm = "iMC RESTful Web Services"
        self.nonce = None
        self.nc = 1

    # ── Digest 认证 ──────────────────────────────────────────────────────────

    def _get_digest_auth_header(self, method: str, uri: str) -> str | None:
        """构建 HTTP Digest 认证头，首次调用自动握手获取 nonce"""
        if not self.nonce:
            try:
                resp = self.session.get(
                    f"{self.base_url}{uri}",
                    verify=self.verify_ssl,
                    headers={"Accept": "application/json"},
                    timeout=(self.connect_timeout, self.read_timeout),
                )
                if resp.status_code == 401 and 'WWW-Authenticate' in resp.headers:
                    auth_parts = {}
                    for part in resp.headers['WWW-Authenticate'].split(','):
                        if '=' in part:
                            k, v = part.split('=', 1)
                            auth_parts[k.strip()] = v.strip(' "')
                    self.nonce = auth_parts.get('nonce', '')
                    self.realm = auth_parts.get('realm', self.realm)
                else:
                    logger.error(f"获取 nonce 失败，状态码: {resp.status_code}")
                    return None
            except requests.Timeout:
                raise TimeoutError("iMC 认证超时")
            except Exception as e:
                logger.error(f"获取 nonce 异常: {e}")
                return None

        cnonce = hashlib.md5(str(time.time()).encode()).hexdigest()[:16]
        ha1 = hashlib.md5(f"{self.username}:{self.realm}:{self.password}".encode()).hexdigest()
        ha2 = hashlib.md5(f"{method}:{uri}".encode()).hexdigest()
        response_hash = hashlib.md5(
            f"{ha1}:{self.nonce}:{self.nc:08d}:{cnonce}:auth:{ha2}".encode()
        ).hexdigest()

        auth_value = (
            f'Digest username="{self.username}", realm="{self.realm}", '
            f'nonce="{self.nonce}", uri="{uri}", response="{response_hash}", '
            f'qop=auth, nc={self.nc:08d}, cnonce="{cnonce}"'
        )
        self.nc += 1
        return auth_value

    def _clear_auth(self):
        """清除认证状态（nonce 过期时调用，下次请求自动重新握手）"""
        self.nonce = None
        self.nc = 1

    # ── 重启 ONU ─────────────────────────────────────────────────────────────

    def reboot_onu(self, mac: str) -> dict:
        """
        远程重启 ONU 设备。
        使用 per-MAC 锁防止同一设备并发重启。
        """
        imc_mac = _normalize_mac(mac)

        # 并发锁
        with _reboot_lock:
            if imc_mac not in _reboot_locks:
                _reboot_locks[imc_mac] = threading.Lock()
        lock = _reboot_locks[imc_mac]

        if not lock.acquire(blocking=False):
            return {"success": False, "message": "该设备正在重启中，请稍后再试"}

        try:
            for retry in range(2):
                try:
                    uri = f"/imcrs/epon/onu/reboot?mac={imc_mac}"
                    auth = self._get_digest_auth_header("POST", uri)
                    if not auth:
                        return {"success": False, "message": "认证失败，无法发送重启请求"}

                    resp = self.session.post(
                        f"{self.base_url}{uri}",
                        headers={
                            "Accept": "application/xml",
                            "Content-Type": "application/xml",
                            "Content-Length": "0",
                            "Authorization": auth,
                        },
                        verify=self.verify_ssl,
                        timeout=(self.connect_timeout, self.read_timeout),
                    )

                    if resp.status_code == 200:
                        m = re.search(r"<errorCode>(\d+)</errorCode>", resp.text)
                        if m:
                            code = m.group(1)
                            msg = REBOOT_ERROR_CODES.get(code, f"未知错误(代码: {code})")
                            return {"success": False, "message": f"重启失败: {msg}"}
                        return {"success": True, "message": "设备正在重启，请稍后..."}

                    elif resp.status_code == 401:
                        self._clear_auth()
                        continue

                    return {"success": False, "message": f"重启请求失败(HTTP {resp.status_code})"}

                except TimeoutError:
                    return {"success": False, "message": "iMC 接口超时，请稍后重试"}
                except Exception as e:
                    logger.error(f"重启异常 (retry={retry}): {e}")
                    if retry == 0:
                        time.sleep(2)
                        continue
                    return {"success": False, "message": f"重启异常: {e}"}

            return {"success": False, "message": "重启失败，已达最大重试次数"}
        finally:
            lock.release()

    # ── 光功率查询 ────────────────────────────────────────────────────────────

    def get_optical_power(self, mac: str) -> dict | None:
        """
        获取 ONU 设备光功率信息。
        返回 dict 或 None（失败时）。
        """
        imc_mac = _normalize_mac(mac)

        for retry in range(2):
            try:
                uri = f"/imcrs/epon/onu/onuLightWaneInfo?mac={imc_mac}"
                auth = self._get_digest_auth_header("GET", uri)
                if not auth:
                    logger.error("生成认证头失败，无法获取光功率")
                    return None

                resp = self.session.get(
                    f"{self.base_url}{uri}",
                    headers={
                        "Accept": "application/json",
                        "Content-Type": "application/json",
                        "Authorization": auth,
                    },
                    verify=self.verify_ssl,
                    timeout=(self.connect_timeout, self.read_timeout),
                )

                if resp.status_code == 200:
                    try:
                        data = resp.json()
                        return {
                            "powerIn": data.get("powerIn"),
                            "powerOut": data.get("powerOut"),
                            "bindMac": data.get("bindMac"),
                            "devId": data.get("devId"),
                            "eponDevName": data.get("eponDevName"),
                            "oltIfName": data.get("oltIfName"),
                            "onuIfDesc": data.get("onuIfDesc"),
                        }
                    except json.JSONDecodeError as e:
                        logger.error(f"解析光功率 JSON 失败: {e}, 内容: {resp.text}")
                        return None

                elif resp.status_code == 401:
                    self._clear_auth()
                    continue

                logger.error(f"光功率 API 失败，状态码: {resp.status_code}")
                return None

            except requests.Timeout:
                raise TimeoutError("iMC 光功率接口请求超时")
            except Exception as e:
                logger.error(f"获取光功率异常 (retry={retry}): {e}")
                if retry == 0:
                    time.sleep(2)
                    continue
                return None

        return None