Skip to content

Commit

Permalink
feat(waf): show country code
Browse files Browse the repository at this point in the history
  • Loading branch information
uubulb committed Jan 4, 2025
1 parent 3999d1f commit 5454579
Show file tree
Hide file tree
Showing 5 changed files with 35 additions and 3 deletions.
18 changes: 15 additions & 3 deletions cmd/dashboard/controller/waf.go
Original file line number Diff line number Diff line change
@@ -1,11 +1,14 @@
package controller

import (
"net"
"slices"
"strconv"

"github.com/gin-gonic/gin"

"github.com/nezhahq/nezha/model"
"github.com/nezhahq/nezha/pkg/geoip"
"github.com/nezhahq/nezha/pkg/utils"
"github.com/nezhahq/nezha/service/singleton"
)
Expand All @@ -21,7 +24,7 @@ import (
// @Produce json
// @Success 200 {object} model.PaginatedResponse[[]model.WAFApiMock, model.WAFApiMock]
// @Router /waf [get]
func listBlockedAddress(c *gin.Context) (*model.Value[[]*model.WAF], error) {
func listBlockedAddress(c *gin.Context) (*model.Value[[]*model.WAFApiMock], error) {
limit, err := strconv.Atoi(c.Query("limit"))
if err != nil || limit < 1 {
limit = 25
Expand All @@ -42,8 +45,17 @@ func listBlockedAddress(c *gin.Context) (*model.Value[[]*model.WAF], error) {
return nil, err
}

return &model.Value[[]*model.WAF]{
Value: waf,
return &model.Value[[]*model.WAFApiMock]{
Value: slices.Collect(utils.Map(slices.Values(waf), func(e *model.WAF) *model.WAFApiMock {
return &model.WAFApiMock{
IP: net.IP(e.IP).String(),
BlockIdentifier: e.BlockIdentifier,
BlockReason: e.BlockReason,
BlockTimestamp: e.BlockTimestamp,
Count: e.Count,
Country: geoip.Lookup1(e.IP),
}
})),
Pagination: model.Pagination{
Offset: offset,
Limit: limit,
Expand Down
1 change: 1 addition & 0 deletions model/waf.go
Original file line number Diff line number Diff line change
Expand Up @@ -31,6 +31,7 @@ type WAFApiMock struct {
BlockReason uint8 `json:"block_reason,omitempty"`
BlockTimestamp uint64 `json:"block_timestamp,omitempty"`
Count uint64 `json:"count,omitempty"`
Country string `json:"country,omitempty"`
}

type WAF struct {
Expand Down
Binary file modified pkg/geoip/geoip.db
Binary file not shown.
8 changes: 8 additions & 0 deletions pkg/geoip/geoip.go
Original file line number Diff line number Diff line change
Expand Up @@ -50,3 +50,11 @@ func Lookup(ip net.IP) (string, error) {

return "", fmt.Errorf("IP not found")
}

func Lookup1(ip net.IP) string {
c, err := Lookup(ip)
if err != nil {
return ""
}
return c
}
11 changes: 11 additions & 0 deletions pkg/utils/utils.go
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@ package utils
import (
"crypto/rand"
"errors"
"iter"
"maps"
"math/big"
"net/netip"
Expand Down Expand Up @@ -164,3 +165,13 @@ func Unique[T comparable](s []T) []T {
}
return ret
}

func Map[T, U any](seq iter.Seq[T], f func(e T) U) iter.Seq[U] {
return func(yield func(U) bool) {
for e := range seq {
if !yield(f(e)) {
return
}
}
}
}

0 comments on commit 5454579

Please sign in to comment.