package main

import (
	"flag"
	"fmt"
	"log"
	"net"
	"os"
	"os/signal"
	"syscall"
	"time"

	"git.kingecg.top/kingecg/vnic"
)

// 命令行参数
var (
	ipAddress  string
	packetSize int
	interval   int
	verbose    bool
)

func init() {
	// 解析命令行参数
	flag.StringVar(&ipAddress, "ip", "192.168.100.1/24", "虚拟网卡的IP地址 (CIDR格式)")
	flag.IntVar(&packetSize, "size", 64, "测试数据包大小 (字节)")
	flag.IntVar(&interval, "interval", 5, "发送数据包的间隔 (秒)")
	flag.BoolVar(&verbose, "verbose", false, "启用详细日志")
	flag.Parse()
}

func main() {
	// 配置日志格式
	log.SetFlags(log.Ldate | log.Ltime | log.Lmicroseconds)
	log.Println("虚拟网卡完整示例启动")

	// 验证IP地址格式
	if _, _, err := net.ParseCIDR(ipAddress); err != nil {
		log.Fatalf("无效的IP地址格式: %v", err)
	}

	// 创建虚拟网卡配置
	config := vnic.Config{
		IP: ipAddress,
	}

	// 创建虚拟网卡
	log.Printf("正在创建虚拟网卡 (IP: %s)...", ipAddress)
	nic, err := vnic.New(config)
	if err != nil {
		log.Fatalf("创建虚拟网卡失败: %v", err)
	}

	// 设置清理函数
	defer cleanupResources(nic)

	log.Printf("虚拟网卡创建成功:")
	log.Printf("- 接口名称: %s", nic.Name())
	log.Printf("- IP地址: %s", ipAddress)

	// 创建一个通道来监听系统信号
	sigCh := make(chan os.Signal, 1)
	signal.Notify(sigCh, syscall.SIGINT, syscall.SIGTERM)

	// 启动数据包处理器
	stopCh := make(chan struct{})
	go packetReceiver(nic, stopCh)
	go packetSender(nic, stopCh)

	// 等待中断信号
	sig := <-sigCh
	fmt.Printf("\n收到信号 %v，正在关闭...\n", sig)

	// 通知所有goroutine停止
	close(stopCh)

	// 给goroutine一些时间来完成清理
	time.Sleep(500 * time.Millisecond)
}

// cleanupResources 负责清理资源
func cleanupResources(nic *vnic.VirtualNIC) {
	log.Println("正在关闭虚拟网卡...")
	if err := nic.Close(); err != nil {
		log.Printf("关闭虚拟网卡时出错: %v", err)
	} else {
		log.Println("虚拟网卡已成功关闭")
	}
}

// packetReceiver 从虚拟网卡读取数据包
func packetReceiver(nic *vnic.VirtualNIC, stopCh <-chan struct{}) {
	log.Println("数据包接收器已启动")

	// 创建缓冲区
	buffer := make([]byte, 2048)

	for {
		// 检查是否收到停止信号
		select {
		case <-stopCh:
			log.Println("数据包接收器正在关闭")
			return
		default:
			// 设置读取超时，以便我们可以定期检查停止信号
			// 注意：在实际应用中，你可能需要使用更复杂的方法来处理这个问题
			nic.Read(buffer)

			// 解析IP数据包
			if verbose && len(buffer) > 0 {
				parseIPPacket(buffer)
			}
		}
	}
}

// packetSender 定期向虚拟网卡发送测试数据包
func packetSender(nic *vnic.VirtualNIC, stopCh <-chan struct{}) {
	log.Println("数据包发送器已启动")

	// 解析IP地址以获取网络信息
	ip, ipNet, _ := net.ParseCIDR(ipAddress)

	// 计算网络中的第二个IP地址作为目标
	// 这只是一个示例，在实际应用中，你可能需要一个有效的目标IP
	dstIP := calculateNextIP(ip, ipNet)

	// 创建测试数据包
	packet := createTestPacket(ip, dstIP, packetSize)

	// 设置定时器
	ticker := time.NewTicker(time.Duration(interval) * time.Second)
	defer ticker.Stop()

	for {
		select {
		case <-stopCh:
			log.Println("数据包发送器正在关闭")
			return
		case <-ticker.C:
			if _, err := nic.Write(packet); err != nil {
				log.Printf("发送数据包时出错: %v", err)
			} else if verbose {
				log.Printf("已发送 %d 字节的测试数据包到 %s", len(packet), dstIP)
			}
		}
	}
}

// calculateNextIP 计算网络中的下一个IP地址
func calculateNextIP(ip net.IP, ipNet *net.IPNet) net.IP {
	// 复制IP以避免修改原始IP
	nextIP := make(net.IP, len(ip))
	copy(nextIP, ip)

	// 将最后一个字节加1
	nextIP[len(nextIP)-1]++

	// 确保IP仍然在网络范围内
	if !ipNet.Contains(nextIP) {
		// 如果超出范围，使用网络的第一个可用IP
		copy(nextIP, ipNet.IP)
		nextIP[len(nextIP)-1]++
	}

	return nextIP
}

// createTestPacket 创建一个测试IP数据包
func createTestPacket(srcIP, dstIP net.IP, size int) []byte {
	// 确保最小大小
	if size < 20 {
		size = 20 // IP头部的最小大小
	}

	// 创建一个基本的IP头部
	packet := make([]byte, size)

	// IP版本(4)和头部长度(5*4=20字节)
	packet[0] = 0x45

	// 服务类型
	packet[1] = 0x00

	// 总长度 (大端字节序)
	packet[2] = byte(size >> 8)
	packet[3] = byte(size)

	// 标识
	packet[4] = 0x00
	packet[5] = 0x00

	// 标志和片偏移
	packet[6] = 0x00
	packet[7] = 0x00

	// TTL
	packet[8] = 64

	// 协议 (1 = ICMP)
	packet[9] = 0x01

	// 头部校验和 (暂时为0)
	packet[10] = 0x00
	packet[11] = 0x00

	// 源IP地址
	copy(packet[12:16], srcIP.To4())

	// 目标IP地址
	copy(packet[16:20], dstIP.To4())

	// 计算IP头部校验和
	checksum := calculateIPChecksum(packet[:20])
	packet[10] = byte(checksum >> 8)
	packet[11] = byte(checksum)

	// 如果有足够的空间，添加一些ICMP数据
	if size >= 28 {
		// ICMP Echo请求
		packet[20] = 0x08 // 类型: Echo
		packet[21] = 0x00 // 代码: 0
		packet[22] = 0x00 // 校验和 (暂时为0)
		packet[23] = 0x00

		// ICMP标识符和序列号
		packet[24] = 0x00
		packet[25] = 0x01
		packet[26] = 0x00
		packet[27] = 0x01

		// 填充剩余空间
		for i := 28; i < size; i++ {
			packet[i] = byte(i % 256)
		}
	}

	return packet
}

// calculateIPChecksum 计算IP头部的校验和
func calculateIPChecksum(header []byte) uint16 {
	var sum uint32

	// 将头部视为16位整数的序列并求和
	for i := 0; i < len(header); i += 2 {
		if i+1 < len(header) {
			sum += uint32(header[i])<<8 | uint32(header[i+1])
		} else {
			sum += uint32(header[i]) << 8
		}
	}

	// 将进位加到结果中
	for sum > 0xffff {
		sum = (sum & 0xffff) + (sum >> 16)
	}

	// 取反
	return ^uint16(sum)
}

// parseIPPacket 解析并显示IP数据包的基本信息
func parseIPPacket(packet []byte) {
	if len(packet) < 20 {
		log.Println("收到无效的IP数据包 (太短)")
		return
	}

	// 提取版本和头部长度
	version := packet[0] >> 4
	headerLen := int(packet[0]&0x0F) * 4

	if version != 4 {
		log.Printf("收到非IPv4数据包 (版本: %d)", version)
		return
	}

	// 提取总长度
	totalLen := int(packet[2])<<8 | int(packet[3])

	// 提取协议
	protocol := packet[9]

	// 提取源IP和目标IP
	srcIP := net.IPv4(packet[12], packet[13], packet[14], packet[15])
	dstIP := net.IPv4(packet[16], packet[17], packet[18], packet[19])

	// 显示基本信息
	log.Printf("收到IP数据包: 长度=%d, 协议=%d, %s -> %s",
		totalLen, protocol, srcIP, dstIP)

	// 如果是ICMP协议，显示更多信息
	if protocol == 1 && len(packet) >= headerLen+2 {
		icmpType := packet[headerLen]
		icmpCode := packet[headerLen+1]
		log.Printf("  ICMP: 类型=%d, 代码=%d", icmpType, icmpCode)
	}
}