#!/usr/bin/env python3
import sys
import json
from pysnmp.hlapi import *

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()
                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 main():
    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
    print(port_filter, mode, switch_ip, community)
    # 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表
    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)
        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
            mac_details.setdefault(port_name, []).append(standard_mac)
    # 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:
            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))

if __name__ == "__main__":
    main()