From a3059eb1cbd7a866293dc58d7fdc06ff79492c32 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E8=80=97=E5=AD=90?= <i@haozi.net>
Date: Sun, 28 Jul 2024 16:06:06 +0800
Subject: [PATCH] =?UTF-8?q?fix:=20ufw=E6=B7=BB=E5=8A=A0=E5=B8=A6ip?=
 =?UTF-8?q?=E7=9A=84=E8=A7=84=E5=88=99=E5=90=8E=E6=8A=A5=E9=94=99?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 app/http/controllers/safe_controller.go | 39 ++++++++++++++++++-------
 1 file changed, 29 insertions(+), 10 deletions(-)

diff --git a/app/http/controllers/safe_controller.go b/app/http/controllers/safe_controller.go
index 818b2207..a26d7fde 100644
--- a/app/http/controllers/safe_controller.go
+++ b/app/http/controllers/safe_controller.go
@@ -100,10 +100,17 @@ func (r *SafeController) GetFirewallRules(ctx http.Context) http.Response {
 		ports := strings.Split(match[1], " ")
 		for _, port := range ports {
 			rule := strings.Split(port, "/")
-			rules = append(rules, map[string]string{
-				"port":     rule[0],
-				"protocol": rule[1],
-			})
+			if len(rule) < 2 {
+				rules = append(rules, map[string]string{
+					"port":     rule[0],
+					"protocol": "all",
+				})
+			} else {
+				rules = append(rules, map[string]string{
+					"port":     rule[0],
+					"protocol": rule[1],
+				})
+			}
 		}
 	} else {
 		out, err := shell.Execf("ufw status | grep -v '(v6)' | grep ALLOW | awk '{print $1}'")
@@ -119,10 +126,17 @@ func (r *SafeController) GetFirewallRules(ctx http.Context) http.Response {
 		}
 		for _, port := range strings.Split(out, "\n") {
 			rule := strings.Split(port, "/")
-			rules = append(rules, map[string]string{
-				"port":     rule[0],
-				"protocol": rule[1],
-			})
+			if len(rule) < 2 {
+				rules = append(rules, map[string]string{
+					"port":     rule[0],
+					"protocol": "all",
+				})
+			} else {
+				rules = append(rules, map[string]string{
+					"port":     rule[0],
+					"protocol": rule[1],
+				})
+			}
 		}
 	}
 
@@ -200,16 +214,21 @@ func (r *SafeController) DeleteFirewallRule(ctx http.Context) http.Response {
 	if port == "" || protocol == "" {
 		return Error(ctx, http.StatusUnprocessableEntity, "参数错误")
 	}
+	if protocol == "all" {
+		protocol = ""
+	} else {
+		protocol = "/" + protocol
+	}
 
 	if os.IsRHEL() {
-		if out, err := shell.Execf("firewall-cmd --remove-port=%s/%s --permanent", port, protocol); err != nil {
+		if out, err := shell.Execf("firewall-cmd --remove-port=%s%s --permanent", port, protocol); err != nil {
 			return Error(ctx, http.StatusInternalServerError, out)
 		}
 		if out, err := shell.Execf("firewall-cmd --reload"); err != nil {
 			return Error(ctx, http.StatusInternalServerError, out)
 		}
 	} else {
-		if out, err := shell.Execf("ufw delete allow %s/%s", port, protocol); err != nil {
+		if out, err := shell.Execf("ufw delete allow %s%s", port, protocol); err != nil {
 			return Error(ctx, http.StatusInternalServerError, out)
 		}
 		if out, err := shell.Execf("ufw reload"); err != nil {