优化主进程子进程获取参数
This commit is contained in:
parent
cb0f58f297
commit
46f3d1f264
22
cap.go
22
cap.go
@ -29,7 +29,7 @@ func printAvailableInterfaces() {
|
||||
|
||||
// 判断 IP 是否在链表中
|
||||
func isIPInList(ip net.IP) bool {
|
||||
for e := ipList.Front(); e != nil; e = e.Next() {
|
||||
for e := IpList.Front(); e != nil; e = e.Next() {
|
||||
if e.Value.(net.IP).Equal(ip) {
|
||||
return true
|
||||
}
|
||||
@ -43,10 +43,10 @@ func printPacketInfo(packet gopacket.Packet) {
|
||||
|
||||
if ipLayer != nil {
|
||||
ip, _ := ipLayer.(*layers.IPv4)
|
||||
ipMutex.Lock()
|
||||
defer ipMutex.Unlock()
|
||||
IpMutex.Lock()
|
||||
defer IpMutex.Unlock()
|
||||
if !isIPInList(ip.SrcIP) {
|
||||
ipList.PushBack(ip.SrcIP)
|
||||
IpList.PushBack(ip.SrcIP)
|
||||
log.Printf("\033[31m 已添加源 IP: %s 到链表 \033[0m\n", ip.SrcIP)
|
||||
|
||||
}
|
||||
@ -54,12 +54,10 @@ func printPacketInfo(packet gopacket.Packet) {
|
||||
}
|
||||
|
||||
func startPacketCapture() {
|
||||
handle, err := pcap.OpenLive(*interfaceName, 65535, true, pcap.BlockForever)
|
||||
handle, err := pcap.OpenLive(*InterfaceName, 65535, true, pcap.BlockForever)
|
||||
if err != nil {
|
||||
log.Fatalf("打开网络接口 %s 出错: %v", *interfaceName, err)
|
||||
log.Fatalf("打开网络接口 %s 出错: %v", *InterfaceName, err)
|
||||
}
|
||||
defer handle.Close()
|
||||
|
||||
defer func() {
|
||||
fmt.Println("清理资源...")
|
||||
if handle != nil {
|
||||
@ -67,14 +65,14 @@ func startPacketCapture() {
|
||||
}
|
||||
}()
|
||||
|
||||
err = handle.SetBPFFilter(*filter)
|
||||
err = handle.SetBPFFilter(*Protocol)
|
||||
if err != nil {
|
||||
log.Fatalf("设置 BPF 过滤器出错: %v", err)
|
||||
}
|
||||
|
||||
var pcapWriter *pcapgo.Writer
|
||||
if *outputFile != "" {
|
||||
file, err := os.Create(*outputFile)
|
||||
if *PcapFile != "" {
|
||||
file, err := os.Create(*PcapFile)
|
||||
if err != nil {
|
||||
log.Fatalf("创建输出文件出错: %v", err)
|
||||
}
|
||||
@ -88,7 +86,7 @@ func startPacketCapture() {
|
||||
}
|
||||
|
||||
packetSource := gopacket.NewPacketSource(handle, handle.LinkType())
|
||||
fmt.Printf("正在监听网络接口 %s，使用过滤器 '%s'...\n", *interfaceName, *filter)
|
||||
fmt.Printf("正在监听网络接口 %s，使用过滤器 '%s'...\n", *InterfaceName, *Protocol)
|
||||
|
||||
sigChan := make(chan os.Signal, 1)
|
||||
signal.Notify(sigChan, os.Interrupt, syscall.SIGTERM)
|
||||
|
||||
118
main.go
118
main.go
@ -20,24 +20,31 @@ var (
|
||||
child = flag.Bool("child", false, "子进程模式")
|
||||
)
|
||||
|
||||
// 全局变量
|
||||
var (
|
||||
listInterfaces bool
|
||||
interfaceName *string
|
||||
outputFile *string
|
||||
filter *string
|
||||
interfacesList bool // 是否列出网络接口
|
||||
InterfaceName *string // 网络接口名称
|
||||
PcapFile *string // 输出文件名
|
||||
Protocol *string // BPF 过滤器
|
||||
|
||||
SIG int
|
||||
MAX_SIG = 10
|
||||
SIG_NAME *string
|
||||
SIG int // 当前使用的 ipset 集合编号
|
||||
MAX_SIG = 10 // 最大 ipset 集合数量
|
||||
SIG_NAME string // 当前使用的 ipset 集合名称
|
||||
|
||||
// 存储 IPv4 地址的链表
|
||||
ipList = list.New()
|
||||
ipMutex sync.Mutex // 用于保护 ipList 的互斥锁
|
||||
IpList = list.New() // 存储 IPv4 地址的链表
|
||||
IpMutex sync.Mutex // 保护 ipList 的互斥锁
|
||||
)
|
||||
|
||||
// 启动子进程
|
||||
func startChildProcess() (*exec.Cmd, error) {
|
||||
cmd := exec.Command(os.Args[0], "-child=true", "-i", *interfaceName)
|
||||
args := []string{}
|
||||
for _, arg := range os.Args[1:] {
|
||||
if !strings.HasPrefix(arg, "-child") { // 只过滤 -child 标志
|
||||
args = append(args, arg)
|
||||
}
|
||||
}
|
||||
args = append(args, "-child=true")
|
||||
cmd := exec.Command(os.Args[0], args...)
|
||||
cmd.Stdout = os.Stdout
|
||||
cmd.Stderr = os.Stderr
|
||||
|
||||
@ -56,10 +63,9 @@ func stopChildProcess(cmd *exec.Cmd) error {
|
||||
|
||||
// 尝试优雅地停止子进程
|
||||
if err := cmd.Process.Signal(syscall.SIGTERM); err != nil {
|
||||
return fmt.Errorf("停止子进程失败: %w", err)
|
||||
return fmt.Errorf("发送 SIGTERM 信号失败: %w", err)
|
||||
}
|
||||
|
||||
// 等待一段时间让子进程退出
|
||||
done := make(chan error, 1)
|
||||
go func() {
|
||||
done <- cmd.Wait()
|
||||
@ -70,8 +76,8 @@ func stopChildProcess(cmd *exec.Cmd) error {
|
||||
if err != nil {
|
||||
return fmt.Errorf("等待子进程退出时出错: %w", err)
|
||||
}
|
||||
case <-time.After(5 * time.Second):
|
||||
fmt.Println("子进程未在规定时间内退出，强制终止...")
|
||||
case <-time.After(5 * time.Second): // 超时时间调整为5秒
|
||||
fmt.Println("子进程未在规定时间内退出，尝试强制终止...")
|
||||
if err := cmd.Process.Kill(); err != nil {
|
||||
return fmt.Errorf("强制终止子进程失败: %w", err)
|
||||
}
|
||||
@ -84,46 +90,63 @@ func stopChildProcess(cmd *exec.Cmd) error {
|
||||
// 等待信号并优雅退出
|
||||
func waitForSignalAndCleanUp(cmd *exec.Cmd) {
|
||||
sigChan := make(chan os.Signal, 1)
|
||||
signal.Notify(sigChan, os.Interrupt, syscall.SIGTERM)
|
||||
signal.Notify(sigChan, syscall.SIGHUP, syscall.SIGINT, syscall.SIGTERM, syscall.SIGQUIT)
|
||||
|
||||
for sig := range sigChan {
|
||||
fmt.Printf("主进程收到信号: %v\n", sig)
|
||||
if cmd != nil && cmd.Process != nil {
|
||||
_ = cmd.Process.Signal(sig) // 转发信号到子进程
|
||||
}
|
||||
if sig == syscall.SIGINT || sig == syscall.SIGTERM {
|
||||
break
|
||||
}
|
||||
}
|
||||
|
||||
<-sigChan
|
||||
fmt.Println("主进程收到终止信号，准备清理子进程...")
|
||||
if err := stopChildProcess(cmd); err != nil {
|
||||
log.Fatalf("清理子进程时遇到错误: %v", err)
|
||||
}
|
||||
fmt.Println("主进程退出")
|
||||
}
|
||||
|
||||
// 守护进程模式（简化版）
|
||||
// 守护进程模式
|
||||
func startDaemon() {
|
||||
// 注意：此方法不会创建真正的守护进程，它只是以分离的方式启动一个新的进程。
|
||||
// 对于创建真正的守护进程，你应该使用 systemd 或其他服务管理器。
|
||||
cmd := exec.Command(os.Args[0], "-d=false", "-child=false", "-i", *interfaceName)
|
||||
// 创建一个新的实例并让当前进程退出。注意这并不是守护进程的标准实现。
|
||||
args := []string{}
|
||||
for _, arg := range os.Args[1:] {
|
||||
if !strings.HasPrefix(arg, "-d") && !strings.HasPrefix(arg, "-child") { // 过滤掉 -d 和 -child 标志
|
||||
args = append(args, arg)
|
||||
}
|
||||
}
|
||||
args = append(args, "-d=false", "-child=false")
|
||||
|
||||
cmd := exec.Command(os.Args[0], args...)
|
||||
cmd.Stdout = os.Stdout
|
||||
cmd.Stderr = os.Stderr
|
||||
err := cmd.Start()
|
||||
if err != nil {
|
||||
if err := cmd.Start(); err != nil {
|
||||
log.Fatalf("无法启动新实例: %v", err)
|
||||
}
|
||||
fmt.Println("新实例已启动，当前进程将退出")
|
||||
os.Exit(0)
|
||||
}
|
||||
|
||||
func runChildProcess() { // 子进程逻辑
|
||||
//fmt.Println("子进程启动中...")
|
||||
// 子进程逻辑
|
||||
func runChildProcess() {
|
||||
fmt.Println("子进程启动中...")
|
||||
|
||||
sigChan := make(chan os.Signal, 1)
|
||||
signal.Notify(sigChan, os.Interrupt, syscall.SIGTERM)
|
||||
signal.Notify(sigChan, syscall.SIGHUP, syscall.SIGINT, syscall.SIGTERM, syscall.SIGQUIT)
|
||||
|
||||
ticker := time.NewTicker(3 * time.Second)
|
||||
defer ticker.Stop()
|
||||
|
||||
for {
|
||||
select {
|
||||
case <-sigChan:
|
||||
fmt.Println("子进程收到终止信号，退出中...")
|
||||
case sig := <-sigChan:
|
||||
fmt.Printf("子进程收到信号: %v，准备退出...\n", sig)
|
||||
return
|
||||
case <-ticker.C:
|
||||
//fmt.Println("子进程!!!")
|
||||
fmt.Println("子进程正在运行...")
|
||||
fmt.Println("子进程接收到的参数:", os.Args)
|
||||
}
|
||||
}
|
||||
}
|
||||
@ -138,7 +161,7 @@ func runMainProcess() { // 主进程逻辑
|
||||
}
|
||||
|
||||
SIG = 0
|
||||
SIG_NAME := fmt.Sprintf("root%d", SIG)
|
||||
SIG_NAME = fmt.Sprintf("root%d", SIG)
|
||||
if return_value := IsIpset(SIG_NAME); return_value == 1 {
|
||||
createIPSet(SIG_NAME)
|
||||
}
|
||||
@ -149,10 +172,9 @@ func runMainProcess() { // 主进程逻辑
|
||||
// 启动IP地域判断，管理
|
||||
go func() {
|
||||
for {
|
||||
ipMutex.Lock() // 锁定互斥锁
|
||||
if ipList.Len() > 0 { // 链表不为空
|
||||
e1 := ipList.Front() // 获取链表第一个元素
|
||||
//fmt.Println(e1.Value) // 输出第一个元素
|
||||
IpMutex.Lock() // 锁定互斥锁
|
||||
if IpList.Len() > 0 { // 链表不为空
|
||||
e1 := IpList.Front() // 获取链表第一个元素
|
||||
|
||||
region, _ := ip2region(e1.Value.(net.IP).String())
|
||||
|
||||
@ -171,16 +193,16 @@ func runMainProcess() { // 主进程逻辑
|
||||
log.Printf("curl_ %s 位置 %s\n", e1.Value.(net.IP).String(), position) // 打印地域
|
||||
AddIPSet(SIG_NAME, e1.Value.(net.IP).String()) // 添加 Ip 到 ipset 集合
|
||||
|
||||
ipList.Remove(e1) // 移除第一个元素
|
||||
IpList.Remove(e1) // 移除第一个元素
|
||||
}
|
||||
} else { // 这时是国内地址
|
||||
ipList.Remove(e1)
|
||||
IpList.Remove(e1)
|
||||
}
|
||||
|
||||
log.Printf("Ip 链表长度:%d\n", ipList.Len()) // 链表长度
|
||||
log.Printf("Ip 链表长度:%d\n", IpList.Len()) // 链表长度
|
||||
time.Sleep(2 * time.Second)
|
||||
}
|
||||
ipMutex.Unlock() // 链表解锁互斥锁
|
||||
IpMutex.Unlock() // 链表解锁互斥锁
|
||||
|
||||
}
|
||||
|
||||
@ -188,9 +210,7 @@ func runMainProcess() { // 主进程逻辑
|
||||
|
||||
// 启动防火墙管理
|
||||
go func() {
|
||||
|
||||
for {
|
||||
|
||||
if ipset_len, _ := NumIPSet(SIG_NAME); ipset_len >= 65534 {
|
||||
//log.Printf("ipset %s 列表已满 %dd\n", SIG_NAME, ipset_len)
|
||||
log.Printf("\033[31m ipset %s 列表已满 %d \033[0m\n", SIG_NAME, ipset_len)
|
||||
@ -222,10 +242,10 @@ func handleCmd() {
|
||||
// 定义命令行标志
|
||||
var instruction string
|
||||
var help bool
|
||||
interfaceName = flag.String("i", "", "指定要使用的网络接口")
|
||||
flag.BoolVar(&listInterfaces, "l", false, "列出可用的网络接口")
|
||||
filter = flag.String("f", "tcp", "指定 BPF 过滤器")
|
||||
outputFile = flag.String("o", "", "保存捕获数据的输出文件（可选）")
|
||||
InterfaceName = flag.String("i", "", "指定要使用的网络接口")
|
||||
flag.BoolVar(&interfacesList, "l", false, "列出可用的网络接口")
|
||||
Protocol = flag.String("f", "tcp", "指定 BPF 过滤器")
|
||||
PcapFile = flag.String("o", "", "保存捕获数据的输出文件（可选）")
|
||||
flag.StringVar(&instruction, "s", "", "-s start 启动 Iptables 规则\n-s stop 停止 Iptables 规则")
|
||||
flag.BoolVar(&help, "h", false, "")
|
||||
flag.BoolVar(&help, "help", false, "帮助信息")
|
||||
@ -264,12 +284,12 @@ func handleCmd() {
|
||||
}
|
||||
}
|
||||
|
||||
if listInterfaces {
|
||||
if interfacesList {
|
||||
printAvailableInterfaces()
|
||||
os.Exit(0)
|
||||
}
|
||||
|
||||
if *interfaceName == "" {
|
||||
if *InterfaceName == "" {
|
||||
log.Fatal("请使用 -i 标志指定网络接口，或者使用 -l 列出接口。")
|
||||
}
|
||||
}
|
||||
@ -284,10 +304,10 @@ func main() {
|
||||
|
||||
// 子进程逻辑
|
||||
if *child {
|
||||
runChildProcess()
|
||||
runChildProcess() // 子进程逻辑
|
||||
return
|
||||
}
|
||||
|
||||
runMainProcess()
|
||||
runMainProcess() // 主进程逻辑
|
||||
|
||||
}
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user