//go:build !windows

package tcpkill

import (
	"b612.me/bcap"
	"b612.me/bcap/nfq"
	"b612.me/starlog"
	"context"
	"fmt"
	"github.com/florianl/go-nfqueue/v2"
	"github.com/google/gopacket"
	"github.com/google/gopacket/layers"
	netm "github.com/shirou/gopsutil/v4/net"
	"net"
	"os"
	"os/exec"
	"os/signal"
	"strconv"
	"strings"
	"sync"
	"syscall"
	"time"
)

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

func (t *TCPKill) doIptables() error {
	if t.SrcPort != 0 {
		if _, err := exec.Command("iptables", "-t", "raw", "-A", "PREROUTING", "-p", "tcp", "--dport", strconv.Itoa(t.SrcPort), "-j", "NFQUEUE", "--queue-num", strconv.Itoa(t.NFQNums)).CombinedOutput(); err != nil {
			return err
		}
		t.cache = append(t.cache, "-t raw -D PREROUTING -p tcp --dport "+strconv.Itoa(t.SrcPort)+" -j NFQUEUE --queue-num "+strconv.Itoa(t.NFQNums))
		if _, err := exec.Command("iptables", "-t", "raw", "-A", "OUTPUT", "-p", "tcp", "--sport", strconv.Itoa(t.SrcPort), "-j", "NFQUEUE", "--queue-num", strconv.Itoa(t.NFQNums)).CombinedOutput(); err != nil {
			return err
		}
		t.cache = append(t.cache, "-t raw -D OUTPUT -p tcp --sport "+strconv.Itoa(t.SrcPort)+" -j NFQUEUE --queue-num "+strconv.Itoa(t.NFQNums))
		return nil
	}
	if t.DstPort != 0 {
		if _, err := exec.Command("iptables", "-t", "raw", "-A", "PREROUTING", "-p", "tcp", "--sport", strconv.Itoa(t.DstPort), "-j", "NFQUEUE", "--queue-num", strconv.Itoa(t.NFQNums)).CombinedOutput(); err != nil {
			return err
		}
		t.cache = append(t.cache, "-t raw -D PREROUTING -p tcp --sport "+strconv.Itoa(t.DstPort)+" -j NFQUEUE --queue-num "+strconv.Itoa(t.NFQNums))
		if _, err := exec.Command("iptables", "-t", "raw", "-A", "OUTPUT", "-p", "tcp", "--dport", strconv.Itoa(t.DstPort), "-j", "NFQUEUE", "--queue-num", strconv.Itoa(t.NFQNums)).CombinedOutput(); err != nil {
			return err
		}
		t.cache = append(t.cache, "-t raw -D OUTPUT -p tcp --dport "+strconv.Itoa(t.DstPort)+" -j NFQUEUE --queue-num "+strconv.Itoa(t.NFQNums))
		return nil
	}
	return fmt.Errorf("No src or dst port detect,it is too dangerous to set iptables automatically without port,please operate manually")
}

func (t *TCPKill) undoIptables() {
	for _, cmd := range t.cache {
		exec.Command("iptables", strings.Fields(cmd)...).CombinedOutput()
	}
}

func (t *TCPKill) Run() error {
	if err := t.PreRun(); err != nil {
		return err
	}
	stopSignal := make(chan os.Signal)
	starlog.Noticef("Starting nfqueue capture\n")
	nf := nfq.NewNfQueue(t.ctx, uint16(t.NFQNums), 65535)
	nf.SetRecall(t.handleRoute)
	go func() {
		if err := nf.Run(); err != nil {
			starlog.Errorln(err)
			stopSignal <- syscall.SIGKILL
		}
	}()
	if t.AutoIptables {
		if _, err := exec.Command("iptables", "-V").CombinedOutput(); err != nil {
			starlog.Warningln("iptables not found, cannot auto set iptables")
			return fmt.Errorf("iptables not found")
		}
		defer t.undoIptables()
		if err := t.doIptables(); err != nil {
			starlog.Errorf("Failed to set iptables:%v\n", err)
			return err
		}
		starlog.Infof("Set iptables success\n")
	}
	signal.Notify(stopSignal, syscall.SIGINT, syscall.SIGTERM, syscall.SIGKILL)
	go t.SynSent()
	go t.handleNfResult()
	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) handleNfResult() {
	for {
		select {
		case <-t.ctx.Done():
			return
		case info := <-t.requests:
			if info == nil {
				continue
			}
			info.p.SetVerdict(info.id, <-info.fin)
		}
	}
}

func (t *TCPKill) handleRoute(id uint32, q *nfqueue.Nfqueue, p nfq.Packet) {
	if t.requests == nil {
		q.SetVerdict(id, nfqueue.NfAccept)
		return
	}
	info, err := t.cap.ParsePacket(p.Packet)
	if err != nil {
		q.SetVerdict(id, nfqueue.NfAccept)
		return
	}
	fin := make(chan int)
	t.requests <- &handler{
		id:     id,
		p:      q,
		packet: p.Packet,
		attr:   p.Attr,
		fin:    fin,
	}
	go func() {
		fin <- t.handlePacket(info, p)
	}()
}

type handler struct {
	id     uint32
	p      *nfqueue.Nfqueue
	packet gopacket.Packet
	attr   nfqueue.Attribute
	fin    chan int
}

func (t *TCPKill) handlePacket(info bcap.PacketInfo, p nfq.Packet) int {
	if p.Packet == nil {
		return nfqueue.NfAccept
	}
	tcpLayer := p.Packet.Layer(layers.LayerTypeTCP)
	if tcpLayer == nil {
		return nfqueue.NfAccept
	}
	tcp := tcpLayer.(*layers.TCP)
	if tcp.SYN && !tcp.ACK {
		//starlog.Debugf("SYN packet:%v\n", p.Packet)
		return nfqueue.NfAccept
	}
	if tcp.RST {
		starlog.Warningf("RST packet:%v <==> %v\n", info.SrcIP+":"+info.SrcPort, info.DstIP+":"+info.DstPort)
		return nfqueue.NfAccept
	}
	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 nfqueue.NfAccept
	}
	RealSendRST(info, t.KillType, t.RstNumbers)
	return nfqueue.NfAccept
}

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 err := SendSYN(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(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)
			}
			starlog.Infof("Send SYN to %v <==> %v\n", conn.Laddr.String(), conn.Raddr.String())
			return true
		})
	}
}

func RealSendRST(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(info.SrcIP, info.SrcPort, info.DstIP, info.DstPort, seq)
		}
		if target == "both" || target == "reverse" {
			SendRST(info.DstIP, info.DstPort, info.SrcIP, info.SrcPort, info.TcpAck()+(uint32(i)*uint32(info.TcpWindow())))
		}
	}
}

func SendRST(srcIP, srcPort, dstIP, dstPort string, seq uint32) error {
	if net.ParseIP(dstIP).To4() != nil {
		return sendIPv4(srcIP, srcPort, dstIP, dstPort, seq, false)
	}
	return sendIPv6(srcIP, srcPort, dstIP, dstPort, seq, false)
}

func SendSYN(srcIP, srcPort, dstIP, dstPort string, seq uint32) error {
	if net.ParseIP(dstIP).To4() != nil {
		return sendIPv4(srcIP, srcPort, dstIP, dstPort, seq, true)
	}
	return sendIPv6(srcIP, srcPort, dstIP, dstPort, seq, true)
}

func sendIPv4(srcIP, srcPort, dstIP, dstPort string, seq uint32, isSyn bool) error {
	fd, err := syscall.Socket(syscall.AF_INET, syscall.SOCK_RAW, syscall.IPPROTO_RAW)
	if err != nil {
		return err
	}
	defer syscall.Close(fd)

	dstNetIP := net.ParseIP(dstIP)
	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 err = gopacket.SerializeLayers(buffer, options, &iPv4, &tcp); err != nil {
		return err
	}
	addr := syscall.SockaddrInet4{
		Port: dPort,
		Addr: [4]byte{dstNetIP.To4()[0], dstNetIP.To4()[1], dstNetIP.To4()[2], dstNetIP.To4()[3]},
	}
	return syscall.Sendto(fd, buffer.Bytes(), 0, &addr)
}

func sendIPv6(srcIP, srcPort, dstIP, dstPort string, seq uint32, isSyn bool) error {
	fd, err := syscall.Socket(syscall.AF_INET6, syscall.SOCK_RAW, syscall.IPPROTO_RAW)
	if err != nil {
		return err
	}
	defer syscall.Close(fd)

	dstNetIP := net.ParseIP(dstIP)
	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 err = gopacket.SerializeLayers(buffer, options, &iPv6, &tcp); err != nil {
		return err
	}
	addr := syscall.SockaddrInet6{
		Port: dPort,
		Addr: [16]byte(dstNetIP.To16()),
	}
	return syscall.Sendto(fd, buffer.Bytes(), 0, &addr)
}