#!/usr/bin/env /opt/miniconda3/envs/py37/bin/python
import sys
import json
from datetime import timedelta
from pysnmp.hlapi import *
from mylog import mylog as mylog

log = mylog("./zabbix.log")


class SNMPHelper:
    """封装SNMP操作的工具类"""

    @staticmethod
    def snmp_walk(host, oid, format="str", strip_prefix=True, community="zabbix"):
        """执行SNMP walk操作并返回处理后的数据"""
        res = {}
        for errorIndication, errorStatus, errorIndex, varBinds in nextCmd(
            SnmpEngine(),
            CommunityData(community),
            UdpTransportTarget((host, 161), timeout=4.0, retries=3),
            ContextData(),
            ObjectType(ObjectIdentity(oid)),
            lookupMib=False,
            lexicographicMode=False,
        ):
            if errorIndication:
                raise ConnectionError(
                    f'SNMP error: "{errorIndication}". Status={errorStatus}'
                )
            elif errorStatus:
                raise ConnectionError(
                    f'SNMP error: "{errorStatus}" at index {errorIndex}'
                )

            for varBind in varBinds:
                oid_str = str(varBind[0])
                value = varBind[1]

                # 处理OID前缀
                if strip_prefix:
                    oid_str = oid_str.replace(oid, "", 1).lstrip(".")

                # 格式化返回值
                if format == "hex":
                    res[oid_str] = value.asOctets().hex().upper()
                elif format == "str":
                    res[oid_str] = str(value)
                else:
                    res[oid_str] = value
        return res


class MACUtils:
    """MAC地址相关工具函数"""

    @staticmethod
    def machex(mac_part):
        """将OID格式的MAC转换为标准格式"""
        parts = [int(x) for x in mac_part.split(".")]
        return "-".join(f"{x:02x}" for x in parts[-6:]).upper()

    @staticmethod
    def format_mac(mac_str: str) -> str:
        """格式化MAC地址字符串"""
        segments = [mac_str[i : i + 2] for i in range(0, len(mac_str), 2)]
        return "-".join(segments).upper()


class NetworkUtils:
    """网络相关工具函数"""

    @staticmethod
    def read_ipv4_from_oid(oid):
        """从OID中提取IPv4地址"""
        parts = [int(x) for x in oid.split(".")]
        return ".".join(str(x) for x in parts[-4:])

    @staticmethod
    def ms_to_timedelta(milliseconds):
        """毫秒时间转换为可读的时间格式"""
        delta = timedelta(milliseconds=int(milliseconds))
        hours = delta.seconds // 3600
        minutes = (delta.seconds % 3600) // 60
        seconds = delta.seconds % 60
        return f"{hours}:{minutes:02d}:{seconds:02d}"


class SwitchMonitor:
    def __init__(self, switch_ip, community):
        self.switch_ip = switch_ip
        self.community = community
        self.port_mapping = {}
        self.mac_table = {}
        self.port_counts = {}
        self.mac_details = {}

        # 初始化基础数据
        self._init_port_mapping()
        self._init_mac_table()

    def _init_port_mapping(self):
        """初始化端口映射数据"""
        dot1dBasePortIfIndex = SNMPHelper.snmp_walk(
            self.switch_ip, "1.3.6.1.2.1.17.1.4.1.2", "int", True, self.community
        )
        ifDescr = SNMPHelper.snmp_walk(
            self.switch_ip, "1.3.6.1.2.1.2.2.1.2", "str", True, self.community
        )

        for bridge_index, if_index in dot1dBasePortIfIndex.items():
            if str(if_index) in ifDescr:
                self.port_mapping[bridge_index] = ifDescr[str(if_index)]

    def _init_mac_table(self):
        """初始化MAC表并统计各端口MAC数量"""
        self.mac_table = SNMPHelper.snmp_walk(
            self.switch_ip, "1.3.6.1.2.1.17.4.3.1.2", "int", True, self.community
        )

        for mac_oid, bridge_index in self.mac_table.items():
            if str(bridge_index) in self.port_mapping:
                port_name = self.port_mapping[str(bridge_index)]
                self.port_counts[port_name] = self.port_counts.get(port_name, 0) + 1

    def get_mac_details(self):
        """获取详细的MAC地址信息（IP和在线时间）"""
        if not hasattr(self, "_mac_details"):
            self._mac_details = {}
            atPhysAddress = SNMPHelper.snmp_walk(
                self.switch_ip, "1.3.6.1.2.1.3.1.1.2", "hex", True, self.community
            )

            # 获取MAC时间信息
            oid_mac = SNMPHelper.snmp_walk(
                self.switch_ip,
                "1.3.6.1.2.1.55.1.12.1.2.719.254.128.0.0.0.0.0.0",
                "hex",
                True,
                self.community,
            )
            oid_hour = SNMPHelper.snmp_walk(
                self.switch_ip,
                "1.3.6.1.2.1.55.1.12.1.5.719.254.128.0.0.0.0.0.0",
                "str",
                True,
                self.community,
            )

            mac_time = {}
            for oid, mac in oid_mac.items():
                time_val = oid_hour.get(oid, "0")
                try:
                    # 避免字符串拼接方式转换时间
                    formatted_time = NetworkUtils.ms_to_timedelta(time_val + "0")
                except ValueError:
                    formatted_time = "N/A"
                mac_time[MACUtils.format_mac(mac)] = formatted_time
            # 构建MAC-IP映射
            mac_ip_table = {}
            for oid, mac in atPhysAddress.items():
                try:
                    ip = NetworkUtils.read_ipv4_from_oid(oid)
                    mac_ip_table[MACUtils.format_mac(mac)] = ip
                except (ValueError, IndexError):
                    continue
            # 组织MAC详细信息
            for mac_oid, bridge_index in self.mac_table.items():
                standard_mac = MACUtils.machex(mac_oid)
                port_name = self.port_mapping.get(str(bridge_index))

                if port_name:
                    self._mac_details.setdefault(port_name, []).append(
                        {
                            "ip": mac_ip_table.get(standard_mac, "N/A"),
                            "mac": standard_mac,
                            "time": mac_time.get(standard_mac, "N/A"),
                        }
                    )

        return self._mac_details


def main():
    log.info(f"启动参数: {sys.argv}")

    # 参数验证
    if len(sys.argv) < 4:
        print("Usage: ./switch_monitor.py <mode> <switch_ip> <community> [port]")
        sys.exit(1)

    mode = sys.argv[1]
    switch_ip = sys.argv[2]
    community = sys.argv[3]
    port_filter = sys.argv[4] if len(sys.argv) >= 5 else None

    try:
        monitor = SwitchMonitor(switch_ip, community)

        # 处理不同模式
        if mode == "discovery":
            # 输出端口自动发现数据
            discovery = [
                {
                    "{#SWITCH_IP}": switch_ip,
                    "{#PORT_NAME}": port,
                    "{#COMMUNITY}": community,
                }
                for port in monitor.port_counts
            ]
            print(json.dumps({"data": discovery}))

        elif mode == "count":
            # 输出特定端口的MAC数量
            if not port_filter:
                print("0")
                return

            count = monitor.port_counts.get(port_filter, 0)
            print(str(count))

        elif mode == "list":
            # 输出特定端口的MAC列表
            mac_details = monitor.get_mac_details()
            result = mac_details.get(port_filter, []) if port_filter else []
            print(json.dumps(result))

        elif mode == "details":
            # 输出详细MAC信息（用于调试）
            details = {
                port: {"count": len(macs), "macs": macs}
                for port, macs in monitor.get_mac_details().items()
            }
            print(json.dumps(details))
            log.info(f"详细MAC信息: {json.dumps(details)}")

    except Exception as e:
        log.error(f"处理时发生错误: {str(e)}")
        if mode in ["discovery", "details"]:
            print(json.dumps({}))
        else:
            print("0" if mode == "count" else "[]")


if __name__ == "__main__":
    sys.argv = [
        "/usr/lib/zabbix/externalscripts/switch_monitor.py",
        "count",
        "192.168.0.112",
        "zabbix",
        "GigabitEthernet1/0/5",
    ]
    main()