//go:build windows && (amd64 || 386)

package tcpkill

import (
	"b612.me/bcap"
	"b612.me/bcap/libpcap"
	"b612.me/starlog"
	"context"
	"fmt"
	"github.com/google/gopacket"
	"github.com/google/gopacket/layers"
	"github.com/google/gopacket/pcap"
	netm "github.com/shirou/gopsutil/v4/net"
	"net"
	"os"
	"os/signal"
	"strconv"
	"strings"
	"sync"
	"syscall"
	"time"
)

type TCPKill struct {
	NFQNums      int
	AutoIptables bool
	BPF          string
	Eth          string
	Host         string
	SrcIP        string
	SrcPort      int
	DstIP        string
	DstPort      int
	KillType     string
	RstNumbers   int
	WaitMode     bool
	Status       string
	matchConns   sync.Map
	matchCount   uint
	sync.Mutex
	cache    []string
	cap      *bcap.Packets
	ctx      context.Context
	stopFn   context.CancelFunc
	requests chan *handler
	pc       *libpcap.NetCatch
	macCache map[string]net.HardwareAddr
}

type handler struct {
}

func (t *TCPKill) presetWindows() {
	t.macCache = make(map[string]net.HardwareAddr)
	if t.BPF == "" {
		if t.DstIP != "" {
			t.BPF = fmt.Sprintf("tcp and host %s", t.DstIP)
		} else if t.SrcIP != "" {
			t.BPF = fmt.Sprintf("tcp and host %s", t.SrcIP)
		} else if t.SrcPort != 0 {
			t.BPF = fmt.Sprintf("tcp and port %d", t.SrcPort)
		} else if t.DstPort != 0 {
			t.BPF = fmt.Sprintf("tcp and port %d", t.DstPort)
		} else {
			t.BPF = "tcp"
		}
	}
	if t.Eth == "" && t.Host == "" {
		allDevice, err := pcap.FindAllDevs()
		if err != nil {
			starlog.Errorf("get pcap devices failed:%v\n", err)
			return
		}
		if t.SrcIP == "" {
			defer func() {
				t.SrcIP = ""
			}()
			ip := t.DstIP
			if ip == "" {
				ip = "223.6.6.6"
			}
			if strings.Contains(ip, ":") {
				ip = "[" + ip + "]"
			}
			l, err := net.Dial("udp", ip+":53")
			if err == nil {
				t.SrcIP = l.LocalAddr().(*net.UDPAddr).IP.String()
				starlog.Infof("No eth or ip detected,use %s as source ip\n", t.SrcIP)
				l.Close()
			} else {
				starlog.Errorf("Failed to get source ip:%v\n", err)
			}
		}
		for _, v := range allDevice {
			if t.SrcIP == "127.0.0.1" && v.Name == "\\Device\\NPF_Loopback" {
				t.Eth = v.Name
				return
			}
			for _, addr := range v.Addresses {
				if addr.IP.String() == t.SrcIP {
					t.Eth = v.Name
					return
				}
			}
		}

	}
}

func (t *TCPKill) Run() error {
	var err error
	if err = t.PreRun(); err != nil {
		return err
	}
	t.presetWindows()
	stopSignal := make(chan os.Signal)
	starlog.Noticef("Starting capture\n")
	signal.Notify(stopSignal, syscall.SIGINT, syscall.SIGTERM, syscall.SIGKILL)
	if t.Eth != "" {
		t.pc, err = libpcap.NewCatchEth(t.Eth, t.BPF)
		if err != nil {
			starlog.Errorf("Failed to create pcap handle:%v\n", err)
			return err
		}
	}
	if t.Host != "" {
		t.pc, err = libpcap.NewCatch(t.Host, t.BPF)
		if err != nil {
			starlog.Errorf("Failed to create pcap handle:%v\n", err)
			return err
		}
	}
	if t.pc == nil {
		starlog.Errorf("No pcap handle\n")
		return fmt.Errorf("No pcap handle")
	}
	t.pc.SetRecall(t.handlePacket)
	go func() {
		if err := t.pc.Run(); err != nil {
			starlog.Errorln(err)
			stopSignal <- syscall.SIGKILL
		}
	}()
	go t.SynSent()
	select {
	//todo need nf.Stop()
	case <-stopSignal:
		starlog.Warningf("Received stop signal\n")

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

	}
	return nil

}

func (t *TCPKill) SynSent() {
	for {
		time.Sleep(time.Millisecond * 2500)
		realConns, err := netm.Connections("tcp")
		if err != nil {
			continue
		}
		t.matchConns.Range(func(key, value any) bool {
			t.matchConns.Delete(key)
			t.matchCount--
			return true
		})
		for _, conn := range realConns {
			if t.Match(conn) {
				t.matchConns.Store(key(conn), conn)
				t.matchCount++
			}
		}
		if t.matchCount == 0 && !t.WaitMode {
			starlog.Warningln("No matched connection anymore")
			t.stopFn()
			return
		}
		t.matchConns.Range(func(k, v any) bool {
			conn := v.(netm.ConnectionStat)
			if t.macCache[conn.Laddr.IP] == nil || t.macCache[conn.Raddr.IP] == nil {
				target := conn.Raddr.IP + ":" + strconv.Itoa(int(conn.Raddr.Port))
				if strings.Contains(conn.Raddr.IP, ":") {
					target = "[" + conn.Raddr.IP + "]:" + strconv.Itoa(int(conn.Raddr.Port))
				}
				starlog.Warningf("No mac address for %v or %v,try send syn to %v\n", conn.Laddr.IP, conn.Raddr.IP, target)
				tcpConn, err := net.DialTimeout("tcp", target, time.Millisecond*800)
				if err == nil {
					tcpConn.(*net.TCPConn).SetLinger(0)
					tcpConn.Close()
				}
				time.Sleep(time.Millisecond * 600)
			}
			t.Lock()
			if err := SendSYN(t.pc.Handle, t.macCache[conn.Laddr.IP], t.macCache[conn.Raddr.IP], conn.Laddr.IP, strconv.Itoa(int(conn.Laddr.Port)), conn.Raddr.IP, strconv.Itoa(int(conn.Raddr.Port)), 1127); err != nil {
				starlog.Errorf("Failed to send SYN:%v\n", err)
			}
			if err := SendSYN(t.pc.Handle, t.macCache[conn.Raddr.IP], t.macCache[conn.Laddr.IP], conn.Raddr.IP, strconv.Itoa(int(conn.Raddr.Port)), conn.Laddr.IP, strconv.Itoa(int(conn.Laddr.Port)), 1127); err != nil {
				starlog.Errorf("Failed to send SYN:%v\n", err)
			}
			t.Unlock()
			starlog.Infof("Send SYN to %v <==> %v\n", conn.Laddr.String(), conn.Raddr.String())
			return true
		})
	}
}

func (t *TCPKill) handlePacket(p gopacket.Packet) {
	t.Lock()
	info, err := t.cap.ParsePacket(p)
	t.Unlock()
	if err != nil {
		return
	}
	tcpLayer := p.Layer(layers.LayerTypeTCP)
	if tcpLayer == nil {
		return
	}
	tcp := tcpLayer.(*layers.TCP)
	if tcp.SYN && !tcp.ACK {
		return
	}
	t.Lock()
	//fmt.Println(info.SrcIP, hex.EncodeToString(info.SrcMac), info.DstIP, hex.EncodeToString(info.DstMac))
	t.macCache[info.SrcIP] = info.SrcMac
	t.macCache[info.DstIP] = info.DstMac
	t.Unlock()
	if tcp.RST {
		starlog.Warningf("RST packet:%v <==> %v\n", info.SrcIP+":"+info.SrcPort, info.DstIP+":"+info.DstPort)
		return
	}
	//starlog.Debugf("packet:%v <==> %v\n", info.SrcIP+":"+info.SrcPort, info.DstIP+":"+info.DstPort)
	if !t.Match(netm.ConnectionStat{
		Status: "PCAP",
		Laddr: netm.Addr{
			IP:   info.SrcIP,
			Port: uint32(tcp.SrcPort),
		},
		Raddr: netm.Addr{
			IP:   info.DstIP,
			Port: uint32(tcp.DstPort),
		},
	}) && !t.Match(netm.ConnectionStat{
		Status: "PCAP",
		Raddr: netm.Addr{
			IP:   info.SrcIP,
			Port: uint32(tcp.SrcPort),
		},
		Laddr: netm.Addr{
			IP:   info.DstIP,
			Port: uint32(tcp.DstPort),
		},
	}) {
		return
	}
	starlog.Noticef("Sending RST... %v:%v <==> %v:%v SEQ=%d ACK=%d\n", info.SrcIP, info.SrcPort, info.DstIP, info.DstPort, info.TcpSeq(), info.TcpAck())
	RealSendRST(t.pc.Handle, info, t.KillType, t.RstNumbers)
}

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