#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# Read-only L2 topology collector: ARP + BRIDGE-MIB FDB + ifName/ifAlias.
# No device config changes. Outputs /tmp/fdb_hosts.tsv, /tmp/fdb_links.tsv, /tmp/fdb_summary.txt
import json, urllib.request, subprocess, re, sys, os, time
import concurrent.futures as cf
import l2cfg

TOK = 'd3abb125eb301a021ef80a2847a803a7cd2d0b1d3c8337a14d09db9d571369b5'
API = 'http://192.168.87.20:8080/api_jsonrpc.php'

def call(m, p):
    d = json.dumps({"jsonrpc": "2.0", "method": m, "params": p, "id": 1}).encode()
    q = urllib.request.Request(API, data=d, headers={'Content-Type': 'application/json', 'Authorization': 'Bearer ' + TOK})
    return json.loads(urllib.request.urlopen(q, timeout=60).read().decode())

AUTH = {0: 'MD5', 1: 'SHA', 2: 'SHA-224', 3: 'SHA-256', 4: 'SHA-384', 5: 'SHA-512'}
PRIV = {0: 'DES', 1: 'AES', 2: 'AES-192', 3: 'AES-256', 4: 'AES', 5: 'AES-256'}

def get_hosts():
    hs = l2cfg.hosts(call, output=['hostid', 'name'])
    res = []
    for h in hs:
        for i in h['interfaces']:
            if i.get('type') != '2':
                continue
            det = i.get('details') or {}
            ver = str(det.get('version', '1'))
            args = None
            if ver == '3':
                lvl = int(det.get('securitylevel', 0))
                args = ['-v3', '-l', {0: 'noAuthNoPriv', 1: 'authNoPriv', 2: 'authPriv'}.get(lvl, 'authPriv'),
                        '-u', det.get('securityname', '')]
                if lvl >= 1:
                    args += ['-a', AUTH.get(int(det.get('authprotocol', 0)), 'SHA'), '-A', det.get('authpassphrase', '')]
                if lvl >= 2:
                    args += ['-x', PRIV.get(int(det.get('privprotocol', 0)), 'AES'), '-X', det.get('privpassphrase', '')]
            else:
                args = ['-v2c', '-c', det.get('community', 'public')]
            res.append({'ip': i['ip'], 'zname': h['name'], 'args': args})
            break
    return res

def walk(args, ip, oid):
    try:
        r = subprocess.run(['snmpbulkwalk', '-On'] + args + ['-t', '3', '-r', '1', ip, oid],
                           capture_output=True, text=True, timeout=30)
        return r.stdout
    except Exception:
        return ''

def parse(text, base):
    out = []
    for line in text.splitlines():
        if ' = ' not in line:
            continue
        left, right = line.split(' = ', 1)
        o = left.strip().lstrip('.')
        if not o.startswith(base):
            continue
        suf = o[len(base):].lstrip('.')
        out.append((suf, right.strip()))
    return out

def num(v):
    m = re.search(r'-?\d+', v)
    return int(m.group()) if m else None

def strval(v):
    # strip type prefix like STRING: / INTEGER: / Hex-STRING:
    if ':' in v:
        return v.split(':', 1)[1].strip().strip('"')
    return v.strip()

def normmac(s):
    s = s.strip().lower().replace('-', ':').replace(' ', ':')
    parts = re.split(r'[:.]', s)
    if len(parts) == 6:
        try:
            return ':'.join('%02x' % int(p, 16) for p in parts)
        except Exception:
            return s
    return s

BASE_FDB = '1.3.6.1.2.1.17.4.3.1'
BASE_BPIF = '1.3.6.1.2.1.17.1.4.1.2'
OID_IFNAME = '1.3.6.1.2.1.31.1.1.1.1'
OID_IFALIAS = '1.3.6.1.2.1.31.1.1.1.18'
OID_ARP = '1.3.6.1.2.1.4.22.1.2'
OID_SYSNAME = '1.3.6.1.2.1.1.5.0'
OID_BRIDGE = '1.3.6.1.2.1.17.1.1.0'

def collect(h):
    ip = h['ip']; args = h['args']
    d = {'ip': ip, 'zname': h['zname'], 'name': h['zname'], 'ok': False}
    sn = walk(args, ip, OID_SYSNAME)
    v = parse(sn, OID_SYSNAME)
    if v:
        d['name'] = strval(v[0][1])
    bm = walk(args, ip, OID_BRIDGE)
    vb = parse(bm, OID_BRIDGE)
    d['bridge_mac'] = normmac(strval(vb[0][1])) if vb else ''
    # base port -> ifIndex
    bp = {}
    for suf, val in parse(walk(args, ip, BASE_BPIF), BASE_BPIF):
        port = num(suf.split('.')[-1])
        ifx = num(val)
        if port is not None and ifx is not None:
            bp[port] = ifx
    # ifName / ifAlias
    ifn = {}
    for suf, val in parse(walk(args, ip, OID_IFNAME), OID_IFNAME):
        ifn[num(suf.split('.')[-1])] = strval(val)
    ifa = {}
    for suf, val in parse(walk(args, ip, OID_IFALIAS), OID_IFALIAS):
        ifa[num(suf.split('.')[-1])] = strval(val)
    # FDB: address table .1 idx=mac, port table .2 idx=mac -> bridgePort
    fdb_addr = {}
    for suf, val in parse(walk(args, ip, BASE_FDB + '.1'), BASE_FDB + '.1'):
        mac = normmac(strval(val))
        fdb_addr[suf] = mac
    fdb = {}  # mac -> ifIndex
    portcnt = {}
    for suf, val in parse(walk(args, ip, BASE_FDB + '.2'), BASE_FDB + '.2'):
        port = num(val)
        mac = fdb_addr.get(suf)
        if mac is None:
            mac = normmac(suf)  # suf is m1.m2..m6
        if port is None:
            continue
        portcnt[port] = portcnt.get(port, 0) + 1
        ifx = bp.get(port)
        if ifx is None:
            continue
        fdb[mac] = ifx
    # ARP
    arp = {}
    for suf, val in parse(walk(args, ip, OID_ARP), OID_ARP):
        parts = suf.split('.')
        if len(parts) >= 5:
            ipx = '.'.join(parts[-4:])
            ipx = '.'.join(str(int(x)) for x in parts[-4:])
            mac = normmac(strval(val))
            if mac and mac != '0:0:0:0:0:0':
                arp[ipx] = mac
    ifx_cnt = {}
    for mac, ifx in fdb.items():
        ifx_cnt[ifx] = ifx_cnt.get(ifx, 0) + 1
    d.update({'ok': True, 'bp': bp, 'ifn': ifn, 'ifa': ifa, 'fdb': fdb,
              'portcnt': portcnt, 'ifx_cnt': ifx_cnt, 'arp': arp})
    return d

def main():
    hosts = get_hosts()
    sys.stderr.write("hosts=%d\n" % len(hosts))
    results = []
    with cf.ThreadPoolExecutor(max_workers=24) as ex:
        futs = {ex.submit(collect, h): h for h in hosts}
        for k, f in enumerate(cf.as_completed(futs)):
            try:
                r = f.result()
            except Exception as e:
                h = futs[f]; r = {'ip': h['ip'], 'zname': h['zname'], 'name': h['zname'], 'ok': False}
            results.append(r)
            if (k + 1) % 20 == 0:
                sys.stderr.write("done %d/%d\n" % (k + 1, len(hosts)))
    ok = [r for r in results if r.get('ok')]
    # switch bridge mac index
    sw_by_bmac = {}
    for r in ok:
        bm = r.get('bridge_mac')
        if bm:
            sw_by_bmac[bm] = r
    # mac -> locations
    mac_loc = {}
    for r in ok:
        for mac, ifx in r['fdb'].items():
            mac_loc.setdefault(mac, []).append((r, ifx, r['ifx_cnt'].get(ifx, 0)))
    def portname(r, ifx):
        return r['ifn'].get(ifx, 'if%d' % ifx)
    def portalias(r, ifx):
        return r['ifa'].get(ifx, '')
    # helper to get port num for ifIndex (reverse bp)
    # ip -> mac globally
    ip2mac = {}
    for r in ok:
        for ip, mac in r['arp'].items():
            ip2mac.setdefault(ip, mac)
    # zname by ip
    zname = {h['ip']: h['zname'] for h in hosts}
    # host edges: ip -> mac -> edge switch port (min port mac count)
    host_rows = []
    for ip, mac in sorted(ip2mac.items()):
        if mac in sw_by_bmac:
            continue
        locs = mac_loc.get(mac, [])
        if not locs:
            continue
        best = min(locs, key=lambda t: t[2])
        r, ifx, pc = best
        host_rows.append([ip, zname.get(ip, ''), mac, r['name'], r['ip'], portname(r, ifx), portalias(r, ifx), str(pc)])
    # switch links: for switch S, fdb mac that is another switch bridge mac
    links = []
    seen = set()
    for r in ok:
        for mac, ifx in r['fdb'].items():
            tgt = sw_by_bmac.get(mac)
            if tgt is None or tgt['ip'] == r['ip']:
                continue
            key = tuple(sorted([r['ip'], tgt['ip']])) + (portname(r, ifx),)
            if key in seen:
                continue
            seen.add(key)
            links.append([r['name'], r['ip'], portname(r, ifx), portalias(r, ifx), tgt['name'], tgt['ip']])
    # candidate uplinks: ports (non-loop) with high mac count
    uplinks = []
    for r in ok:
        # invert fdb to port -> count already in portcnt per bridge port; map to ifIndex
        inv = {}
        for mac, ifx in r['fdb'].items():
            inv[ifx] = inv.get(ifx, 0) + 1
        for ifx, c in inv.items():
            if c >= 8:
                uplinks.append([r['name'], r['ip'], portname(r, ifx), portalias(r, ifx), str(c)])
    uplinks.sort(key=lambda x: -int(x[4]))
    with open('/tmp/fdb_hosts.tsv', 'w', encoding='utf-8') as f:
        f.write('ip\tzabbix_name\tmac\tswitch_name\tswitch_ip\tport\tport_alias\tport_mac_count\n')
        for row in host_rows:
            f.write('\t'.join(row) + '\n')
    with open('/tmp/fdb_links.tsv', 'w', encoding='utf-8') as f:
        f.write('a_name\ta_ip\ta_port\ta_alias\tb_name\tb_ip\n')
        for row in links:
            f.write('\t'.join(row) + '\n')
    with open('/tmp/fdb_uplinks.tsv', 'w', encoding='utf-8') as f:
        f.write('switch_name\tswitch_ip\tport\tport_alias\tmac_count\n')
        for row in uplinks:
            f.write('\t'.join(row) + '\n')
    with open('/tmp/fdb_summary.txt', 'w', encoding='utf-8') as f:
        f.write('switches_ok=%d / total=%d\n' % (len(ok), len(hosts)))
        f.write('host_edges=%d\n' % len(host_rows))
        f.write('switch_links=%d\n' % len(links))
        f.write('uplink_ports(>=8mac)=%d\n' % len(uplinks))
        f.write('ips_with_mac=%d\n' % len(ip2mac))
    print(open('/tmp/fdb_summary.txt').read())

if __name__ == '__main__':
    main()