#!/usr/bin/env python3
"""
纯Python 抓包工具 (在Pi上运行, 无需 tcpdump)

用 AF_PACKET 原始套接字抓所有以太网包, 过滤出指定的 UDP 流, 写成标准 pcap 格式
用法: sudo python3 _raw_capture.py <输出pcap> <时长秒> [--filter-host IP] [--filter-port PORT]
"""
import socket, struct, sys, time, argparse

ETH_P_ALL = 0x0003

PCAP_MAGIC = 0xa1b2c3d4
LINKTYPE_ETHERNET = 1


def write_pcap_global_header(f):
    f.write(struct.pack('<IHHIIII',
        PCAP_MAGIC, 2, 4, 0, 0, 65535, LINKTYPE_ETHERNET))


def write_pcap_packet(f, ts, pkt):
    ts_sec = int(ts)
    ts_usec = int((ts - ts_sec) * 1e6)
    f.write(struct.pack('<IIII', ts_sec, ts_usec, len(pkt), len(pkt)))
    f.write(pkt)


def parse_packet(pkt):
    """返回 (src_ip, dst_ip, proto, src_port, dst_port) 或 None"""
    if len(pkt) < 14: return None
    eth_type = struct.unpack('>H', pkt[12:14])[0]
    if eth_type != 0x0800: return None  # 只要IPv4

    # IP 头
    if len(pkt) < 14 + 20: return None
    ihl = (pkt[14] & 0x0f) * 4
    proto = pkt[14 + 9]
    if proto != 17: return None  # 只要 UDP
    src_ip = '.'.join(str(b) for b in pkt[14+12:14+16])
    dst_ip = '.'.join(str(b) for b in pkt[14+16:14+20])

    # UDP 头
    udp_off = 14 + ihl
    if len(pkt) < udp_off + 8: return None
    src_port, dst_port = struct.unpack('>HH', pkt[udp_off:udp_off+4])
    return src_ip, dst_ip, src_port, dst_port


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('output', help='输出pcap路径')
    parser.add_argument('duration', type=float, help='时长(秒)')
    parser.add_argument('--filter-host', help='只抓涉及此IP的')
    parser.add_argument('--filter-port', type=int, help='只抓涉及此端口的')
    parser.add_argument('--iface', default=None, help='绑定接口名(如eth0)')
    args = parser.parse_args()

    sock = socket.socket(socket.AF_PACKET, socket.SOCK_RAW, socket.htons(ETH_P_ALL))
    if args.iface:
        sock.bind((args.iface, 0))
        print(f"  抓包接口: {args.iface}", flush=True)
    else:
        print(f"  抓包接口: ALL", flush=True)
    sock.settimeout(0.5)

    out = open(args.output, 'wb')
    write_pcap_global_header(out)

    end = time.time() + args.duration
    n_total = 0
    n_match = 0
    n_ipv4 = 0
    n_udp = 0
    from collections import Counter
    flow_counter = Counter()  # 统计 UDP 流向

    print(f"  抓包 {args.duration}s...", flush=True)
    while time.time() < end:
        try:
            pkt, _ = sock.recvfrom(65535)
        except socket.timeout:
            continue
        n_total += 1
        # 粗略判断: 14字节以太网头 + IP头
        if len(pkt) >= 14:
            eth_type = struct.unpack('>H', pkt[12:14])[0]
            if eth_type == 0x0800:
                n_ipv4 += 1
                if len(pkt) >= 14+20 and pkt[14+9] == 17:
                    n_udp += 1
        info = parse_packet(pkt)
        if info is None: continue
        src_ip, dst_ip, sp, dp = info
        flow_counter[(src_ip, sp, dst_ip, dp)] += 1

        if args.filter_host:
            if args.filter_host not in (src_ip, dst_ip): continue
        if args.filter_port:
            if args.filter_port not in (sp, dp): continue
        # 每命中第一个新流就打印一下, 方便诊断
        flow = (src_ip, sp, dst_ip, dp)
        if not hasattr(main, '_seen'):
            main._seen = set()
        if flow not in main._seen:
            main._seen.add(flow)
            print(f"  新流: {src_ip}:{sp} → {dst_ip}:{dp}", flush=True)

        write_pcap_packet(out, time.time(), pkt)
        n_match += 1

    out.close()
    sock.close()
    print(f"  总包: {n_total}, IPv4: {n_ipv4}, UDP: {n_udp}, 命中: {n_match}", flush=True)
    print(f"  Top UDP 流 (前20):", flush=True)
    for flow, cnt in flow_counter.most_common(20):
        sip, spt, dip, dpt = flow
        print(f"    {sip}:{spt} → {dip}:{dpt}  ({cnt})", flush=True)


if __name__ == '__main__':
    main()