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

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


def snmp_walk(host, oid, format="str", strip_prefix=True, community="zabbix"):
    res = {}
    # 使用正确的 nextCmd 函数
    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: "{str(errorIndication)}". Status={str(errorStatus)}'
            )
        elif errorStatus:
            raise ConnectionError(
                f'SNMP error: "{str(errorStatus)}" at index {str(errorIndex)}'
            )
        else:
            for varBind in varBinds:
                oid_str = str(varBind[0])
                value = varBind[1]

                # 如果指定了去除前缀
                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


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:])


def read_ipv4_from_oid_tail(oid, with_len=True):
    parts = [int(x) for x in oid.split(".")]
    if with_len:
        assert parts[-5] == 4  # number of elements
    return ".".join([str(x) for x in parts[-4:]])


from datetime import timedelta


def ms_to_timedelta(milliseconds):
    delta = timedelta(milliseconds=milliseconds)
    hours = delta.seconds // 3600
    minutes = (delta.seconds % 3600) // 60
    seconds = delta.seconds % 60
    return f"{hours}:{minutes:02d}:{seconds:02d}"


def format_mac(mac_str: str) -> str:
    # 每2个字符分割一次，生成子字符串列表
    segments = [mac_str[i : i + 2] for i in range(0, len(mac_str), 2)]
    # 用连字符连接子字符串，并添加末尾换行符
    return "-".join(segments)


def main():
    log.info(str(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
    log.info(sys.argv)
    # 1. 获取端口映射表
    port_mapping = {}
    dot1dBasePortIfIndex = snmp_walk(
        switch_ip, "1.3.6.1.2.1.17.1.4.1.2", "int", True, community
    )
    ifDescr = snmp_walk(switch_ip, "1.3.6.1.2.1.2.2.1.2", "str", True, community)
    for bridge_index, if_index in dot1dBasePortIfIndex.items():
        if str(if_index) in ifDescr:
            port_mapping[str(bridge_index)] = ifDescr[str(if_index)]
    # 2. 获取MAC表和ip地址
    mac_table = snmp_walk(switch_ip, "1.3.6.1.2.1.17.4.3.1.2", "int", True, community)

    # 3. 统计端口MAC数量
    port_counts = {}
    mac_details = {}
    for mac_oid, bridge_index in mac_table.items():
        standard_mac = machex(mac_oid).upper()
        if str(bridge_index) in port_mapping:
            port_name = port_mapping[str(bridge_index)]
            port_counts[port_name] = port_counts.get(port_name, 0) + 1

    if mode == "list":
        atPhysAddress = snmp_walk(
            switch_ip, "1.3.6.1.2.1.3.1.1.2", "hex", True, community
        )
        oid_mac = snmp_walk(
            switch_ip,
            "1.3.6.1.2.1.55.1.12.1.2.719.254.128.0.0.0.0.0.0",
            "hex",
            True,
            community,
        )
        oid_hour = snmp_walk(
            switch_ip,
            "1.3.6.1.2.1.55.1.12.1.5.719.254.128.0.0.0.0.0.0",
            "str",
            True,
            community,
        )
        mac_time = {}
        # ms_to_timedelta
        for oid, mac in oid_mac.items():
            f_mac = format_mac(mac)
            time = oid_hour.get(oid, "0")
            time = ms_to_timedelta(int(time + "0"))
            mac_time[f_mac] = time
        mac_ip_table = {}
        for oid, mac in atPhysAddress.items():
            ip = read_ipv4_from_oid_tail(oid, with_len=False)
            mac_ip_table[format_mac(mac)] = ip
        for mac_oid, bridge_index in mac_table.items():
            standard_mac = machex(mac_oid).upper()
            if str(bridge_index) in port_mapping:
                port_name = port_mapping[str(bridge_index)]
                ip = mac_ip_table.get(standard_mac, False)
                mac_details.setdefault(port_name, []).append(
                    {
                        "ip": ip,
                        "mac": standard_mac,
                        "time": mac_time.get(standard_mac, "N/A"),
                    }
                )
    # 3. 如果指定了端口过滤，进行过滤
    # 4. 根据模式输出
    if mode == "discovery":
        # 输出端口自动发现数据
        discovery = []
        for port in port_counts.keys():
            discovery.append(
                {
                    "{#SWITCH_IP}": switch_ip,
                    "{#PORT_NAME}": port,
                    "{#COMMUNITY}": community,
                }
            )
        print(json.dumps({"data": discovery}))

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

        print(str(port_counts.get(port_filter, 0)))

    elif mode == "list":
        # 输出特定端口的MAC列表
        if not port_filter:
            print("[]")
            return
        if port_filter in mac_details:
            pass
            print(json.dumps(mac_details[port_filter]))
        else:
            print("[]")

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


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