From abb57571affde54180d9dd946417cc374de64f9f Mon Sep 17 00:00:00 2001 From: xxxsen Date: Thu, 14 Mar 2024 20:44:19 +0800 Subject: [PATCH] feat: more rule --- cleaner/cleaner.go | 38 +++++++++++++++++++------- cleaner/config.go | 20 ++++++++++++-- cleaner/rule.go | 65 ++++++++++++++++++++++++++++++++++++-------- cleaner/rule_test.go | 37 +++++++++++++++++++++++++ config/config.go | 10 ++++--- go.mod | 4 +++ go.sum | 1 + main.go | 4 ++- 8 files changed, 150 insertions(+), 29 deletions(-) create mode 100644 cleaner/rule_test.go diff --git a/cleaner/cleaner.go b/cleaner/cleaner.go index 6126a2e..762510a 100644 --- a/cleaner/cleaner.go +++ b/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 } diff --git a/cleaner/config.go b/cleaner/config.go index 57cd59a..71f1544 100644 --- a/cleaner/config.go +++ b/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...) } } diff --git a/cleaner/rule.go b/cleaner/rule.go index 938d123..d329824 100644 --- a/cleaner/rule.go +++ b/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 } diff --git a/cleaner/rule_test.go b/cleaner/rule_test.go new file mode 100644 index 0000000..c642ea8 --- /dev/null +++ b/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")) +} diff --git a/config/config.go b/config/config.go index 2ee50c2..2669d4b 100644 --- a/config/config.go +++ b/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) { diff --git a/go.mod b/go.mod index 1725731..5c7dd66 100644 --- a/go.mod +++ b/go.mod @@ -3,6 +3,7 @@ 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 @@ -10,9 +11,12 @@ require ( ) 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 ) diff --git a/go.sum b/go.sum index 2e750c0..887f5e4 100644 --- a/go.sum +++ b/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= diff --git a/main.go b/main.go index 0b5d2ce..0f75544 100644 --- a/main.go +++ b/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))