import sys
from pysnmp.hlapi import *
import pysnmp.proto.rfc1902 as rfc1902


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 split_numbers(oid):
    return [int(x) for x in oid.split(".")]


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


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


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


def machex(getvar):
    macs = getvar.split(".")
    i = 0
    ma = []
    for x in range(0, 6):
        maca = macs[i]
        if len(maca) == 1:
            a = hex(int(maca)).replace("x", "")
        else:
            a = hex(int(maca))[2:]
        ma.append(a)
        i = i + 1
    return ma[0] + ":" + ma[1] + ":" + ma[2] + ":" + ma[3] + ":" + ma[4] + ":" + ma[5]


if __name__ == "__main__":

    # Read ARP table
    print(" - Reading device ARP table...", file=sys.stderr)
    atPhysAddress = snmp_walk(
        "192.168.0.112", "1.3.6.1.2.1.3.1.1.2", "hex", community="zabbix"
    )
    for oid, mac in atPhysAddress.items():
        ip = read_ipv4_from_oid_tail(oid, with_len=False)
        print(ip)
        print(mac)
    # Read dot1dBasePortIfIndex table
    print(" - Reading device dot1dBasePortIfIndex table...", file=sys.stderr)
    dot1dBasePortIfIndex = snmp_walk(
        "192.168.0.112", "1.3.6.1.2.1.17.1.4.1.2", "int", community="zabbix"
    )
    dot1dBasePort = {}
    for bid, id in dot1dBasePortIfIndex.items():
        ip = read_ipv4_from_oid_tail(bid, with_len=False)
        print("bid=", bid)
        print("id=", id)
        dot1dBasePort[bid] = str(id)
    print(dot1dBasePort)

    # Read ifDescr table
    print(" - Reading device ifDescr table...", file=sys.stderr)
    ifDescr = snmp_walk(
        "192.168.0.112", "1.3.6.1.2.1.2.2.1.2", "str", community="zabbix"
    )
    Descr = {}
    for id, desc in ifDescr.items():
        ip = read_ipv4_from_oid_tail(id, with_len=False)
        print("id=", id)
        print("desc=", desc)
        Descr[id] = desc
    print(Descr)

    dot1dBasePortDescr = {}
    for key in dot1dBasePort.keys():
        dot1dBasePortDescr[key] = Descr[dot1dBasePort[key]]

    print(dot1dBasePortDescr)

    # Read dot1qTpFdbPort table
    print(" - Reading device dot1qTpFdbPort table...", file=sys.stderr)
    dot1qTpFdbPort = snmp_walk(
        "192.168.0.112", "1.3.6.1.2.1.17.4.3.1.2", "int", community="zabbix"
    )
    dot1qTpFdb = {}
    for mac, bid in dot1qTpFdbPort.items():
        macdec = read_mac_from_oid_tail(mac, with_len=False)
        print("machex=", machex(macdec))
        print("bid=", bid)
        dot1qTpFdb[machex(macdec)] = str(bid)
    print(dot1qTpFdb)

    dot1qTpFdbDescr = {}
    for key in dot1qTpFdb.keys():
        if dot1qTpFdb[key] in dot1dBasePortDescr.keys():
            dot1qTpFdbDescr[key] = dot1dBasePortDescr[dot1qTpFdb[key]]

    print(dot1qTpFdbDescr)