//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
	Pids         []int32
	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())
}