You cannot select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
star/tcpkill/tcpkill.go

109 lines
2.8 KiB
Go

//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)
}