From 885330e94812f5ddaa61a762cf78c11b5e7be4a1 Mon Sep 17 00:00:00 2001
From: UUBulb <35923940+uubulb@users.noreply.github.com>
Date: Sat, 23 Nov 2024 16:22:22 +0800
Subject: [PATCH] update profile api (#16)

* update profile api

* rename

* fix realip assertion

* add waf api
---
 cmd/dashboard/controller/controller.go |  3 ++
 cmd/dashboard/controller/server.go     |  2 +-
 cmd/dashboard/controller/user.go       | 12 ++++---
 cmd/dashboard/controller/waf.go        | 50 ++++++++++++++++++++++++++
 cmd/dashboard/rpc/rpc.go               |  3 +-
 model/user.go                          |  5 +++
 model/waf.go                           | 17 ++++++++-
 7 files changed, 85 insertions(+), 7 deletions(-)
 create mode 100644 cmd/dashboard/controller/waf.go

diff --git a/cmd/dashboard/controller/controller.go b/cmd/dashboard/controller/controller.go
index 2efd00a..4d398b4 100644
--- a/cmd/dashboard/controller/controller.go
+++ b/cmd/dashboard/controller/controller.go
@@ -123,6 +123,9 @@ func routers(r *gin.Engine) {
 	auth.PATCH("/nat/:id", commonHandler(updateNAT))
 	auth.POST("/batch-delete/nat", commonHandler(batchDeleteNAT))
 
+	auth.GET("/waf", commonHandler(listBlockedAddress))
+	auth.POST("/batch-delete/waf", commonHandler(batchDeleteBlockedAddress))
+
 	auth.PATCH("/setting", commonHandler(updateConfig))
 
 	r.NoRoute(fallbackToFrontend)
diff --git a/cmd/dashboard/controller/server.go b/cmd/dashboard/controller/server.go
index 75f2f18..041e459 100644
--- a/cmd/dashboard/controller/server.go
+++ b/cmd/dashboard/controller/server.go
@@ -116,8 +116,8 @@ func batchDeleteServer(c *gin.Context) (any, error) {
 				delete(singleton.AlertsCycleTransferStatsStore[alert.ID].NextUpdate, sid)
 			}
 		}
-		singleton.DB.Unscoped().Delete(&model.Transfer{}, "server_id = ?", sid)
 	}
+	singleton.DB.Unscoped().Delete(&model.Transfer{}, "server_id in (?)", servers)
 	singleton.AlertsLock.Unlock()
 
 	singleton.OnServerDelete(servers)
diff --git a/cmd/dashboard/controller/user.go b/cmd/dashboard/controller/user.go
index a4a839c..f55e8b4 100644
--- a/cmd/dashboard/controller/user.go
+++ b/cmd/dashboard/controller/user.go
@@ -2,9 +2,10 @@ package controller
 
 import (
 	"github.com/gin-gonic/gin"
+	"golang.org/x/crypto/bcrypt"
+
 	"github.com/naiba/nezha/model"
 	"github.com/naiba/nezha/service/singleton"
-	"golang.org/x/crypto/bcrypt"
 )
 
 // Get profile
@@ -14,14 +15,17 @@ import (
 // @Description Get profile
 // @Tags auth required
 // @Produce json
-// @Success 200 {object} model.CommonResponse[model.User]
+// @Success 200 {object} model.CommonResponse[model.Profile]
 // @Router /profile [get]
-func getProfile(c *gin.Context) (*model.User, error) {
+func getProfile(c *gin.Context) (*model.Profile, error) {
 	auth, ok := c.Get(model.CtxKeyAuthorizedUser)
 	if !ok {
 		return nil, singleton.Localizer.ErrorT("unauthorized")
 	}
-	return auth.(*model.User), nil
+	return &model.Profile{
+		User:    *auth.(*model.User),
+		LoginIP: c.GetString(model.CtxKeyRealIPStr),
+	}, nil
 }
 
 // List user
diff --git a/cmd/dashboard/controller/waf.go b/cmd/dashboard/controller/waf.go
new file mode 100644
index 0000000..c27dcdb
--- /dev/null
+++ b/cmd/dashboard/controller/waf.go
@@ -0,0 +1,50 @@
+package controller
+
+import (
+	"github.com/gin-gonic/gin"
+
+	"github.com/naiba/nezha/model"
+	"github.com/naiba/nezha/service/singleton"
+)
+
+// List blocked addresses
+// @Summary List blocked addresses
+// @Security BearerAuth
+// @Schemes
+// @Description List server
+// @Tags auth required
+// @Produce json
+// @Success 200 {object} model.CommonResponse[[]model.WAF]
+// @Router /waf [get]
+func listBlockedAddress(c *gin.Context) ([]*model.WAF, error) {
+	var waf []*model.WAF
+	if err := singleton.DB.Find(&waf).Error; err != nil {
+		return nil, err
+	}
+
+	return waf, nil
+}
+
+// Batch delete blocked addresses
+// @Summary Edit server
+// @Security BearerAuth
+// @Schemes
+// @Description Edit server
+// @Tags auth required
+// @Accept json
+// @Param request body []string true "block list"
+// @Produce json
+// @Success 200 {object} model.CommonResponse[any]
+// @Router /batch-delete/waf [patch]
+func batchDeleteBlockedAddress(c *gin.Context) (any, error) {
+	var list []string
+	if err := c.ShouldBindJSON(&list); err != nil {
+		return nil, err
+	}
+
+	if err := model.BatchClearIP(singleton.DB, list); err != nil {
+		return nil, newGormError("%v", err)
+	}
+
+	return nil, nil
+}
diff --git a/cmd/dashboard/rpc/rpc.go b/cmd/dashboard/rpc/rpc.go
index 415606c..f760038 100644
--- a/cmd/dashboard/rpc/rpc.go
+++ b/cmd/dashboard/rpc/rpc.go
@@ -27,7 +27,8 @@ func ServeRPC() *grpc.Server {
 }
 
 func waf(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (interface{}, error) {
-	if err := model.CheckIP(singleton.DB, ctx.Value(model.CtxKeyRealIP{}).(string)); err != nil {
+	realip, _ := ctx.Value(model.CtxKeyRealIP{}).(string)
+	if err := model.CheckIP(singleton.DB, realip); err != nil {
 		return nil, err
 	}
 	return handler(ctx, req)
diff --git a/model/user.go b/model/user.go
index a09c8a6..e1f297b 100644
--- a/model/user.go
+++ b/model/user.go
@@ -5,3 +5,8 @@ type User struct {
 	Username string `json:"username,omitempty" gorm:"uniqueIndex"`
 	Password string `json:"password,omitempty" gorm:"type:char(72)"`
 }
+
+type Profile struct {
+	User
+	LoginIP string `json:"login_ip,omitempty"`
+}
diff --git a/model/waf.go b/model/waf.go
index e74135d..657ca13 100644
--- a/model/waf.go
+++ b/model/waf.go
@@ -56,7 +56,22 @@ func ClearIP(db *gorm.DB, ip string) error {
 	if err != nil {
 		return err
 	}
-	return db.Delete(&WAF{}, "ip = ?", ipBinary).Error
+	return db.Unscoped().Delete(&WAF{}, "ip = ?", ipBinary).Error
+}
+
+func BatchClearIP(db *gorm.DB, ip []string) error {
+	if len(ip) < 1 {
+		return nil
+	}
+	ips := make([][]byte, 0, len(ip))
+	for _, s := range ip {
+		ipBinary, err := utils.IPStringToBinary(s)
+		if err != nil {
+			continue
+		}
+		ips = append(ips, ipBinary)
+	}
+	return db.Unscoped().Delete(&WAF{}, "ip in (?)", ips).Error
 }
 
 func BlockIP(db *gorm.DB, ip string, reason uint8) error {