add max_rec fea.

kkkgo · Oct 19, 2023 · e094e19 · e094e19
1 parent 5dec223
commit e094e19
Show file tree

Hide file tree

Showing 3 changed files with 58 additions and 8 deletions.
diff --git a/FILES/usr/bin/ppg.sh b/FILES/usr/bin/ppg.sh
@@ -82,7 +82,13 @@ load_netrec() {
         if ps | grep -v "grep" | grep "wsPort"; then
             echo "PPGW REC RUNNING."
         else
-            /usr/bin/ppgw -wsPort="$clash_web_port" -secret="$clash_web_password" >/dev/tty0 2>&1 &
+            if [ -f /tmp/ppgw.ini ]; then
+                . /tmp/ppgw.ini 2>/dev/tty0
+            fi
+            if [ -z "$max_rec" ]; then
+                max_rec="5000"
+            fi
+            /usr/bin/ppgw -wsPort="$clash_web_port" -secret="$clash_web_password" -net_rec_num="$max_rec" >/dev/tty0 2>&1 &
         fi
     fi
 }
@@ -252,7 +258,7 @@ load_ovpn() {
 gen_hash() {
     if [ -f /tmp/ppgw.ini ]; then
         . /tmp/ppgw.ini 2>/dev/tty0
-        str="ppgw""$fake_cidr""$dns_ip""$dns_port""$openport""$sleeptime""$clash_web_port""$clash_web_password""$mode""$udp_enable""$socks5_ip""$socks5_port""$ovpnfile""$ovpn_username""$ovpn_password""$yamlfile""$suburl""$subtime""$fast_node""$test_node_url""$ext_node""$cpudelay""$dns_burn""$ex_dns""$net_rec"
+        str="ppgw""$fake_cidr""$dns_ip""$dns_port""$openport""$sleeptime""$clash_web_port""$clash_web_password""$mode""$udp_enable""$socks5_ip""$socks5_port""$ovpnfile""$ovpn_username""$ovpn_password""$yamlfile""$suburl""$subtime""$fast_node""$test_node_url""$ext_node""$cpudelay""$dns_burn""$ex_dns""$net_rec""$max_rec"
         echo "$str" | md5sum | grep -Eo "[a-z0-9]{32}" | head -1
     else
         echo "INI does not exist"
@@ -635,6 +641,7 @@ while true; do
         old_clash_web_port=$clash_web_port
         old_clash_web_password=$clash_web_password
         old_net_rec=$net_rec
+        old_max_rec=$max_rec
     fi
     try_conf "ppgw.ini" "ini"
     hash=$(gen_hash)
@@ -664,6 +671,9 @@ while true; do
             if [ "$old_net_rec" != "$net_rec" ]; then
                 kill_netrec
             fi
+            if [ "$old_max_rec" != "$max_rec" ]; then
+                kill_netrec
+            fi
             if [ "$mode" = "suburl" ]; then
                 get_conf "$suburl" "yaml"
             fi

diff --git a/ReadMe.md b/ReadMe.md
@@ -107,6 +107,7 @@ ex_dns="223.5.5.5:53"
 
 # Network traffic records
 net_rec=no
+max_rec=5000
 ```
 下面逐项来解释选项的用法：
 - 1 配置文件第一行必须以`#paopao-gateway`开头。配置格式为`选项="值"`。
@@ -137,12 +138,12 @@ net_rec=no
     - 1、节点使用了分区域解析，只有使用了境内DNS才能连接，参见[issue](https://github.com/kkkgo/PaoPaoGateWay/issues/20)，`dns_burn`功能可以额外对节点进行解析。
     - 2、节点DNS解析存在多个解析入口，`dns_burn`功能会把所有可能的入口都作为新节点加入到配置文件中，在测速的时候就可以选择到速度最好的入口，而不是随机选择。
     - 3、节点的所有可能的解析结果都会被临时硬编码到配置文件中，除非所有节点都测速失败或者订阅更新，该配置文件不会变化，可以减少节点的DNS查询，使用IP直连，并有效避免节点临时出现可能的DNS污染或者DNS故障的情况，比如节点域名忘记续费导致解析失败。
-- 15 `net_rec`选项：网络流量记录功能，可以记录网关连接了哪些域名、上传下载消耗了多少流量、客户端IP，并默认按照消耗流量的大小排序，该功能可以根据实际情况方便地调整分析自己的域名规则列表。设置为`yes`开启该功能后可以在web界面点击[记录]，选择下载表格或者在线加载。注意事项：
+- 15 `net_rec`选项：网络流量记录功能，可以记录网关连接了哪些域名、上传下载消耗了多少流量、客户端IP，并默认按照消耗流量的大小排序，该功能可以根据实际情况方便地调整分析自己的域名规则列表。设置为`yes`开启该功能后可以在web界面点击[记录]，选择下载表格或者在线加载。其中`max_rec`选项指定最大记录数，默认为5000，当记录的内容超过`max_rec`的2倍后，仅保留前`max_rec`项记录。注意事项：  
     - 1、重启或者修改ppgw的密码、Web端口、关了再开`net_rec`选项会导致数据清空。
     - 2、如果流量太小连接持续时间过短，有可能在记录之前连接已经关闭，流量会显示为0B。
     - 3、客户端IP记录的是最后一次连接该域名的客户端IP。
     - 4、数据下载没有身份验证。
-    - 5、理论上会略微增加资源占用，取决于你的并发连接数量，可适当增加运行资源。
+    - 5、理论上会略微增加资源占用，取决于你的并发连接数量以及`max_rec`，可适当增加运行资源。
 
 ## 使用docker定制ISO镜像:ppwgiso
 ![pull](https://img.shields.io/docker/pulls/sliamb/ppgwiso.svg) ![size](https://img.shields.io/docker/image-size/sliamb/ppgwiso)   

diff --git a/ppgw.go b/ppgw.go
@@ -50,6 +50,7 @@ var (
 	reload                bool
 	closeall              bool
 	wsPort                string
+	net_rec_num           string
 )
 
 var orange = "\033[38;5;208m"
@@ -153,10 +154,15 @@ func main() {
 	flag.BoolVar(&closeall, "closeall", false, "close all connections.")
 	//ws catch
 	flag.StringVar(&wsPort, "wsPort", "", "wsPort")
+	flag.StringVar(&net_rec_num, "net_rec_num", "", "net_rec_num")
 
 	flag.Parse()
 	//net_rec
 	if wsPort != "" && secret != "" {
+		max_rec, err := strconv.Atoi(net_rec_num)
+		if err != nil {
+			max_rec = 5000
+		}
 		wsURL := "ws://127.0.0.1:" + wsPort + "/connections?token=" + secret
 		fmt.Printf("\n" + green + "[PaoPaoGW REC]" + reset + "Start NET REC :" + wsPort + " \n")
 
@@ -224,22 +230,42 @@ func main() {
 							}
 						}
 					}
+
+					if len(domainInfoMap) > max_rec*2 {
+						var trimmedList []DomainInfo
+						for _, info := range domainInfoMap {
+							trimmedList = append(trimmedList, *info)
+						}
+						sort.Slice(trimmedList, func(i, j int) bool {
+							return trimmedList[i].Download+trimmedList[i].Upload > trimmedList[j].Download+trimmedList[j].Upload
+						})
+						if len(trimmedList) > max_rec {
+							trimmedList = trimmedList[:max_rec]
+						}
+						domainInfoMap = make(map[string]*DomainInfo)
+						for i := range trimmedList {
+							domainInfoMap[trimmedList[i].Domain] = &trimmedList[i]
+						}
+					}
 					mutex.Unlock()
 				}
 			}
 		}()
-
-		for range time.Tick(1 * time.Second) {
+		for range time.Tick(3 * time.Second) {
 			mutex.Lock()
 			domainInfoList = nil
 			for _, info := range domainInfoMap {
 				domainInfoList = append(domainInfoList, *info)
 			}
 			sort.Sort(domainInfoList)
-			mutex.Unlock()
-			file, err := os.Create("/etc/config/clash/clash-dashboard/data.csv")
+
+			newFilePath := "/etc/config/clash/clash-dashboard/data.csv.new"
+			oldFilePath := "/etc/config/clash/clash-dashboard/data.csv.old"
+			currentFilePath := "/etc/config/clash/clash-dashboard/data.csv"
+			file, err := os.Create(newFilePath)
 			if err != nil {
 				fmt.Printf("\n" + red + "[PaoPaoGW REC]" + reset + "Failed to create CSV file.\n")
+				mutex.Unlock()
 				continue
 			}
 			defer file.Close()
@@ -259,6 +285,19 @@ func main() {
 			if err := writer.Error(); err != nil {
 				fmt.Printf("\n" + red + "[PaoPaoGW REC]" + reset + "Error flushing CSV data.\n")
 			}
+			file.Close()
+			if _, err := os.Stat(currentFilePath); err == nil {
+				os.Rename(currentFilePath, oldFilePath)
+			}
+
+			if err := os.Rename(newFilePath, currentFilePath); err != nil {
+				fmt.Printf("\n" + red + "[PaoPaoGW REC]" + reset + "Failed to replace CSV file.\n")
+				os.Rename(oldFilePath, currentFilePath)
+			}
+			if _, err := os.Stat(oldFilePath); err == nil {
+				os.Remove(oldFilePath)
+			}
+			mutex.Unlock()
 		}
 		os.Exit(0)
 	}