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.
109 lines
2.8 KiB
Go
109 lines
2.8 KiB
Go
5 days ago
|
//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)
|
||
|
}
|