feat: more rule

cleaner/cleaner.go

@@ -16,9 +16,11 @@ const (
 )
 
 type Cleaner struct {
-	api  *qbapi.QBAPI
-	c    *config
-	rset *uaRuleSet
+	api        *qbapi.QBAPI
+	c          *config
+	uaRule     *strRuleSet
+	regionRule *strRuleSet
+	ipRule     *ipRuleSet
 }
 
 func New(opts ...Option) (*Cleaner, error) {
@@ -33,11 +35,19 @@ func New(opts ...Option) (*Cleaner, error) {
 	if err != nil {
 		return nil, fmt.Errorf("new api fail, err:%w", err)
 	}
-	rst, err := makeUserAgentRuleSet(c.rs)
+	uaRule, err := makeStrRuleSet(c.uaRs)
 	if err != nil {
-		return nil, fmt.Errorf("make rule set failed, err:%w", err)
+		return nil, fmt.Errorf("make ua rule set failed, err:%w", err)
 	}
-	return &Cleaner{api: api, c: c, rset: rst}, nil
+	regionRule, err := makeStrRuleSet(c.regionRs)
+	if err != nil {
+		return nil, fmt.Errorf("make region rule set failed, err:%w", err)
+	}
+	ipRule, err := makeIPRuleSet(c.ipRs)
+	if err != nil {
+		return nil, fmt.Errorf("make ip rule set failed, err:%w", err)
+	}
+	return &Cleaner{api: api, c: c, uaRule: uaRule, regionRule: regionRule, ipRule: ipRule}, nil
 }
 
 func (c *Cleaner) Start() error {
@@ -93,9 +103,17 @@ func (c *Cleaner) ensureLogin(ctx context.Context) error {
 	return nil
 }
 
-func (c *Cleaner) isBlackListUserAgent(info *qbapi.TorrentPeerItem) bool {
-	ua := info.Client
-	return c.rset.isMatch(ua)
+func (c *Cleaner) isInBlackList(info *qbapi.TorrentPeerItem) bool {
+	if c.uaRule.isMatch(info.Client) {
+		return true
+	}
+	if c.regionRule.isMatch(info.CountryCode) {
+		return true
+	}
+	if c.ipRule.isMatch(info.Ip) {
+		return true
+	}
+	return false
 }
 
 func (c *Cleaner) doCheckLogic(ctx context.Context, item *qbapi.TorrentListItem) error {
@@ -107,7 +125,7 @@ func (c *Cleaner) doCheckLogic(ctx context.Context, item *qbapi.TorrentListItem)
 	pd := peers.Peers
 	for addr, info := range pd {
 		//检查ban客户端
-		if c.isBlackListUserAgent(info) {
+		if c.isInBlackList(info) {
 			banConns[addr] = info
 			continue
 		}

cleaner/config.go

@@ -5,7 +5,9 @@ import (
 )
 
 type config struct {
-	rs []string
+	uaRs     []string
+	ipRs     []string
+	regionRs []string
 
 	//
 	interval time.Duration
@@ -23,9 +25,21 @@ func WithQBConfig(host, u, p string) Option {
 	}
 }
 
-func WithAddRules(rs ...string) Option {
+func WithAddUaRule(rs ...string) Option {
 	return func(c *config) {
-		c.rs = append(c.rs, rs...)
+		c.uaRs = append(c.uaRs, rs...)
+	}
+}
+
+func WithAddIPRule(rs ...string) Option {
+	return func(c *config) {
+		c.ipRs = append(c.ipRs, rs...)
+	}
+}
+
+func WithAddRegionRule(rs ...string) Option {
+	return func(c *config) {
+		c.regionRs = append(c.regionRs, rs...)
 	}
 }
 

cleaner/rule.go

@@ -2,47 +2,48 @@ package cleaner
 
 import (
 	"fmt"
+	"net"
 	"regexp"
 	"strings"
 )
 
 const (
-	defaultFullPrefix   = "full:"
+	defaultPrefixPrefix = "prefix:"
 	defaultRegexpPrefix = "regexp:"
 )
 
-type uaRuleSet struct {
+type strRuleSet struct {
 	regexp []*regexp.Regexp
 	full   map[string]struct{}
 	prefix []string
 }
 
-func (r *uaRuleSet) isMatch(ua string) bool {
-	if _, ok := r.full[ua]; ok {
+func (r *strRuleSet) isMatch(str string) bool {
+	if _, ok := r.full[str]; ok {
 		return true
 	}
 	for _, prefix := range r.prefix {
-		if strings.HasPrefix(ua, prefix) {
+		if strings.HasPrefix(str, prefix) {
 			return true
 		}
 	}
 	for _, reg := range r.regexp {
-		if reg.MatchString(ua) {
+		if reg.MatchString(str) {
 			return true
 		}
 	}
 	return false
 }
 
-func makeUserAgentRuleSet(rs []string) (*uaRuleSet, error) {
-	set := &uaRuleSet{
+func makeStrRuleSet(rs []string) (*strRuleSet, error) {
+	set := &strRuleSet{
 		regexp: nil,
 		full:   make(map[string]struct{}),
 		prefix: nil,
 	}
 	for _, r := range rs {
-		if strings.HasPrefix(r, defaultFullPrefix) {
-			set.full[r[len(defaultFullPrefix):]] = struct{}{}
+		if strings.HasPrefix(r, defaultPrefixPrefix) {
+			set.prefix = append(set.prefix, r[len(defaultPrefixPrefix):])
 			continue
 		}
 		if strings.HasPrefix(r, defaultRegexpPrefix) {
@@ -52,8 +53,50 @@ func makeUserAgentRuleSet(rs []string) (*uaRuleSet, error) {
 				return nil, fmt.Errorf("compile regexp failed, err:%w", err)
 			}
 			set.regexp = append(set.regexp, expr)
+			continue
+		}
+		set.full[r] = struct{}{}
+	}
+	return set, nil
+}
+
+type ipRuleSet struct {
+	ip   map[string]struct{}
+	cidr []*net.IPNet
+}
+
+func (r *ipRuleSet) isMatch(strip string) bool {
+	if _, ok := r.ip[strip]; ok {
+		return true
+	}
+	ip := net.ParseIP(strip)
+	if ip == nil {
+		return false
+	}
+	for _, cidr := range r.cidr {
+		if cidr.Contains(ip) {
+			return true
+		}
+	}
+	return false
+}
+
+func makeIPRuleSet(rs []string) (*ipRuleSet, error) {
+	set := &ipRuleSet{
+		ip:   make(map[string]struct{}),
+		cidr: nil,
+	}
+	for _, item := range rs {
+		if strings.Contains(item, "/") {
+			//TODO: 合并cidr
+			_, cidr, err := net.ParseCIDR(item)
+			if err != nil {
+				return nil, err
+			}
+			set.cidr = append(set.cidr, cidr)
+			continue
 		}
-		set.prefix = append(set.prefix, r)
+		set.ip[item] = struct{}{}
 	}
 	return set, nil
 }

cleaner/rule_test.go

@@ -0,0 +1,37 @@
+package cleaner
+
+import (
+	"github.com/stretchr/testify/assert"
+	"testing"
+)
+
+func TestIP(t *testing.T) {
+	set, err := makeIPRuleSet([]string{
+		"127.0.0.1",
+		"192.168.0.0/16",
+		"10.0.0.0/8",
+	})
+	assert.NoError(t, err)
+	assert.True(t, set.isMatch("127.0.0.1"))
+	assert.True(t, set.isMatch("192.168.50.6"))
+	assert.True(t, set.isMatch("10.1.2.3"))
+	assert.False(t, set.isMatch("11.1.2.3"))
+	assert.False(t, set.isMatch("192.167.50.1"))
+	assert.False(t, set.isMatch("127.0.0.2"))
+
+}
+
+func TestStrSet(t *testing.T) {
+	set, err := makeStrRuleSet([]string{
+		"prefix:abc.com",
+		"ccc.com",
+		"regexp:^hhhh.*$",
+	})
+	assert.NoError(t, err)
+	assert.True(t, set.isMatch("abc.com.qqq.cc"))
+	assert.True(t, set.isMatch("ccc.com"))
+	assert.True(t, set.isMatch("hhhhru ok ?"))
+	assert.False(t, set.isMatch("aaa.abc.com"))
+	assert.False(t, set.isMatch("cccc.com"))
+	assert.False(t, set.isMatch("ahhhhqqq"))
+}

config/config.go

@@ -13,10 +13,12 @@ type QBConfig struct {
 }
 
 type Config struct {
-	QBConfig  QBConfig         `json:"qb_config"`
-	LogConfig logger.LogConfig `json:"log_config"`
-	UaList    []string         `json:"ua_list"`
-	Interval  int              `json:"interval"`
+	QBConfig        QBConfig         `json:"qb_config"`
+	LogConfig       logger.LogConfig `json:"log_config"`
+	BlacklistUa     []string         `json:"blacklist_ua"`
+	BlacklistRegion []string         `json:"blacklist_region"`
+	BlacklistIP     []string         `json:"blacklist_ip"`
+	Interval        int              `json:"interval"`
 }
 
 func Parse(file string) (*Config, error) {

go.mod

@@ -3,16 +3,20 @@ module qb-helper
 go 1.21
 
 require (
+	github.com/stretchr/testify v1.8.4
 	github.com/xxxsen/common v0.1.2
 	github.com/xxxsen/qbapi v0.0.1
 	github.com/xxxsen/runner v0.0.1
 	go.uber.org/zap v1.23.0
 )
 
 require (
+	github.com/davecgh/go-spew v1.1.1 // indirect
+	github.com/pmezard/go-difflib v1.0.0 // indirect
 	go.uber.org/atomic v1.7.0 // indirect
 	go.uber.org/multierr v1.6.0 // indirect
 	golang.org/x/net v0.17.0 // indirect
 	golang.org/x/sync v0.6.0 // indirect
 	gopkg.in/natefinch/lumberjack.v2 v2.0.0 // indirect
+	gopkg.in/yaml.v3 v3.0.1 // indirect
 )

go.sum

@@ -38,6 +38,7 @@ golang.org/x/sys v0.0.0-20210423082822-04245dca01da/go.mod h1:h1NjWce9XRLGQEsW7w
 golang.org/x/term v0.0.0-20201126162022-7de9c90e9dd1/go.mod h1:bj7SfCRtBDWHUb9snDiAeCFNEtKQo2Wmx5Cou7ajbmo=
 golang.org/x/text v0.3.6/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ=
 golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ=
+gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
 gopkg.in/natefinch/lumberjack.v2 v2.0.0 h1:1Lc07Kr7qY4U2YPouBjpCLxpiyxIVoxqXgkXLknAOE8=
 gopkg.in/natefinch/lumberjack.v2 v2.0.0/go.mod h1:l0ndWWf7gzL7RNwBG7wST/UCcT4T24xpD6X8LsfU/+k=

main.go

@@ -24,7 +24,9 @@ func main() {
 	svc, err := cleaner.New(
 		cleaner.WithQBConfig(conf.QBConfig.Host, conf.QBConfig.Username, conf.QBConfig.Password),
 		cleaner.WithInterval(time.Duration(conf.Interval)*time.Second),
-		cleaner.WithAddRules(conf.UaList...),
+		cleaner.WithAddUaRule(conf.BlacklistUa...),
+		cleaner.WithAddIPRule(conf.BlacklistIP...),
+		cleaner.WithAddRegionRule(conf.BlacklistRegion...),
 	)
 	if err != nil {
 		logkit.Fatal("init cleaner failed", zap.Error(err))