//go:build windows && !arm64

package tcm

import (
	"b612.me/bcap"
	"b612.me/bcap/libpcap"
	"b612.me/starlog"
	"context"
	"encoding/hex"
	"fmt"
	"github.com/google/gopacket"
	"github.com/google/gopacket/layers"
	"github.com/google/gopacket/pcap"
	"github.com/google/gopacket/pcapgo"
	"net"
	"os"
	"os/signal"
	"strconv"
	"strings"
	"sync"
	"syscall"
	"time"
)

type Libpcap struct {
	count uint64
	// 按连接数最后一个历史报文
	cap *bcap.Packets
	// 监控目标ip地址列表
	target []string
	// 将军，下达命令吧！
	// 监控命令列表
	targetCmd   []string
	targetAsHex bool
	//保存为pcap格式的数据文件
	saveFile string
	// 写缓存
	packetCache chan gopacket.Packet
	//交互模式
	interactive bool
	//展示所有报文信息，包括未在追踪列表的
	showAll bool
	//展示payload报文（utf-8直接输出）
	showPayload bool
	//payload展示为hex
	showAsHex bool
	//payload允许展示的最大字符
	maxShowPayloadSize int
	//不展示任何信息
	noShowMode bool
	// 触发封禁词后，使用RST重置链路
	useRST bool
	// RST模式，target=目标端单向RST，reverse=对端反向RST，both=双向RST
	rstMode    string
	printColor []*starlog.Color
	logCache   chan loged
	ctx        context.Context
	fn         context.CancelFunc
	bpf        string
	eth        string
	host       string
	handle     *libpcap.NetCatch
	blockMap   sync.Map
}

type loged struct {
	str        string
	stateLevel int
	logLevel   string
}

func (t *Libpcap) Run() error {
	var err error
	stopSignal := make(chan os.Signal)
	starlog.Noticef("Starting libpcap capture\n")
	if t.bpf == "" {
		starlog.Errorln("请输入一个抓包模式")
		os.Exit(3)
	}
	if t.host == "" && t.eth == "" {
		starlog.Errorln("请输入eth网卡名或host名")
		ifs, err := libpcap.FindAllDevs()
		if err == nil {
			fmt.Println("网卡名如下:\n----------\n")
			for k, v := range ifs {
				var ips []string
				for _, vv := range v.Addresses {
					ips = append(ips, vv.IP.String())
				}
				fmt.Printf("%d.\t%v\t%s\t%v\n", k+1, v.Name, strings.Join(ips, " , "), v.Description)
			}
			fmt.Println()
		}
		return fmt.Errorf("请输入eth网卡名或host名")
	}
	if t.host == "" {
		t.handle, err = libpcap.NewCatchEth(t.eth, t.bpf)
	} else {
		t.handle, err = libpcap.NewCatch(t.host, t.bpf)
	}
	if err != nil {
		starlog.Errorln("failed to listen:", err)
		return err
	}
	t.handle.SetRecall(t.handlePacket)
	if t.targetAsHex {
		for k, v := range t.targetCmd {
			tmp, _ := hex.DecodeString(v)
			t.targetCmd[k] = string(tmp)
		}
	}
	go func() {
		if err = t.handle.Run(); err != nil {
			starlog.Errorln(err)
			stopSignal <- syscall.SIGKILL
		}
	}()
	if t.saveFile != "" {
		f, err := os.Create(t.saveFile)
		if err != nil {
			starlog.Errorln("创建写入文件失败", err)
			os.Exit(4)
		}
		defer f.Close()
		go t.pcapWriter(t.ctx, f)
	}
	signal.Notify(stopSignal, syscall.SIGINT, syscall.SIGTERM, syscall.SIGKILL)
	go t.logPrint(t.ctx)
	select {
	//todo need nf.Stop()
	case <-stopSignal:
		starlog.Warningf("Received stop signal\n")

	case <-t.ctx.Done():
		starlog.Infoln("TCPMonitor Task Finished")

	}
	return nil

}

func NewLibpcap() *Libpcap {
	var nf = new(Libpcap)
	nf.packetCache = make(chan gopacket.Packet, 2048)
	nf.logCache = make(chan loged, 2048)
	nf.printColor = []*starlog.Color{
		starlog.NewColor(starlog.FgWhite),                 //0=unknown
		starlog.NewColor(starlog.FgHiCyan),                //1=tcp_connect_1,
		starlog.NewColor(starlog.FgHiCyan),                //2=tcp_connect_2,
		starlog.NewColor(starlog.FgHiCyan),                //3=tcp_connect_3,
		starlog.NewColor(starlog.FgCyan),                  //4=tcp_bye_bye
		starlog.NewColor(starlog.FgCyan),                  //5=tcp_bye_bye
		starlog.NewColor(starlog.FgCyan),                  //6=tcp_bye_bye
		starlog.NewColor(starlog.FgCyan),                  //7=tcp_bye_bye
		starlog.NewColor(starlog.FgCyan),                  //8=tcp_bye_bye
		starlog.NewColor(starlog.FgGreen),                 //9=tcp_ok
		starlog.NewColor(starlog.BgRed, starlog.FgYellow), //10=tcp_retrans
		starlog.NewColor(starlog.FgHiMagenta),             //ece
		starlog.NewColor(starlog.FgHiMagenta),             //cwr
		starlog.NewColor(starlog.FgRed),                   //rst
		starlog.NewColor(starlog.FgHiGreen),               //keepalive
		starlog.NewColor(starlog.FgWhite),                 //0=unknown
		starlog.NewColor(starlog.FgWhite),                 //0=unknown
		starlog.NewColor(starlog.FgWhite),                 //0=unknown
		starlog.NewColor(starlog.FgWhite),                 //0=unknown
		starlog.NewColor(starlog.FgWhite),                 //0=unknown
		starlog.NewColor(starlog.FgWhite),                 //20=udp
		starlog.NewColor(starlog.FgWhite),                 //0=unknown
	}
	nf.cap = bcap.NewPackets()
	nf.ctx, nf.fn = context.WithCancel(context.Background())
	return nf
}

func (n *Libpcap) logPrint(ctx context.Context) {
	for {
		select {
		case <-ctx.Done():
			return
		case l := <-n.logCache:
			if n.noShowMode {
				fmt.Printf("已捕获报文数量:%v个\r", n.count)
				continue
			}
			switch l.logLevel {
			case "info":
				starlog.Info(l.str)
			case "notice":
				starlog.Notice(l.str)
			case "warning":
				starlog.Warning(l.str)
			case "error":
				starlog.Error(l.str)
			case "critical":
				starlog.Critical(l.str)
			case "debug":
				starlog.Debug(l.str)
			case "payload":
				fmt.Println(l.str)
			default:
				n.printColor[l.stateLevel].Print(l.str)
			}
		}
	}
}

func (n *Libpcap) ipInRange(ip string) bool {
	if len(n.target) == 0 {
		return false
	}
	for _, v := range n.target {
		if v == ip {
			return true
		}
	}
	return false
}

func (n *Libpcap) strInRange(str string) bool {
	if len(n.targetCmd) == 0 {
		return false
	}
	for _, v := range n.targetCmd {
		if v == "" {
			continue
		}
		if strings.Contains(str, v) {
			return true
		}
	}
	return false
}

func (n *Libpcap) handlePacket(p gopacket.Packet) {
	n.count++
	if n.saveFile != "" {
		n.packetCache <- p
	}
	var layer gopacket.Layer
	for _, layerType := range []gopacket.LayerType{
		layers.LayerTypeTCP, layers.LayerTypeUDP, layers.LayerTypeICMPv4, layers.LayerTypeICMPv6,
		layers.LayerTypeIPv4, layers.LayerTypeIPv6, layers.LayerTypeARP,
	} {
		layer = p.Layer(layerType)
		if layer == nil {
			continue
		}
		break
	}
	if layer == nil {
		n.logCache <- loged{
			str:        "无法定义layer类型\n",
			stateLevel: 0,
			logLevel:   "error",
		}
		return
	}
	info, err := n.cap.ParsePacket(p)
	if err != nil {
		starlog.Errorln(err)
		return
	}
	var dec string
	switch info.StateDescript() {
	case 0:
		dec = "未知状态"
	case 1:
		dec = "SYN tcp建立一次握手"
	case 2:
		dec = "SYN,ACK tcp建立二次握手"
	case 3:
		dec = "ACK tcp建立三次握手"
	case 4:
		dec = "FIN tcp断开一次挥手"
	case 5:
		dec = "ACK tcp断开二次挥手"
	case 6:
		dec = "FIN,ACK tcp断开二次三次挥手"
	case 7:
		dec = "FIN tcp断开三次挥手"
	case 8:
		dec = "ACK tcp断开四次挥手"
	case 9:
		dec = "tcp报文"
	case 10:
		dec = "TCP重传"
	case 11:
		dec = "TCP ece"
	case 12:
		dec = "TCP cwr"
	case 13:
		dec = "TCP RST重置"
	case 14:
		dec = "TCP Keepalive"
	}
	if !n.showAll && !n.ipInRange(info.SrcIP) && !n.ipInRange(info.DstIP) {
		return
	}
	n.logCache <- loged{
		str: fmt.Sprintf("%s %v:%v -> %v:%v %s seq=%v ack=%v win=%v len=%v\n", time.Now().Format("2006-01-02 15:04:05.000000"), info.SrcIP, info.SrcPort,
			info.DstIP, info.DstPort, dec, info.TcpSeq(), info.TcpAck(), info.TcpWindow(), info.TcpPayloads()),
		stateLevel: int(info.StateDescript()),
		logLevel:   "",
	}

	if n.showPayload {
		str := string(layer.LayerPayload())
		if n.maxShowPayloadSize > 0 {
			if len(str) > n.maxShowPayloadSize {
				str = str[:n.maxShowPayloadSize]
			}
		}
		if n.showAsHex {
			str = hex.EncodeToString([]byte(str))
		}
		n.logCache <- loged{
			str:        str,
			stateLevel: int(info.StateDescript()),
			logLevel:   "payload",
		}
	}
	if n.ipInRange(info.SrcIP) || n.ipInRange(info.DstIP) {
		_, ok := n.blockMap.Load(info.Key)
		if n.useRST && (ok || n.strInRange(string(layer.LayerPayload()))) {
			n.blockMap.Store(info.Key, true)
			starlog.Warningf("触发封禁关键词 RST重置\n")
			RealSendRST(n.handle.Handle, info, n.rstMode, 3)
		}
	}
}

func (n *Libpcap) pcapWriter(stopCtx context.Context, fp *os.File) error {
	w := pcapgo.NewWriter(fp)
	err := w.WriteFileHeader(65535, layers.LinkTypeEthernet)
	if err != nil {
		return err
	}
	for {
		select {
		case <-stopCtx.Done():
			return nil
		case p := <-n.packetCache:
			w.WritePacket(p.Metadata().CaptureInfo, p.Data())
		}
	}
}

func RealSendRST(p *pcap.Handle, info bcap.PacketInfo, target string, number int) {
	for i := 0; i < number; i++ {
		if target == "both" || target == "target" {
			SendRST(p, info.SrcMac, info.DstMac, info.SrcIP, info.SrcPort, info.DstIP, info.DstPort, info.TcpSeq()+(uint32(i)*uint32(info.TcpWindow())))
			//SendRST(p, info.DstMac, info.SrcMac, info.SrcIP, info.SrcPort, info.DstIP, info.DstPort, info.TcpSeq()+(uint32(i)*uint32(info.TcpWindow())))
		}
		if target == "both" || target == "reverse" {
			SendRST(p, info.DstMac, info.SrcMac, info.DstIP, info.DstPort, info.SrcIP, info.SrcPort, info.TcpAck()+(uint32(i)*uint32(info.TcpWindow())))
			//SendRST(p, info.SrcMac, info.DstMac, info.DstIP, info.DstPort, info.SrcIP, info.SrcPort, info.TcpAck()+(uint32(i)*uint32(info.TcpWindow())))
		}
	}
}

func SendRST(p *pcap.Handle, srcMac, dstMac []byte, srcIP, srcPort, dstIP, dstPort string, seq uint32) error {
	if net.ParseIP(dstIP).To4() != nil {
		return sendIPv4(p, srcMac, dstMac, srcIP, srcPort, dstIP, dstPort, seq, false)
	}
	return sendIPv6(p, srcMac, dstMac, srcIP, srcPort, dstIP, dstPort, seq, false)
}

func SendSYN(p *pcap.Handle, srcMac, dstMac []byte, srcIP, srcPort, dstIP, dstPort string, seq uint32) error {
	if net.ParseIP(dstIP).To4() != nil {
		return sendIPv4(p, srcMac, dstMac, srcIP, srcPort, dstIP, dstPort, seq, true)
	}
	return sendIPv6(p, srcMac, dstMac, srcIP, srcPort, dstIP, dstPort, seq, true)
}

func sendIPv4(p *pcap.Handle, srcMac, dstMac []byte, srcIP, srcPort, dstIP, dstPort string, seq uint32, isSyn bool) error {

	dstNetIP := net.ParseIP(dstIP)
	eth := layers.Ethernet{
		SrcMAC:       srcMac,
		DstMAC:       dstMac,
		EthernetType: layers.EthernetTypeIPv4,
	}
	iPv4 := layers.IPv4{
		SrcIP:    net.ParseIP(srcIP),
		DstIP:    dstNetIP,
		Version:  4,
		TTL:      64,
		Protocol: layers.IPProtocolTCP,
	}
	sPort, err := strconv.Atoi(srcPort)
	if err != nil {
		return err
	}
	dPort, err := strconv.Atoi(dstPort)
	if err != nil {
		return err
	}
	tcp := layers.TCP{
		SrcPort: layers.TCPPort(sPort),
		DstPort: layers.TCPPort(dPort),
		Seq:     seq,
		RST:     !isSyn,
		SYN:     isSyn,
	}

	if err = tcp.SetNetworkLayerForChecksum(&iPv4); err != nil {
		return err
	}
	buffer := gopacket.NewSerializeBuffer()
	options := gopacket.SerializeOptions{
		FixLengths:       true,
		ComputeChecksums: true,
	}
	if srcMac == nil && dstMac == nil {
		if err = gopacket.SerializeLayers(buffer, options, &iPv4, &tcp); err != nil {
			return err
		}
	} else {
		if err = gopacket.SerializeLayers(buffer, options, &eth, &iPv4, &tcp); err != nil {
			return err
		}
	}
	return p.WritePacketData(buffer.Bytes())
}

func sendIPv6(p *pcap.Handle, srcMac, dstMac []byte, srcIP, srcPort, dstIP, dstPort string, seq uint32, isSyn bool) error {
	dstNetIP := net.ParseIP(dstIP)
	eth := layers.Ethernet{
		SrcMAC:       srcMac,
		DstMAC:       dstMac,
		EthernetType: layers.EthernetTypeIPv6,
	}
	iPv6 := layers.IPv6{
		SrcIP:      net.ParseIP(srcIP),
		DstIP:      dstNetIP,
		Version:    6,
		NextHeader: layers.IPProtocolTCP,
		HopLimit:   64,
	}
	sPort, err := strconv.Atoi(srcPort)
	if err != nil {
		return err
	}
	dPort, err := strconv.Atoi(dstPort)
	if err != nil {
		return err
	}
	tcp := layers.TCP{
		SrcPort: layers.TCPPort(sPort),
		DstPort: layers.TCPPort(dPort),
		Seq:     seq,
		RST:     !isSyn,
		SYN:     isSyn,
	}
	if err := tcp.SetNetworkLayerForChecksum(&iPv6); err != nil {
		return err
	}
	buffer := gopacket.NewSerializeBuffer()
	options := gopacket.SerializeOptions{
		FixLengths:       true,
		ComputeChecksums: true,
	}
	if srcMac == nil && dstMac == nil {
		if err = gopacket.SerializeLayers(buffer, options, &iPv6, &tcp); err != nil {
			return err
		}
	} else {
		if err = gopacket.SerializeLayers(buffer, options, &eth, &iPv6, &tcp); err != nil {
			return err
		}
	}
	return p.WritePacketData(buffer.Bytes())
}