//go:build !(windows && arm64)

package tcpkill

import (
	"b612.me/bcap"
	"b612.me/starlog"
	"context"
	"fmt"
	netm "github.com/shirou/gopsutil/v4/net"
	"net"
	"runtime"
)

func (t *TCPKill) PreRun() error {
	if t.cap == nil {
		t.cap = bcap.NewPackets()
		t.ctx, t.stopFn = context.WithCancel(context.Background())
	}
	t.requests = make(chan *handler, 1024)
	if t.KillType == "" {
		t.KillType = "both"
	}
	conns, err := netm.Connections("tcp")
	if err != nil {
		starlog.Errorf("Failed to get system connections:%v\n", err)
		return err
	}
	starlog.Infof("Got %d connections\n", len(conns))
	for _, conn := range conns {
		//fmt.Printf("Connection: %v => %v PID=%v Status=%v\n", conn.Laddr, conn.Raddr, conn.Pid, conn.Status)
		if t.Match(conn) {
			fmt.Printf("Matched connection: %v:%v => %v:%v PID=%v Status=%v\n", conn.Laddr.IP, conn.Laddr.Port, conn.Raddr.IP, conn.Raddr.Port, conn.Pid, conn.Status)
			t.matchConns.Store(key(conn), conn)
			t.matchCount++
		}
	}
	if t.matchCount == 0 && !t.WaitMode {
		starlog.Warningln("No matched connection")
		return fmt.Errorf("No matched connection")
	}
	starlog.Infof("Matched %d connections\n", t.matchCount)
	return nil
}

func (t *TCPKill) Match(info netm.ConnectionStat) bool {
	if info.Status != "PCAP" {
		if t.Status != "" && t.Status != info.Status {
			return false
		}
		if info.Status == "DELETE" || info.Status == "CLOSED" || info.Status == "LISTEN" {
			return false
		}
		if runtime.GOOS == "windows" && info.Status == "TIME_WAIT" {
			return false
		}
	}
	if _, ok := t.matchConns.Load(key(info)); ok {
		return true
	}
	if t.SrcIP == "" && t.SrcPort == 0 && t.DstIP == "" && t.DstPort == 0 {
		if t.Status != "" && info.Status != "PCAP" {
			return true
		}
		return false
	}
	innerCheck := func(srcIP string, srcPort int, conns netm.Addr) bool {
		sIp := net.ParseIP(srcIP)
		if sIp == nil {
			return false
		}
		lAddr := net.ParseIP(conns.IP)
		if lAddr != nil {
			if sIp.To16() != nil && lAddr.To16() != nil && !lAddr.To16().Equal(sIp.To16()) {
				return false
			}
			if sIp.To4() != nil && lAddr.To4() != nil && !lAddr.To4().Equal(sIp.To4()) {
				return false
			}
			if (sIp.To4() != nil && lAddr.To4() == nil) || (sIp.To4() == nil && lAddr.To4() != nil) {
				return false
			}
			if srcPort != 0 && uint32(srcPort) != conns.Port {
				return false
			}
		}
		return true
	}
	if t.SrcIP != "" {
		if !innerCheck(t.SrcIP, t.SrcPort, info.Laddr) {
			return false
		}
	} else if t.SrcPort != 0 && t.SrcPort != int(info.Laddr.Port) {
		return false
	}
	if t.DstIP != "" {
		if !innerCheck(t.DstIP, t.DstPort, info.Raddr) {
			return false
		}
	} else if t.DstPort != 0 && t.DstPort != int(info.Raddr.Port) {
		return false
	}
	return true
}

func key(i netm.ConnectionStat) string {
	return fmt.Sprintf("tcp://%v:%v-%v:%v", i.Laddr.IP, i.Laddr.Port, i.Raddr.IP, i.Raddr.Port)
}