#!/usr/bin/env python3
"""
三层交换机上行流量检测工具
连接核心交换机 172.88.88.1 (sywz-core-sw)，检测：
1. 各 uplink 口 (tengige 万兆口) 的流量速率、带宽利用率、错误包
2. ARP 表 VLAN 分布
3. 接口状态概览
4. 生成流量健康报告

用法: python check_uplink_traffic.py
"""
import telnetlib
import re
import json
import os
import sys
import time
from datetime import datetime

# ===== 配置 =====
SWITCH_HOST = "172.88.88.1"
SWITCH_USER = "admin"
SWITCH_PASS = "sywz8208wzwz"
OUTPUT_DIR = os.path.dirname(os.path.abspath(__file__))

# 所有 tengige 口
TENGIGE_PORTS = [f"tengige3_{i}" for i in range(32)]


def telnet_session(host, user, password):
    """建立 Telnet 会话并登录"""
    tn = telnetlib.Telnet(host, 23, timeout=15)
    result = tn.read_until(b'Username:', 10)
    if not result:
        result = tn.read_until(b'login:', 5)
    tn.write((user + "\n").encode('ascii'))
    time.sleep(0.5)
    result = tn.read_until(b'Password:', 8)
    tn.write((password + "\n").encode('ascii'))
    output = b""
    for _ in range(20):
        time.sleep(0.5)
        try:
            chunk = tn.read_very_eager()
            if chunk:
                output += chunk
        except Exception:
            pass
        if b'>' in output or b'#' in output:
            break
    return tn


def send_cmd(tn, cmd, wait=10, auto_more=False):
    """
    发送命令并获取完整输出。
    使用 \\n 而不是 \\r\\n，避免交换机退格行为。
    """
    # 先清缓冲区
    time.sleep(0.3)
    try:
        while tn.read_very_eager():
            pass
    except Exception:
        pass

    tn.write((cmd + "\n").encode('ascii'))
    time.sleep(0.4)

    output = ""
    more_sent = 0
    for _ in range(wait * 2):
        time.sleep(0.5)
        try:
            chunk = tn.read_very_eager().decode('ascii', errors='ignore')
            if chunk:
                output += chunk
        except Exception:
            pass
        # 处理分页
        if auto_more and '--More' in output and more_sent < 100:
            more_sent += 1
            tn.write(b' ')
            time.sleep(0.3)
            continue
        # 检测命令提示符: 以 > 或 # 结尾的行
        stripped = output.rstrip()
        lines = stripped.split('\n')
        if lines and (lines[-1].strip().endswith('>') or lines[-1].strip().endswith('#') or lines[-1].strip().endswith(']')):
            break
    return output


def clean_output(output):
    """清洗控制字符和 More 标记"""
    # 移除 More 行
    cleaned = re.sub(r'.*--More.*--.*\n?', '', output)
    # 移除退格符及其覆盖的字符
    cleaned = re.sub(r'.[\x08]+.', '', cleaned)
    return cleaned


# ─── 解析函数 ───


def parse_status_table(output):
    """
    解析 show interface status 输出
    返回: [{name, description, speed, duplex, auto, state, vlan}, ...]
    """
    cleaned = clean_output(output)
    ports = []
    in_table = False
    for line in cleaned.split('\n'):
        line = line.strip()
        if not line:
            continue
        if line.startswith('port-name'):
            in_table = True
            continue
        if in_table:
            if line.startswith('auto') or line.startswith('<') or 'show interface' in line.lower():
                break
            parts = line.split()
            if len(parts) >= 6:
                ports.append({
                    'name': parts[0],
                    'description': parts[1],
                    'speed': parts[2],
                    'duplex': parts[3],
                    'auto': parts[4],
                    'state': parts[5],
                    'vlan': ' '.join(parts[6:]) if len(parts) > 6 else '',
                })
    return ports


def parse_interface_detail(output):
    """
    解析 show interface tengigeX_X 输出
    返回: {status, rx_bps, tx_bps, rx_bytes, tx_bytes, rx_errors, tx_errors,
           crc_errors, bandwidth_mbps, rx_packets, tx_packets, pvid, vlan_belong, ...}
    """
    cleaned = clean_output(output)
    result = {
        'status': 'unknown',
        'admin_state': 'unknown',
        'line_state': 'unknown',
        'rx_bps': 0, 'tx_bps': 0,
        'rx_mbps': 0, 'tx_mbps': 0,
        'rx_bytes': 0, 'tx_bytes': 0,
        'rx_packets': 0, 'tx_packets': 0,
        'rx_errors': 0, 'tx_errors': 0,
        'crc_errors': 0,
        'length_errors': 0,
        'frame_errors': 0,
        'rx_lost_pct': 0, 'tx_lost_pct': 0,
        'bandwidth_mbps': 10000,
        'utilization_rx': 0, 'utilization_tx': 0,
        'pvid': '',
        'vlan_belong': '',
        'description': '',
    }

    for line in cleaned.split('\n'):
        line = line.strip()
        if not line or 'show interface' in line.lower():
            continue

        # 管理状态
        m = re.search(r'administration\s+state\s+is\s+(\w+)', line)
        if m:
            result['admin_state'] = m.group(1)
        m = re.search(r'line\s+state\s+is\s+(\w+)', line)
        if m:
            result['line_state'] = m.group(1)

        # 综合状态
        if result['admin_state'] == 'UP' and result['line_state'] == 'UP':
            result['status'] = 'up'
        elif result['admin_state'] == 'UP' and result['line_state'] != 'UP':
            result['status'] = 'line_down'
        else:
            result['status'] = result['admin_state'].lower()

        # 描述
        m = re.search(r'Description:\s*(.+)', line)
        if m:
            result['description'] = m.group(1).strip()

        # 实时速率 (两种格式)
        m = re.search(r'Input\s+speed:\s*[\d,]+\s*pps,\s*([\d,]+)\s*bps', line)
        if m:
            result['rx_bps'] = int(m.group(1).replace(',', ''))
        m = re.search(r'Output\s+speed:\s*[\d,]+\s*pps,\s*([\d,]+)\s*bps', line)
        if m:
            result['tx_bps'] = int(m.group(1).replace(',', ''))

        # 收包统计
        m = re.search(r'Input\(normal\):\s*([\d,]+)\s*packets.*?([\d,.]+)%\s*(?:rxpackets\s*lost)?', line)
        if m:
            result['rx_packets'] = int(m.group(1).replace(',', ''))
            result['rx_lost_pct'] = float(m.group(2).replace(',', ''))
        m = re.search(r'Input\(normal\).*?([\d,]+)\s*bytes', line)
        if m:
            result['rx_bytes'] = int(m.group(1).replace(',', ''))

        # 发包统计
        m = re.search(r'Output\(normal\):\s*([\d,]+)\s*packets.*?([\d,.]+)%\s*(?:txpackets\s*lost)?', line)
        if m:
            result['tx_packets'] = int(m.group(1).replace(',', ''))
            result['tx_lost_pct'] = float(m.group(2).replace(',', ''))
        m = re.search(r'Output\(normal\).*?([\d,]+)\s*bytes', line)
        if m:
            result['tx_bytes'] = int(m.group(1).replace(',', ''))

        # 入向错误
        m = re.search(r'Input:\s*(\d+)\s*input\s*errors', line)
        if m:
            result['rx_errors'] = int(m.group(1))
        m = re.search(r'(\d+)\s*length_errors', line)
        if m:
            result['length_errors'] = int(m.group(1))
        m = re.search(r'(\d+)\s*crc_errors', line)
        if m:
            result['crc_errors'] = int(m.group(1))
        m = re.search(r'(\d+)\s*frame_errors', line)
        if m:
            result['frame_errors'] = int(m.group(1))

        # 出向错误
        m = re.search(r'Output:\s*(\d+)\s*output\s*errors', line)
        if m:
            result['tx_errors'] = int(m.group(1))

        # 速率/双工
        m = re.search(r'(\d+)Mbps-speed', line)
        if m:
            result['bandwidth_mbps'] = int(m.group(1))

        # PVID 和 VLAN 归属
        m = re.search(r'pvid\s+(\d+)', line)
        if m:
            result['pvid'] = m.group(1)
        m = re.search(r'vlan\s+belong\s+(.+)', line)
        if m:
            result['vlan_belong'] = m.group(1).strip()

    # 计算 Mbps 和利用率
    result['rx_mbps'] = round(result['rx_bps'] / 1_000_000, 3)
    result['tx_mbps'] = round(result['tx_bps'] / 1_000_000, 3)
    bw = max(result['bandwidth_mbps'], 1)
    result['utilization_rx'] = round(result['rx_mbps'] / bw * 100, 2)
    result['utilization_tx'] = round(result['tx_mbps'] / bw * 100, 2)

    return result


def parse_arp(output):
    """解析 show arp all 输出"""
    cleaned = clean_output(output)
    records = []
    in_table = False
    for line in cleaned.split('\n'):
        line = line.strip()
        if not line:
            continue
        if 'IP Address' in line or 'Ip Address' in line:
            in_table = True
            continue
        if not in_table:
            continue
        if line.startswith('[') or 'show arp' in line.lower() or 'Type: S-Static' in line:
            continue
        parts = line.split()
        if len(parts) >= 5:
            records.append({
                'ip': parts[0],
                'mac': parts[1],
                'vlan': parts[2],
                'interface': parts[3],
                'age': parts[4],
            })
    return records


def health_check(interfaces):
    """评估接口健康状态"""
    issues = []
    warnings = []

    for name, info in interfaces.items():
        status = info.get('status', 'unknown')

        if status != 'up':
            issues.append(f"⚠️  {name}: 状态={status}（非UP状态）")
            continue

        rx_util = info.get('utilization_rx', 0)
        tx_util = info.get('utilization_tx', 0)

        if rx_util > 80:
            warnings.append(f"🔴 {name}: 入向利用率 {rx_util}% (>80%)")
        elif rx_util > 60:
            warnings.append(f"🟡 {name}: 入向利用率 {rx_util}% (>60%)")

        if tx_util > 80:
            warnings.append(f"🔴 {name}: 出向利用率 {tx_util}% (>80%)")
        elif tx_util > 60:
            warnings.append(f"🟡 {name}: 出向利用率 {tx_util}% (>60%)")

        rx_err = info.get('rx_errors', 0)
        tx_err = info.get('tx_errors', 0)
        crc_err = info.get('crc_errors', 0)

        if rx_err > 100:
            warnings.append(f"🔴 {name}: 入向错误 {rx_err}")
        elif rx_err > 0:
            warnings.append(f"🟡 {name}: 入向错误 {rx_err}")

        if tx_err > 100:
            warnings.append(f"🔴 {name}: 出向错误 {tx_err}")
        elif tx_err > 0:
            warnings.append(f"🟡 {name}: 出向错误 {tx_err}")

        if crc_err > 10:
            warnings.append(f"🔴 {name}: CRC错误 {crc_err}（物理线路可能有问题）")
        elif crc_err > 0:
            warnings.append(f"🟡 {name}: CRC错误 {crc_err}")

        lost = info.get('rx_lost_pct', 0) + info.get('tx_lost_pct', 0)
        if lost > 1:
            warnings.append(f"🔴 {name}: 丢包率 {lost}% (>1%)")
        elif lost > 0:
            warnings.append(f"🟡 {name}: 丢包率 {lost}%")

    return issues, warnings


# ─── 主流程 ───


def main():
    sys.stdout.reconfigure(encoding='utf-8')
    ts = datetime.now().strftime('%Y%m%d_%H%M%S')
    run_time = datetime.now().strftime('%Y-%m-%d %H:%M:%S')

    print("=" * 70)
    print("  三层交换机上行流量健康检测")
    print(f"  目标: {SWITCH_HOST}")
    print(f"  时间: {run_time}")
    print("=" * 70)

    # ── 连接 ──
    print("\n🔌 连接交换机...")
    tn = telnet_session(SWITCH_HOST, SWITCH_USER, SWITCH_PASS)
    print("  连接成功")

    # 关闭分页
    send_cmd(tn, "terminal length 0", 3)
    send_cmd(tn, "no paging", 3)
    time.sleep(0.3)

    # ── Step 1: 接口状态总览 ──
    print("\n" + "=" * 70)
    print("  [1/4] 采集接口状态总览")
    print("=" * 70)
    status_out = send_cmd(tn, "show interface status", 12, auto_more=True)
    status_ports = parse_status_table(status_out)
    print(f"  总端口: {len(status_ports)}")

    # 从 status 中构建接口状态字典
    all_interfaces = {}
    for p in status_ports:
        if p['name'].startswith('tengige'):
            bw = 10000
            if p['speed'] == '1000M':
                bw = 1000
            elif p['speed'] == '10000M':
                bw = 10000
            all_interfaces[p['name']] = {
                'status': p['state'],
                'speed': p['speed'],
                'vlan': p['vlan'],
                'rx_bps': 0, 'tx_bps': 0,
                'rx_mbps': 0, 'tx_mbps': 0,
                'rx_bytes': 0, 'tx_bytes': 0,
                'rx_packets': 0, 'tx_packets': 0,
                'rx_errors': 0, 'tx_errors': 0,
                'crc_errors': 0,
                'length_errors': 0,
                'frame_errors': 0,
                'rx_lost_pct': 0, 'tx_lost_pct': 0,
                'bandwidth_mbps': bw,
                'utilization_rx': 0, 'utilization_tx': 0,
                'pvid': '',
                'vlan_belong': '',
                'description': '',
            }

    # 打印状态总览
    up_count = sum(1 for p in status_ports if p['state'] == 'up')
    down_count = sum(1 for p in status_ports if p['state'] == 'down')
    print(f"  UP: {up_count}  |  DOWN: {down_count}")
    print(f"\n  {'端口':<18} {'速率':<10} {'状态':<8} {'VLAN归属'}")
    print(f"  {'-'*65}")
    for p in status_ports:
        if p['name'].startswith('tengige'):
            state_icon = '🟢' if p['state'] == 'up' else '🔴'
            print(f"  {p['name']:<18} {p['speed']:<10} {state_icon} {p['state']:<5} {p['vlan']}")

    # ── Step 2: 逐个采集 uplink 流量 ──
    print("\n" + "=" * 70)
    print("  [2/4] 采集上行接口流量详情")
    print("=" * 70)

    up_tengige = [p['name'] for p in status_ports if p['name'].startswith('tengige') and p['state'] == 'up']
    print(f"  共有 {len(up_tengige)} 个 UP 状态的 tengige 口需要采集")

    for i, iface in enumerate(up_tengige):
        print(f"  [{i+1}/{len(up_tengige)}] {iface}...", end=' ', flush=True)
        detail_out = send_cmd(tn, f"show interface {iface}", 6)
        detail = parse_interface_detail(detail_out)
        # 合并流量数据 — 状态以 status 表为准，不可覆盖
        for k, v in detail.items():
            if k == 'status' or k == 'admin_state' or k == 'line_state':
                continue
            if v is not None and v != '':
                all_interfaces[iface][k] = v
        print(f"rx={all_interfaces[iface]['rx_mbps']:.2f}Mbps tx={all_interfaces[iface]['tx_mbps']:.2f}Mbps "
              f"err={all_interfaces[iface]['rx_errors']+all_interfaces[iface]['tx_errors']} crc={all_interfaces[iface]['crc_errors']}")
        time.sleep(0.2)

    # ── Step 3: ARP 表 ──
    print("\n" + "=" * 70)
    print("  [3/4] 采集 ARP 表")
    print("=" * 70)
    arp_out = send_cmd(tn, "show arp all", 15, auto_more=True)
    arp_records = parse_arp(arp_out)

    # ARP 按 VLAN 汇总
    vlan_summary = {}
    for r in arp_records:
        vlan = r['vlan']
        if vlan not in vlan_summary:
            vlan_summary[vlan] = {'count': 0, 'by_interface': {}}
        vlan_summary[vlan]['count'] += 1
        iface = r['interface']
        vlan_summary[vlan]['by_interface'][iface] = vlan_summary[vlan]['by_interface'].get(iface, 0) + 1

    # ARP 按接口统计（用于流量对照）
    iface_arp_count = {}
    for r in arp_records:
        iface = r['interface']
        iface_arp_count[iface] = iface_arp_count.get(iface, 0) + 1

    print(f"  ARP 条目: {len(arp_records)}  |  VLAN 数: {len(vlan_summary)}")

    # ── Step 4: VLAN 信息 ──
    print("\n" + "=" * 70)
    print("  [4/4] 采集关键 VLAN 信息")
    print("=" * 70)
    key_vlans = sorted(vlan_summary.keys(), key=lambda v: int(v) if v.isdigit() else 9999)
    for vlan_id in key_vlans[:8]:  # 只查前8个主要 VLAN
        send_cmd(tn, f"show vlan {vlan_id}", 5)
        print(f"  VLAN {vlan_id}: {vlan_summary[vlan_id]['count']} 设备")

    # ── 退出 ──
    tn.write("exit\n".encode('ascii'))
    tn.close()
    print("\n  已断开连接")

    # ════════════ 报告 ════════════
    issues, warnings = health_check(all_interfaces)

    print("\n" + "=" * 70)
    print("  📊 上行流量健康报告")
    print("=" * 70)

    # 流量表格
    print(f"\n  {'接口':<18} {'状态':<6} {'入向Mbps':>10} {'出向Mbps':>10} {'入利用%':>7} {'出利用%':>7} {'错误':>6} {'CRC':>5} {'丢包%':>6}")
    print(f"  {'-'*85}")
    active = 0
    for name in sorted(all_interfaces.keys()):
        info = all_interfaces[name]
        s = info.get('status', '?')
        icon = '🟢' if s == 'up' else '🔴' if s == 'down' else '⚪'
        rx = info.get('rx_mbps', 0)
        tx = info.get('tx_mbps', 0)
        ru = info.get('utilization_rx', 0)
        tu = info.get('utilization_tx', 0)
        err = info.get('rx_errors', 0) + info.get('tx_errors', 0)
        crc = info.get('crc_errors', 0)
        lost = info.get('rx_lost_pct', 0) + info.get('tx_lost_pct', 0)
        arp_n = iface_arp_count.get(name, 0)

        print(f"  {name:<18} {icon} {s:<4} {rx:>10.3f} {tx:>10.3f} {ru:>7.2f} {tu:>7.2f} {err:>6} {crc:>5} {lost:>6.2f}")

        if s == 'up' and (rx > 0 or tx > 0):
            active += 1

    # VLAN ARP 分布
    print(f"\n  🔢 VLAN ARP 分布:")
    print(f"  {'VLAN':<8} {'设备数':<8} {'Top 接口'}")
    print(f"  {'-'*55}")
    for vlan in sorted(vlan_summary.keys(), key=lambda v: int(v) if v.isdigit() else 9999):
        info = vlan_summary[vlan]
        top = sorted(info['by_interface'].items(), key=lambda x: x[1], reverse=True)[:4]
        top_s = ', '.join(f'{i}({c})' for i, c in top)
        print(f"  VLAN{vlan:<4} {info['count']:<8} {top_s}")

    # 告警
    print(f"\n  🚨 告警: {len(issues) + len(warnings)} 条")
    for i in issues:
        print(f"     {i}")
    for w in warnings:
        print(f"     {w}")
    if not issues and not warnings:
        print(f"     ✅ 所有接口运行正常")

    # 汇总
    up_n = sum(1 for v in all_interfaces.values() if v.get('status') == 'up')
    down_n = sum(1 for v in all_interfaces.values() if v.get('status') == 'down')
    print(f"\n  📋 汇总:")
    print(f"     tengige端口: {len(all_interfaces)} (UP:{up_n} DOWN:{down_n} 活跃:{active})")
    print(f"     ARP条目: {len(arp_records)} | VLAN: {len(vlan_summary)}")

    # 保存文件
    raw_dir = os.path.join(OUTPUT_DIR, f"scan_{ts}")
    os.makedirs(raw_dir, exist_ok=True)

    # 原始数据
    with open(os.path.join(raw_dir, "interface_status.txt"), 'w', encoding='utf-8') as f:
        f.write(status_out)
    with open(os.path.join(raw_dir, "arp_all.txt"), 'w', encoding='utf-8') as f:
        f.write(arp_out)

    # JSON 报告
    report = {
        'scan_time': run_time,
        'target': SWITCH_HOST,
        'interfaces': all_interfaces,
        'vlan_summary': {k: {'count': v['count'], 'by_interface': v['by_interface']} for k, v in vlan_summary.items()},
        'arp_total': len(arp_records),
        'iface_arp_count': iface_arp_count,
        'issues': issues,
        'warnings': warnings,
        'raw_dir': raw_dir,
    }
    rf = os.path.join(OUTPUT_DIR, f"traffic_report_{ts}.json")
    with open(rf, 'w', encoding='utf-8') as f:
        json.dump(report, f, ensure_ascii=False, indent=2)
    print(f"\n  📄 报告: {rf}")
    print(f"  📁 原始数据: {raw_dir}/")

    print("\n" + "=" * 70)
    print("  检测完成")
    print("=" * 70)

    return report


if __name__ == '__main__':
    main()