#!/usr/bin/env python3
"""
迪普交换机 Telnet 自动采集工具
连接 172.88.88.1 采集 ARP/MAC/接口/VLAN 信息，导出 JSON
"""
import telnetlib
import re
import json
import os
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__))

# NAS 同步配置
NAS_HOST = "192.168.34.14"
NAS_USER = "lipeng"
NAS_PASS = "Sywz1905"
NAS_TARGET_DIR = "/volume1/soft/web/share"


def telnet_read_until(tn, expected, timeout=10):
    """读取直到期望字符串出现"""
    try:
        result = tn.read_until(expected.encode('ascii'), timeout=timeout)
        return result.decode('ascii', errors='ignore')
    except EOFError:
        return ""


def telnet_session(host, user, password):
    """建立telnet连接并登录"""
    tn = telnetlib.Telnet(host, 23, timeout=15)
    print(f"  连接成功: {host}")

    # 等待 Username 提示
    result = telnet_read_until(tn, 'Username:', 10)
    if not result:
        result = telnet_read_until(tn, 'login:', 5)
    print(f"  登录提示: {result[-80:].strip() if result else '(无)'}")

    tn.write((user + "\n").encode('ascii'))
    print("  -> 发送用户名")

    # 等待 Password 提示
    result = telnet_read_until(tn, 'Password:', 8)
    print(f"  密码提示: {result[-60:].strip() if result else '(无)'}")

    tn.write((password + "\n").encode('ascii'))
    print("  -> 发送密码")

    # 等待登录完成，出现命令提示符
    import time
    output = ""
    for _ in range(20):
        time.sleep(0.5)
        try:
            chunk = tn.read_very_eager().decode('ascii', errors='ignore')
            if chunk:
                output += chunk
        except Exception:
            pass
        if '>' in output or '#' in output:
            break

    print(f"  登录后输出: {output[-150:].strip() if output else '(空)'}")

    if '>' in output or '#' in output:
        print("  登录成功 [OK]")
    else:
        print("  警告: 未检测到提示符")

    return tn


def send_cmd(tn, cmd, wait=8, auto_more=False):
    """发送命令并读取完整输出，auto_more 则自动翻页"""
    import time
    tn.write((cmd + "\n").encode('ascii'))

    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
        # 处理 --More-- 分页
        if auto_more and '--More' in output and more_sent < 50:
            more_sent += 1
            tn.write(b' ')
            time.sleep(0.3)
            continue
        if output.rstrip().endswith('>') or output.rstrip().endswith('#'):
            break

    return output


def parse_arp(output):
    """解析迪普交换机的 show arp 输出"""
    # 先清理 --More-- 行和退格字符
    cleaned = re.sub(r'.*--More.*--.*\n?', '', output)
    cleaned = re.sub(r'.[\x08]+.', '', cleaned)  # 退格+覆盖

    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 'Total:' in line or 'show arp' in line.lower():
            continue
        if line.startswith('<') or line.startswith('Password:'):
            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],
                'type': parts[5] if len(parts) > 5 else 'D'
            })
    return records


def parse_mac(output):
    """解析 MAC 地址表"""
    records = []
    in_table = False
    for line in output.split('\n'):
        line = line.strip()
        if not line:
            continue
        if 'MAC Address' in line:
            in_table = True
            continue
        if not in_table:
            continue
        if line.startswith('[') or 'Total' in line:
            continue
        parts = line.split()
        if len(parts) >= 4:
            records.append({'mac': parts[1], 'vlan': parts[2], 'port': parts[3]})
    return records


def parse_interfaces(output):
    """解析接口信息"""
    ifaces = []
    current = None
    for line in output.split('\n'):
        line = line.strip()
        if not line:
            continue
        if line.startswith('[') or line.startswith('Password:'):
            continue
        if line.startswith('Interface '):
            if current:
                ifaces.append(current)
            name = line.replace('Interface ', '').strip()
            current = {'name': name, 'status': '', 'description': '',
                       'vlan_belong': '', 'pvid': '', 'speed': '',
                       'port_type': '', 'media': ''}
            continue
        if current is None:
            continue
        if 'administration state' in line:
            parts = line.split('line state is')
            a_up = 'UP' in parts[0]
            l_up = len(parts) > 1 and 'UP' in parts[1]
            current['status'] = 'UP' if (a_up and l_up) else 'DOWN'
        elif 'Description:' in line:
            current['description'] = line.split('Description:')[1].strip()
        elif 'vlan belong' in line:
            current['vlan_belong'] = line.split('vlan belong')[1].strip()
        elif line.strip().startswith('pvid'):
            current['pvid'] = line.split('pvid')[1].strip()
        elif 'Mbps-speed mode' in line:
            current['speed'] = line.strip()
        elif 'Layer2 port type' in line:
            current['port_type'] = line.split('Layer2 port type')[1].strip()
        elif 'Media type is' in line:
            current['media'] = line.split('Media type is')[1].strip()
    if current:
        ifaces.append(current)
    return ifaces


def is_monitor_ip(ip):
    """判断是否为监控网段IP"""
    monitor_segments = ['47', '58', '59'] + [str(n) for n in range(70, 78)]
    for seg in monitor_segments:
        if f'.{seg}.' in ip or ip.startswith(f'192.168.{seg}.'):
            return True
    return False


def load_previous_arp_set(json_file):
    """加载上一次扫描的 ARP 数据，返回 (ip+mac) 的集合和学习记录"""
    prev_set = set()
    prev_records = []
    if os.path.exists(json_file):
        try:
            with open(json_file, 'r', encoding='utf-8') as f:
                old_data = json.load(f)
            for row in old_data.get('records', []):
                if len(row) >= 3 and row[0] and row[1] and not row[0].startswith('>>') and not row[0].startswith('ARP') and not row[0].startswith('合计') and not row[0].startswith('其中') and row[0] != '接口 (Interface)':
                    key = row[1] + '|' + row[2]  # ip|mac
                    prev_set.add(key)
                    prev_records.append(row)
        except Exception:
            pass
    return prev_set, prev_records


def save_history(json_file):
    """将当前 arp_data.json 存档到 history/ 目录，带时间戳"""
    if not os.path.exists(json_file):
        return
    history_dir = os.path.join(os.path.dirname(json_file), 'history')
    os.makedirs(history_dir, exist_ok=True)
    timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
    history_file = os.path.join(history_dir, f'arp_{timestamp}.json')
    try:
        import shutil
        shutil.copy2(json_file, history_file)
        print(f"  历史存档: {history_file}")
    except Exception as e:
        print(f"  存档失败: {e}")


def build_json(arp_data, json_file):
    """生成网页JSON数据，包含新增标记和历史存档"""
    # 1. 先存档当前数据（如果存在）
    save_history(json_file)

    # 2. 加载上一次的数据用于对比
    prev_set, _ = load_previous_arp_set(json_file)

    # 3. 找出新增的设备
    new_set = set()
    for rec in arp_data:
        key = rec['ip'] + '|' + rec['mac']
        if key not in prev_set:
            new_set.add(key)

    by_iface = {}
    for rec in arp_data:
        iface = rec['interface']
        if iface not in by_iface:
            by_iface[iface] = []
        by_iface[iface].append(rec)

    scan_date = datetime.now().strftime('%Y-%m-%d')
    scan_time = datetime.now().strftime('%H:%M:%S')
    records_for_web = []
    records_for_web.append([f"ARP 地址解析表  |  采集日期：{scan_date} {scan_time}", '', '', '', '', '', ''])
    records_for_web.append(['接口 (Interface)', 'IP 地址', 'MAC 地址', 'VLAN ID', '类型', '序号', '备注'])

    rn = 0
    new_count = 0
    for iface in sorted(by_iface.keys()):
        records = by_iface[iface]
        # 统计本接口新增数
        iface_new = sum(1 for r in records if (r['ip'] + '|' + r['mac']) in new_set)
        header_extra = f"  🆕 {iface_new}条新增" if iface_new > 0 else ""
        records_for_web.append([f"  >>  {iface}  ({len(records)} 条记录){header_extra}", '', '', '', '', '', ''])
        for i, rec in enumerate(records, 1):
            ip = rec['ip']
            key = ip + '|' + rec['mac']
            is_new = key in new_set
            remark = '监控' if is_monitor_ip(ip) else ('交换机' if ip.startswith('172.') else '')
            if is_new:
                remark = '🆕新增' + (' ' + remark if remark else '')
                new_count += 1
            records_for_web.append([iface, ip, rec['mac'], rec['vlan'], rec['type'], str(i), remark, '1' if is_new else '0'])
            rn += 1

    records_for_web.append([f"合计：{rn} 条记录，{len(by_iface)} 个接口", '', '', '', '', '', ''])
    if new_count > 0:
        records_for_web.append([f"本次新增：{new_count} 条记录", '', '', '', '', '', ''])
    where_monitor = sum(1 for r in arp_data if is_monitor_ip(r['ip']))
    if where_monitor > 0:
        records_for_web.append([f"其中监控网段：{where_monitor} 条记录", '', '', '', '', '', ''])

    # 接口统计
    stats = [["接口统计汇总", '', '', '', '']]
    stats.append(['接口', '记录数', '涉及VLAN', 'IP段示例', '备注'])
    for iface in sorted(by_iface.keys()):
        records = by_iface[iface]
        vlans_set = sorted(set(r['vlan'] for r in records))
        ip_sample = records[0]['ip'] if records else ''
        if len(records) > 1:
            ip_sample += f' ~ {records[-1]["ip"]}'
        monitor_count = sum(1 for r in records if is_monitor_ip(r['ip']))
        switch_count = sum(1 for r in records if r['ip'].startswith('172.'))
        iface_new_count = sum(1 for r in records if (r['ip'] + '|' + r['mac']) in new_set)
        remark_parts = []
        if monitor_count > 0:
            remark_parts.append(f'监控({monitor_count})')
        if switch_count > 0:
            remark_parts.append(f'交换机({switch_count})')
        if iface_new_count > 0:
            remark_parts.append(f'新增({iface_new_count})')
        stats.append([iface, str(len(records)), ','.join(vlans_set), ip_sample, ' '.join(remark_parts)])
    stats.append([f"合计：{len(by_iface)} 个接口", '', f'共 {rn} 条记录', '', ''])

    # 生成历史文件列表
    history_files = []
    history_dir = os.path.join(os.path.dirname(json_file), 'history')
    if os.path.exists(history_dir):
        for fname in sorted(os.listdir(history_dir), reverse=True):
            if fname.endswith('.json'):
                history_files.append(fname)

    with open(json_file, 'w', encoding='utf-8') as f:
        json.dump({
            'records': records_for_web,
            'stats': stats,
            'new_count': new_count,
            'scan_time': f'{scan_date} {scan_time}',
            'history': history_files
        }, f, ensure_ascii=False, indent=2)


def main():
    print("=" * 60)
    print("  迪普交换机 Telnet 数据采集")
    print(f"  目标: {SWITCH_HOST}")
    print(f"  时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
    print("=" * 60)

    tn = telnet_session(SWITCH_HOST, SWITCH_USER, SWITCH_PASS)

    # 尝试关闭分页显示
    send_cmd(tn, "terminal length 0", 3)
    send_cmd(tn, "no paging", 3)

    # 采集ARP
    print("\n采集 ARP 表...")
    output = send_cmd(tn, "show arp all", 15, auto_more=True)
    with open(os.path.join(OUTPUT_DIR, 'debug_arp.txt'), 'w', encoding='utf-8') as f:
        f.write(output)
    print(f"  ARP 原始输出长度: {len(output)}")
    arp_data = parse_arp(output)
    print(f"  ARP: {len(arp_data)} 条")

    if len(arp_data) == 0:
        print("  [WARN] ARP 为空, debug_arp.txt 内容:")
        print("  " + output[:500])

    # 采集接口
    print("采集接口信息...")
    output = send_cmd(tn, "show interface", 45, auto_more=True)
    with open(os.path.join(OUTPUT_DIR, 'showinterface.txt'), 'w', encoding='utf-8') as f:
        f.write(output)
    ifaces = parse_interfaces(output)
    print(f"  接口: {len(ifaces)} 个")

    # 采集MAC表
    print("采集 MAC 地址表...")
    output = send_cmd(tn, "show mac-address-table dynamic", 15)
    with open(os.path.join(OUTPUT_DIR, 'showmac.txt'), 'w', encoding='utf-8') as f:
        f.write(output)
    mac_data = parse_mac(output)
    print(f"  MAC: {len(mac_data)} 条")

    # 采集VLAN
    print("采集 VLAN 信息...")
    output = send_cmd(tn, "show vlan", 10)
    with open(os.path.join(OUTPUT_DIR, 'showvlan.txt'), 'w', encoding='utf-8') as f:
        f.write(output)
    print(f"  VLAN: 已获取")

    # 退出
    tn.write("exit\n".encode('ascii'))
    tn.close()

    # 生成 JSON（含历史存档和新增对比）
    json_file = os.path.join(OUTPUT_DIR, 'arp_data.json')
    build_json(arp_data, json_file)
    print(f"\n[OK] JSON: {json_file}")

    is_mon = sum(1 for r in arp_data if is_monitor_ip(r['ip']))
    iface_count = len(set(r['interface'] for r in arp_data))

    # 重新加载 JSON 获取新增统计
    try:
        with open(json_file, 'r', encoding='utf-8') as f:
            saved = json.load(f)
        new_cnt = saved.get('new_count', 0)
    except Exception:
        new_cnt = 0

    print(f"  总记录: {len(arp_data)} 条, 监控: {is_mon} 条, 接口: {iface_count} 个")
    if new_cnt > 0:
        print(f"  [NEW] 本次新增: {new_cnt} 条")

    # 同步到群晖 NAS
    print("\n同步到群晖 NAS...")
    sync_to_nas()
    print("\nDone!")


def sync_to_nas():
    """通过 SSH 将本地文件同步到 NAS"""
    try:
        import paramiko
        import base64
        import glob
    except ImportError:
        print("  [WARN] 缺少 paramiko 库，跳过同步。pip install paramiko")
        return

    files_to_sync = ['arp_data.json', 'arp.html', 'collect_arp.py']
    try:
        ssh = paramiko.SSHClient()
        ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
        ssh.connect(NAS_HOST, username=NAS_USER, password=NAS_PASS, timeout=15)

        for fname in files_to_sync:
            local_path = os.path.join(OUTPUT_DIR, fname)
            if not os.path.exists(local_path):
                print(f"  [SKIP] {fname} 不存在")
                continue
            try:
                with open(local_path, 'rb') as f:
                    data = base64.b64encode(f.read()).decode('ascii')
                chan = ssh.get_transport().open_session()
                chan.exec_command('echo "' + data + '" | base64 -d > ' + NAS_TARGET_DIR + '/' + fname)
                chan.recv_exit_status()
                print(f"  [OK] {fname}")
            except Exception as e:
                print(f"  [FAIL] {fname}: {e}")

        # 同步历史文件
        history_dir = os.path.join(OUTPUT_DIR, 'history')
        if os.path.exists(history_dir):
            chan = ssh.get_transport().open_session()
            chan.exec_command('mkdir -p ' + NAS_TARGET_DIR + '/history')
            chan.recv_exit_status()
            for fname in os.listdir(history_dir):
                if fname.endswith('.json'):
                    local_path = os.path.join(history_dir, fname)
                    try:
                        with open(local_path, 'rb') as f:
                            data = base64.b64encode(f.read()).decode('ascii')
                        chan = ssh.get_transport().open_session()
                        chan.exec_command('echo "' + data + '" | base64 -d > ' + NAS_TARGET_DIR + '/history/' + fname)
                        chan.recv_exit_status()
                    except Exception as e:
                        print(f"  [FAIL] history/{fname}: {e}")
            print(f"  [OK] 历史文件已同步")

        # 重启 NAS 上的 Web 服务（通过任务计划）
        try:
            import urllib.request
            import ssl
            ctx = ssl.create_default_context()
            ctx.check_hostname = False
            ctx.verify_mode = ssl.CERT_NONE
            login_url = f"https://{NAS_HOST}:5001/webapi/auth.cgi?api=SYNO.API.Auth&version=6&method=login&account={NAS_USER}&passwd={NAS_PASS}&session=FileStation&format=sid"
            r = urllib.request.urlopen(login_url, context=ctx)
            sid = json.loads(r.read())['data']['sid']
            # 运行开机任务来重启服务
            restart_url = f"https://{NAS_HOST}:5001/webapi/entry.cgi?api=SYNO.Core.TaskScheduler&version=1&method=run&id=1001&_sid={sid}"
            urllib.request.urlopen(restart_url, context=ctx)
            print("  [OK] Web 服务已重启")
        except Exception as e:
            print(f"  [WARN] 重启失败（可手动重启）: {e}")

        ssh.close()
    except Exception as e:
        print(f"  [FAIL] 连接 NAS 失败: {e}")
        print("  提示: 请确认 NAS 开机且 SSH 服务已启用")


if __name__ == '__main__':
    main()