diff --git a/plugin/antiabuse/anti.go b/plugin/antiabuse/anti.go index af337959..ad2b2910 100644 --- a/plugin/antiabuse/anti.go +++ b/plugin/antiabuse/anti.go @@ -1,4 +1,4 @@ -// Package antiabuse defines anti_abuse plugin ,support abuse words check and add/remove abuse words +// Package antiabuse defines antiabuse plugin ,support abuse words check and add/remove abuse words package antiabuse import ( @@ -14,33 +14,24 @@ import ( ) func init() { - engine := control.Register("anti_abuse", &ctrl.Options[*zero.Ctx]{ + engine := control.Register("antiabuse", &ctrl.Options[*zero.Ctx]{ DisableOnDefault: false, Help: "违禁词检测", PrivateDataFolder: "anti_abuse", }) onceRule := fcext.DoOnceOnSuccess(func(ctx *zero.Ctx) bool { + managers = ctx.State["managers"].(*ctrl.Control[*zero.Ctx]).Manager db.DBPath = engine.DataFolder() + "anti_abuse.db" err := db.Open(time.Hour * 4) if err != nil { ctx.SendChain(message.Text("open db error: ", err)) return false } - err = db.Create("banUser", &banUser{}) - if err != nil { - ctx.SendChain(message.Text("create table error: ", err)) - return false - } err = db.Create("banWord", &banWord{}) if err != nil { ctx.SendChain(message.Text("create table error: ", err)) return false } - err = recoverUser() - if err != nil { - ctx.SendChain(message.Text("recover data error: ", err)) - return false - } err = recoverWord() if err != nil { ctx.SendChain(message.Text("recover data error: ", err)) @@ -48,29 +39,32 @@ func init() { } return true }) - engine.OnMessage(zero.OnlyGroup, onceRule, banRule) + engine.OnMessage(onceRule, zero.OnlyGroup, banRule) engine.OnCommand("添加违禁词", zero.OnlyGroup, zero.AdminPermission, onceRule).Handle( func(ctx *zero.Ctx) { - if err := insertWord(ctx.Event.GroupID, ctx.State["args"].(string)); err != nil { - ctx.SendChain(message.Text("add ban word error:", err)) + args := ctx.State["args"].(string) + if err := insertWord(ctx.Event.GroupID, args); err != nil { + ctx.SendChain(message.Text("error:", err)) + } else { + ctx.SendChain(message.Text(fmt.Sprintf("添加违禁词 %s 成功", args))) } }) engine.OnCommand("删除违禁词", zero.OnlyGroup, zero.AdminPermission, onceRule).Handle( func(ctx *zero.Ctx) { - if err := deleteWord(ctx.Event.GroupID, ctx.State["args"].(string)); err != nil { - ctx.SendChain(message.Text("add ban word error:", err)) + args := ctx.State["args"].(string) + if err := deleteWord(ctx.Event.GroupID, args); err != nil { + ctx.SendChain(message.Text("error:", err)) + } else { + ctx.SendChain(message.Text(fmt.Sprintf("删除违禁词 %s 成功", args))) } }) engine.OnCommand("查看违禁词", zero.OnlyGroup, onceRule).Handle( func(ctx *zero.Ctx) { - gidPrefix := fmt.Sprintf("%d-", ctx.Event.GroupID) - var words []string - _ = wordSet.Iter(func(s string) error { - trueWord := strings.SplitN(s, gidPrefix, 1)[1] - words = append(words, trueWord) - return nil - }) - ctx.SendChain(message.Text("本群违禁词有:\n", strings.Join(words, " |"))) + if set, ok := wordMap[ctx.Event.GroupID]; !ok { + ctx.SendChain(message.Text("本群无违禁词")) + } else { + ctx.SendChain(message.Text("本群违禁词有:", strings.Join(set.ToSlice(), " |"))) + } }) } diff --git a/plugin/antiabuse/database.go b/plugin/antiabuse/database.go deleted file mode 100644 index 0c53a6e0..00000000 --- a/plugin/antiabuse/database.go +++ /dev/null @@ -1,81 +0,0 @@ -package antiabuse - -import ( - "fmt" - "time" - - sqlite "github.com/FloatTech/sqlite" -) - -var db = &sqlite.Sqlite{} - -type banUser struct { - UUID string `db:"uuid"` - DueTime int64 `db:"due_time"` -} - -func insertUser(gid, uid int64) error { - obj := &banUser{fmt.Sprintf("%d-%d", gid, uid), time.Now().Add(4 * time.Hour).UnixNano()} - return db.Insert("banUser", obj) -} - -func deleteUser(gid, uid int64) error { - sql := fmt.Sprintf("WHERE uuid=%d-%d", gid, uid) - return db.Del("banUser", sql) -} - -func recoverUser() error { - if !db.CanFind("banUser", "") { - return nil - } - obj := &banUser{} - var uuids []string - err := db.FindFor("banUser", obj, "", func() error { - if time.Now().UnixNano() < obj.DueTime { - uuids = append(uuids, obj.UUID) - } else { - if err := db.Del("banUser", "WHERE uuid="+obj.UUID); err != nil { - return err - } - } - return nil - }, - ) - if err != nil { - return err - } - banSet.AddMany(uuids) - return nil -} - -type banWord struct { - GroupWord string `db:"group_word"` -} - -func insertWord(gid int64, word string) error { - obj := &banWord{fmt.Sprintf("%d-%s", gid, word)} - return db.Insert("banWord", obj) -} - -func deleteWord(gid int64, word string) error { - sql := fmt.Sprintf("WHERE group_word = %d-%s", gid, word) - return db.Del("banWord", sql) -} - -func recoverWord() error { - if !db.CanFind("banWord", "") { - return nil - } - obj := &banWord{} - var groupWords []string - err := db.FindFor("banWord", obj, "", func() error { - groupWords = append(groupWords, obj.GroupWord) - return nil - }, - ) - if err != nil { - return err - } - wordSet.AddMany(groupWords) - return nil -} diff --git a/plugin/antiabuse/set.go b/plugin/antiabuse/set.go index ce153842..92e5acef 100644 --- a/plugin/antiabuse/set.go +++ b/plugin/antiabuse/set.go @@ -3,23 +3,27 @@ package antiabuse import "sync" //Set defines HashSet structure -type Set struct { +type Set[T comparable] struct { sync.RWMutex - m map[string]struct{} + m map[T]struct{} } -var banSet = &Set{m: make(map[string]struct{})} -var wordSet = &Set{m: make(map[string]struct{})} +// NewSet creates Set with optional key(s) +func NewSet[T comparable]() *Set[T] { + return &Set[T]{m: make(map[T]struct{})} +} -// Add adds element to Set -func (s *Set) Add(key string) { +// Add adds key(s) to Set +func (s *Set[T]) Add(key ...T) { s.Lock() defer s.Unlock() - s.m[key] = struct{}{} + for _, k := range key { + s.m[k] = struct{}{} + } } -// Include asserts element in Set -func (s *Set) Include(key string) bool { +// Include asserts key in Set +func (s *Set[T]) Include(key T) bool { s.RLock() defer s.RUnlock() _, ok := s.m[key] @@ -27,9 +31,9 @@ func (s *Set) Include(key string) bool { } // Iter calls f when traversing Set -func (s *Set) Iter(f func(string) error) error { - s.Lock() - defer s.Unlock() +func (s *Set[T]) Iter(f func(T) error) error { + s.RLock() + defer s.RUnlock() var err error for key := range s.m { err = f(key) @@ -40,18 +44,19 @@ func (s *Set) Iter(f func(string) error) error { return nil } -// Remove removes element from Set -func (s *Set) Remove(key string) { +// Remove removes key from Set +func (s *Set[T]) Remove(key T) { s.Lock() defer s.Unlock() delete(s.m, key) } -// AddMany adds multiple elements to Set -func (s *Set) AddMany(keys []string) { - s.Lock() - defer s.Unlock() - for _, k := range keys { - s.m[k] = struct{}{} +// ToSlice convert Set to slice +func (s *Set[T]) ToSlice() (res []T) { + s.RLock() + defer s.RUnlock() + for key := range s.m { + res = append(res, key) } + return res } diff --git a/plugin/antiabuse/utils.go b/plugin/antiabuse/utils.go index 329b2840..8e603184 100644 --- a/plugin/antiabuse/utils.go +++ b/plugin/antiabuse/utils.go @@ -1,47 +1,110 @@ package antiabuse import ( + "errors" "fmt" + "hash/crc32" "strings" "time" + sqlite "github.com/FloatTech/sqlite" + "github.com/FloatTech/ttl" + ctrl "github.com/FloatTech/zbpctrl" + "github.com/sirupsen/logrus" zero "github.com/wdvxdr1123/ZeroBot" "github.com/wdvxdr1123/ZeroBot/message" ) +var managers *ctrl.Manager[*zero.Ctx] //managers lazy load +var errBreak = errors.New("break") +var db = &sqlite.Sqlite{} +var crc32Table = crc32.MakeTable(crc32.IEEE) +var wordMap = make(map[int64]*Set[string]) + +func onDel(uid int64, _ struct{}) { + if managers == nil { + return + } + if err := managers.DoUnblock(uid); err != nil { + logrus.Error("do unblock error:", err) + } +} + +var cache = ttl.NewCacheOn[int64, struct{}](4*time.Hour, [4]func(int64, struct{}){ + nil, nil, onDel, nil}) + +type banWord struct { + Crc32ID uint32 `db:"crc32_id"` + GroupID int64 `db:"group_id"` + Word string `db:"word"` +} + func banRule(ctx *zero.Ctx) bool { if !ctx.Event.IsToMe { - return false + return true } - gid := ctx.Event.GroupID uid := ctx.Event.UserID - uuid := fmt.Sprintf("%d-%d", gid, uid) - if banSet.Include(uuid) { - return false + gid := ctx.Event.GroupID + wordSet := wordMap[gid] + if wordSet == nil { + return true } - gidPrefix := fmt.Sprintf("%d-", ctx.Event.GroupID) - var words []string - _ = wordSet.Iter(func(s string) error { - trueWord := strings.SplitN(s, gidPrefix, 1)[1] - words = append(words, trueWord) + err := wordSet.Iter(func(word string) error { + if strings.Contains(ctx.MessageString(), word) { + if err := managers.DoBlock(uid); err != nil { + return err + } + cache.Set(uid, struct{}{}) + return errBreak + } return nil }) - for _, word := range words { - if strings.Contains(ctx.MessageString(), word) { - if err := insertUser(gid, uid); err != nil { - ctx.SendChain(message.Text("ban error: ", err)) - } - banSet.Add(uuid) - ctx.SetGroupBan(gid, uid, 4*3600) - time.AfterFunc(4*time.Hour, func() { - banSet.Remove(uuid) - if err := deleteUser(gid, uid); err != nil { - ctx.SendChain(message.Text("ban error: ", err)) - } - }) - ctx.SendChain(message.Text("检测到违禁词，已封禁/屏蔽4小时")) - return false - } + if err != nil && err != errBreak { + ctx.SendChain(message.Text("block user error:", err)) + return true } - return true + ctx.SetGroupBan(gid, uid, 4*3600) + ctx.SendChain(message.Text("检测到违禁词,已封禁/屏蔽4小时")) + return false +} + +func insertWord(gid int64, word string) error { + str := fmt.Sprintf("%d-%s", gid, word) + checksum := crc32.Checksum([]byte(str), crc32Table) + obj := &banWord{checksum, gid, word} + if _, ok := wordMap[gid]; !ok { + wordMap[gid] = NewSet[string]() + } + wordMap[gid].Add(word) + return db.Insert("banWord", obj) +} + +func deleteWord(gid int64, word string) error { + if _, ok := wordMap[gid]; !ok { + return errors.New("本群还没有违禁词~") + } + if !wordMap[gid].Include(word) { + return errors.New(word + " 不在本群违禁词集合中") + } + wordMap[gid].Remove(word) + str := fmt.Sprintf("%d-%s", gid, word) + checksum := crc32.Checksum([]byte(str), crc32Table) + sql := fmt.Sprintf("WHERE crc32_id = %d", checksum) + return db.Del("banWord", sql) +} + +func recoverWord() error { + if !db.CanFind("banWord", "") { + return nil + } + obj := &banWord{} + err := db.FindFor("banWord", obj, "", func() error { + if _, ok := wordMap[obj.GroupID]; !ok { + wordMap[obj.GroupID] = NewSet[string]() + } + wordMap[obj.GroupID].Add(obj.Word) + return nil + }, + ) + return err } diff --git a/plugin/antiabuse/utils_test.go b/plugin/antiabuse/utils_test.go new file mode 100644 index 00000000..c41334d5 --- /dev/null +++ b/plugin/antiabuse/utils_test.go @@ -0,0 +1,202 @@ +package antiabuse + +import ( + "os" + "sort" + "testing" +) + +func TestInsertWord(t *testing.T) { + wordMap = make(map[int64]*Set[string]) + defer func() { + wordMap = make(map[int64]*Set[string]) + }() + path := "test.db" + defer func() { + err := os.Remove("test.db") + if err != nil { + t.Fatal(err) + } + }() + db.DBPath = path + err := db.Open(0) + defer func() { + err := db.Close() + if err != nil { + t.Fatal(err) + } + }() + err = db.Create("banWord", &banWord{}) + if err != nil { + t.Fatal(err) + } + defer func() { + err := db.Drop("banWord") + if err != nil { + t.Fatal(err) + } + }() + err = insertWord(123, "one") + if err != nil { + t.Fatal(err) + } + if ok := wordMap[123].Include("one"); !ok { + t.Fatal(`wordMap[123] should found "one" but not`) + } + if !db.CanFind("banWord", "WHERE group_id=123 AND word= 'one' ") { + t.Fatal(`db should found 123-one but not`) + } +} + +func TestDeleteWord(t *testing.T) { + wordMap = make(map[int64]*Set[string]) + defer func() { + wordMap = make(map[int64]*Set[string]) + }() + path := "test.db" + defer func() { + err := os.Remove("test.db") + if err != nil { + t.Fatal(err) + } + }() + db.DBPath = path + err := db.Open(0) + defer func() { + err := db.Close() + if err != nil { + t.Fatal(err) + } + }() + err = db.Create("banWord", &banWord{}) + if err != nil { + t.Fatal(err) + } + defer func() { + err := db.Drop("banWord") + if err != nil { + t.Fatal(err) + } + }() + err = insertWord(123, "one") + if err != nil { + t.Fatal(err) + } + err = deleteWord(123, "one") + if err != nil { + t.Fatal(err) + } +} + +func TestShowWord(t *testing.T) { + wordMap = make(map[int64]*Set[string]) + defer func() { + wordMap = make(map[int64]*Set[string]) + }() + path := "test.db" + defer func() { + err := os.Remove("test.db") + if err != nil { + t.Fatal(err) + } + }() + db.DBPath = path + err := db.Open(0) + defer func() { + err := db.Close() + if err != nil { + t.Fatal(err) + } + }() + err = db.Create("banWord", &banWord{}) + if err != nil { + t.Fatal(err) + } + defer func() { + err := db.Drop("banWord") + if err != nil { + t.Fatal(err) + } + }() + err = insertWord(123, "one") + if err != nil { + t.Fatal(err) + } + err = insertWord(123, "one") + if err != nil { + t.Fatal(err) + } + err = insertWord(123, "two") + if err != nil { + t.Fatal(err) + } + var db123 []string + var map123 []string + obj := &banWord{} + err = db.FindFor("banWord", obj, "WHERE group_id=123", func() error { + db123 = append(db123, obj.Word) + return nil + }) + if err != nil { + t.Fatal(err) + } + sort.Strings(db123) + if len(db123) != 2 || db123[0] != "one" || db123[1] != "two" { + t.Fatal("db should found 123-one and 123-two but not") + } + map123 = wordMap[123].ToSlice() + sort.Strings(map123) + if len(map123) != 2 || map123[0] != "one" || map123[1] != "two" { + t.Fatal("wordMap[123] should found 123-one and 123-two but not") + } +} + +func TestRecoverWord(t *testing.T) { + wordMap = make(map[int64]*Set[string]) + defer func() { + wordMap = make(map[int64]*Set[string]) + }() + path := "test.db" + defer func() { + err := os.Remove("test.db") + if err != nil { + t.Fatal(err) + } + }() + db.DBPath = path + err := db.Open(0) + defer func() { + err := db.Close() + if err != nil { + t.Fatal(err) + } + }() + err = db.Create("banWord", &banWord{}) + if err != nil { + t.Fatal(err) + } + defer func() { + err := db.Drop("banWord") + if err != nil { + t.Fatal(err) + } + }() + err = insertWord(123, "one") + if err != nil { + t.Fatal(err) + } + err = insertWord(123, "two") + if err != nil { + t.Fatal(err) + } + wordMap = make(map[int64]*Set[string]) + err = recoverWord() + if err != nil { + t.Fatal(err) + } + map123 := wordMap[123].ToSlice() + sort.Strings(map123) + if len(map123) != 2 || map123[0] != "one" || map123[1] != "two" { + t.Fatal("wordMap[123] should found 123-one and 123-two but not") + } +}