diff --git a/README.md b/README.md index 8242657a..d529dc15 100644 --- a/README.md +++ b/README.md @@ -122,6 +122,12 @@ zerobot -h -t token -u url [-d|w] [-g 监听地址:端口] qq1 qq2 qq3 ... - [x] 添加[涩图/二次元/风景/车万][P站图片ID] - [x] 删除[涩图/二次元/风景/车万][P站图片ID] - [x] > setu status +- **本地涩图** `import _ "github.com/FloatTech/ZeroBot-Plugin/plugin_nativesetu"` + - [x] 来份本地[xxx] + - [x] 刷新本地[xxx] + - [x] 设置本地setu绝对路径[xxx] + - [x] 刷新所有本地setu + - [x] 所有本地setu分类 - **lolicon** `import _ "github.com/FloatTech/ZeroBot-Plugin/plugin_lolicon"` - [x] 来份萝莉 - **搜图** `import _ "github.com/FloatTech/ZeroBot-Plugin/plugin_saucenao"` diff --git a/go.mod b/go.mod index 1a6195dd..0e0bce6d 100644 --- a/go.mod +++ b/go.mod @@ -6,6 +6,7 @@ require ( github.com/FloatTech/AnimeAPI v1.1.10 github.com/FloatTech/ZeroBot-Plugin-Gif v0.2.4 github.com/FloatTech/bot-manager v1.0.1-0.20211112011524-85b9895271ed + github.com/corona10/goimagehash v1.0.3 github.com/fogleman/gg v1.3.0 github.com/fumiama/cron v1.3.0 github.com/fumiama/go-base16384 v1.2.1 diff --git a/go.sum b/go.sum index ac7915fe..4c396560 100644 --- a/go.sum +++ b/go.sum @@ -19,6 +19,8 @@ github.com/antchfx/htmlquery v1.2.3 h1:sP3NFDneHx2stfNXCKbhHFo8XgNjCACnU/4AO5gWz github.com/antchfx/htmlquery v1.2.3/go.mod h1:B0ABL+F5irhhMWg54ymEZinzMSi0Kt3I2if0BLYa3V0= github.com/antchfx/xpath v1.1.6 h1:6sVh6hB5T6phw1pFpHRQ+C4bd8sNI+O58flqtg7h0R0= github.com/antchfx/xpath v1.1.6/go.mod h1:Yee4kTMuNiPYJ7nSNorELQMr1J33uOpXDMByNYhvtNk= +github.com/corona10/goimagehash v1.0.3 h1:NZM518aKLmoNluluhfHGxT3LGOnrojrxhGn63DR/CZA= +github.com/corona10/goimagehash v1.0.3/go.mod h1:VkvE0mLn84L4aF8vCb6mafVajEb6QYMHl2ZJLn0mOGI= github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c= github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= @@ -39,10 +41,6 @@ github.com/fumiama/cron v1.3.0 h1:ZWlwuexF+HQHl3cYytEE5HNwD99q+3vNZF1GrEiXCFo= github.com/fumiama/cron v1.3.0/go.mod h1:bz5Izvgi/xEUI8tlBN8BI2jr9Moo8N4or0KV8xXuPDY= github.com/fumiama/go-base16384 v1.2.1 h1:6OGprW8g/95m2ocmryHi8mipZ7bx9StFMZDKEqLvMiA= github.com/fumiama/go-base16384 v1.2.1/go.mod h1:1HTC0QFL7BjS0DuO5Qm+fBYKQkHqmAapLbRpCxrhPXQ= -github.com/fumiama/gofastTEA v0.0.3 h1:JKcNktWArLkJe88Y+zGmsQGhqlh8IzYqUnWy2ipylh0= -github.com/fumiama/gofastTEA v0.0.3/go.mod h1:+sBZ05nCA2skZkursHNvyr8kULlEetrYTM2y5kA4rQc= -github.com/fumiama/gofastTEA v0.0.4 h1:SOWEIXBkFekhaxZoLEFk/L3rOh2X1G5PeM2TLAZycaQ= -github.com/fumiama/gofastTEA v0.0.4/go.mod h1:+sBZ05nCA2skZkursHNvyr8kULlEetrYTM2y5kA4rQc= github.com/fumiama/gofastTEA v0.0.5 h1:Pd/2eSfLl2V0CqZL8pnu1CIU8Fy4HYpLutpliXU70Ds= github.com/fumiama/gofastTEA v0.0.5/go.mod h1:+sBZ05nCA2skZkursHNvyr8kULlEetrYTM2y5kA4rQc= github.com/fumiama/gotracemoe v0.0.3 h1:iI5EbE9A3UUbfukG6+/soYPjp1S31eCNYf4tw7s6/Jc= @@ -114,6 +112,8 @@ github.com/modern-go/reflect2 v1.0.2-0.20210109003243-333559e1834b h1:6Xjqolv/0D github.com/modern-go/reflect2 v1.0.2-0.20210109003243-333559e1834b/go.mod h1:bx2lNnkwVCuqBIxFjflWJWanXIb3RllmbCylyMrvgv0= github.com/mroth/weightedrand v0.4.1 h1:rHcbUBopmi/3x4nnrvwGJBhX9d0vk+KgoLUZeDP6YyI= github.com/mroth/weightedrand v0.4.1/go.mod h1:3p2SIcC8al1YMzGhAIoXD+r9olo/g/cdJgAD905gyNE= +github.com/nfnt/resize v0.0.0-20180221191011-83c6a9932646 h1:zYyBkD/k9seD2A7fsi6Oo2LfFZAehjjQMERAvZLEDnQ= +github.com/nfnt/resize v0.0.0-20180221191011-83c6a9932646/go.mod h1:jpp1/29i3P1S/RLdc7JQKbRpFeM1dOBd8T9ki5s+AY8= github.com/onsi/ginkgo v1.6.0/go.mod h1:lLunBs/Ym6LB5Z9jYTR76FiuTmxDTDusOGeTQH+WWjE= github.com/onsi/ginkgo v1.7.0/go.mod h1:lLunBs/Ym6LB5Z9jYTR76FiuTmxDTDusOGeTQH+WWjE= github.com/onsi/gomega v1.4.3/go.mod h1:ex+gbHU/CVuBBDIJjb2X0qEXbFg53c61hWP/1CpauHY= diff --git a/main.go b/main.go index 9b614999..a97f5a4f 100644 --- a/main.go +++ b/main.go @@ -48,6 +48,7 @@ import ( _ "github.com/FloatTech/ZeroBot-Plugin/plugin_aiwife" // 随机老婆 _ "github.com/FloatTech/ZeroBot-Plugin/plugin_image_finder" // 关键字搜图 _ "github.com/FloatTech/ZeroBot-Plugin/plugin_lolicon" // lolicon 随机图片 + _ "github.com/FloatTech/ZeroBot-Plugin/plugin_nativesetu" // 本地涩图 _ "github.com/FloatTech/ZeroBot-Plugin/plugin_saucenao" // 以图搜图 _ "github.com/FloatTech/ZeroBot-Plugin/plugin_setutime" // 来份涩图 _ "github.com/FloatTech/ZeroBot-Plugin/plugin_tracemoe" // 搜番 diff --git a/plugin_nativesetu/data.go b/plugin_nativesetu/data.go new file mode 100644 index 00000000..5ae43e5a --- /dev/null +++ b/plugin_nativesetu/data.go @@ -0,0 +1,96 @@ +package nativesetu + +import ( + "image" + "io/fs" + "os" + "sync" + + "github.com/corona10/goimagehash" + "github.com/sirupsen/logrus" + "github.com/wdvxdr1123/ZeroBot/utils/helper" + + "github.com/FloatTech/ZeroBot-Plugin/utils/file" + "github.com/FloatTech/ZeroBot-Plugin/utils/process" + "github.com/FloatTech/ZeroBot-Plugin/utils/sql" +) + +// setuclass holds setus in a folder, which is the class name. +type setuclass struct { + ImgID uint64 `db:"imgid"` // ImgID 图片唯一 id (dhash) + Name string `db:"name"` // Name 图片名 +} + +var ( + setuclasses []string + db = &sql.Sqlite{DBPath: dbfile} + mu sync.RWMutex +) + +func init() { + go func() { + process.SleepAbout1sTo2s() + err := os.MkdirAll(datapath, 0755) + if err != nil { + panic(err) + } + if file.IsExist(cfgfile) { + b, err := os.ReadFile(cfgfile) + if err == nil { + setupath = helper.BytesToString(b) + logrus.Println("[nsetu] set setu dir to", setupath) + } + } + }() +} + +func scanall(path string) error { + setuclasses = setuclasses[:0] + model := &setuclass{} + root := os.DirFS(path) + return fs.WalkDir(root, "./", func(path string, d fs.DirEntry, err error) error { + if err != nil { + return err + } + if d.IsDir() { + clsn := d.Name() + mu.Lock() + err = db.Create(clsn, model) + setuclasses = append(setuclasses, clsn) + mu.Unlock() + if err == nil { + err = scanclass(root, clsn) + if err != nil { + return err + } + } + } + return err + }) +} + +func scanclass(root fs.FS, clsn string) error { + return fs.WalkDir(root, clsn, func(path string, d fs.DirEntry, err error) error { + if err != nil { + return err + } + if !d.IsDir() { + f, e := os.Open(path) + if e != nil { + return e + } + img, _, e := image.Decode(f) + if e != nil { + return e + } + dh, e := goimagehash.DifferenceHash(img) + if e != nil { + return e + } + mu.Lock() + err = db.Insert(clsn, &setuclass{ImgID: dh.GetHash(), Name: d.Name()}) + mu.Unlock() + } + return err + }) +} diff --git a/plugin_nativesetu/main.go b/plugin_nativesetu/main.go new file mode 100644 index 00000000..04a10b8b --- /dev/null +++ b/plugin_nativesetu/main.go @@ -0,0 +1,87 @@ +package nativesetu + +import ( + "fmt" + "os" + + "github.com/FloatTech/ZeroBot-Plugin/control" + "github.com/FloatTech/ZeroBot-Plugin/utils/rule" + zero "github.com/wdvxdr1123/ZeroBot" + "github.com/wdvxdr1123/ZeroBot/message" + "github.com/wdvxdr1123/ZeroBot/utils/helper" +) + +const ( + datapath = "data/nsetu" + dbfile = datapath + "/data.db" + cfgfile = datapath + "/setupath.txt" +) + +var ( + setupath = "/tmp" // 绝对路径，图片根目录 +) + +func init() { + engine := control.Register("nativesetu", &control.Options{ + DisableOnDefault: false, + Help: "本地涩图\n" + + "- 来份本地[xxx]\n" + + "- 刷新本地[xxx]\n" + + "- 设置本地setu绝对路径[xxx]\n" + + "- 刷新所有本地setu\n" + + "- 所有本地setu分类", + }) + engine.OnRegex(`^来份本地(.*)$`, rule.FirstValueInList(setuclasses)).SetBlock(true).SetPriority(20). + Handle(func(ctx *zero.Ctx) { + imgtype := ctx.State["regex_matched"].([]string)[1] + sc := new(setuclass) + mu.RLock() + err := db.Pick(imgtype, sc) + mu.RUnlock() + if err != nil { + ctx.SendChain(message.Text("ERROR: ", err)) + } else { + p := "file:///" + setupath + "/" + imgtype + "/" + sc.Name + ctx.SendChain(message.Text(imgtype, ": ", sc.Name, "\n"), message.Image(p)) + } + }) + engine.OnRegex(`^刷新本地(.*)$`, rule.FirstValueInList(setuclasses), zero.SuperUserPermission).SetBlock(true).SetPriority(20). + Handle(func(ctx *zero.Ctx) { + imgtype := ctx.State["regex_matched"].([]string)[1] + err := scanclass(os.DirFS(setupath), imgtype) + if err == nil { + ctx.SendChain(message.Text("成功！")) + } else { + ctx.SendChain(message.Text("ERROR: ", err)) + } + }) + engine.OnRegex(`^设置本地setu绝对路径(.*)$`, zero.SuperUserPermission).SetBlock(true).SetPriority(20). + Handle(func(ctx *zero.Ctx) { + setupath = ctx.State["regex_matched"].([]string)[1] + err := os.WriteFile(cfgfile, helper.StringToBytes(setupath), 0644) + if err != nil { + ctx.SendChain(message.Text("ERROR: ", err)) + } + }) + engine.OnFullMatch("刷新所有本地setu", zero.SuperUserPermission).SetBlock(true).SetPriority(20). + Handle(func(ctx *zero.Ctx) { + err := scanall(setupath) + if err != nil { + ctx.SendChain(message.Text("ERROR: ", err)) + } + }) + engine.OnFullMatch("所有本地setu分类").SetBlock(true).SetPriority(20). + Handle(func(ctx *zero.Ctx) { + msg := "所有本地setu分类" + mu.RLock() + for i, c := range setuclasses { + n, err := db.Count(c) + if err == nil { + msg += fmt.Sprintf("\n%02d. %s(%d)", i, c, n) + } else { + msg += fmt.Sprintf("\n%02d. %s(error)", i, c) + } + } + mu.RUnlock() + }) +} diff --git a/plugin_setutime/setu_geter.go b/plugin_setutime/setu_geter.go index 6c9b478a..bb0b6a0c 100644 --- a/plugin_setutime/setu_geter.go +++ b/plugin_setutime/setu_geter.go @@ -20,6 +20,7 @@ import ( "github.com/FloatTech/ZeroBot-Plugin/control" fileutil "github.com/FloatTech/ZeroBot-Plugin/utils/file" "github.com/FloatTech/ZeroBot-Plugin/utils/math" + "github.com/FloatTech/ZeroBot-Plugin/utils/rule" "github.com/FloatTech/ZeroBot-Plugin/utils/sql" ) @@ -112,7 +113,7 @@ func init() { // 插件主体 "- 删除[涩图/二次元/风景/车万][P站图片ID]\n" + "- >setu status", }) - engine.OnRegex(`^来份(.*)$`, firstValueInList(pool.List)).SetBlock(true).SetPriority(20). + engine.OnRegex(`^来份(.*)$`, rule.FirstValueInList(pool.List)).SetBlock(true).SetPriority(20). Handle(func(ctx *zero.Ctx) { if !limit.Load(ctx.Event.UserID).Acquire() { ctx.SendChain(message.Text("请稍后重试0x0...")) @@ -125,7 +126,7 @@ func init() { // 插件主体 for i := 0; i < times; i++ { illust := &pixiv.Illust{} // 查询出一张图片 - if err := pool.DB.Find(imgtype, illust, "ORDER BY RANDOM() limit 1"); err != nil { + if err := pool.DB.Pick(imgtype, illust); err != nil { ctx.SendChain(message.Text("ERROR: ", err)) continue } @@ -155,7 +156,7 @@ func init() { // 插件主体 } }) - engine.OnRegex(`^添加(.*?)(\d+)$`, firstValueInList(pool.List), zero.SuperUserPermission).SetBlock(true).SetPriority(21). + engine.OnRegex(`^添加(.*?)(\d+)$`, rule.FirstValueInList(pool.List), zero.SuperUserPermission).SetBlock(true).SetPriority(21). Handle(func(ctx *zero.Ctx) { var ( imgtype = ctx.State["regex_matched"].([]string)[1] @@ -186,7 +187,7 @@ func init() { // 插件主体 ctx.SendChain(message.Text("添加成功")) }) - engine.OnRegex(`^删除(.*?)(\d+)$`, firstValueInList(pool.List), zero.SuperUserPermission).SetBlock(true).SetPriority(22). + engine.OnRegex(`^删除(.*?)(\d+)$`, rule.FirstValueInList(pool.List), zero.SuperUserPermission).SetBlock(true).SetPriority(22). Handle(func(ctx *zero.Ctx) { var ( imgtype = ctx.State["regex_matched"].([]string)[1] @@ -218,19 +219,6 @@ func init() { // 插件主体 }) } -// firstValueInList 判断正则匹配的第一个参数是否在列表中 -func firstValueInList(list []string) zero.Rule { - return func(ctx *zero.Ctx) bool { - first := ctx.State["regex_matched"].([]string)[1] - for i := range list { - if first == list[i] { - return true - } - } - return false - } -} - // size 返回缓冲池指定类型的现有大小 func (p *imgpool) size(imgtype string) int { return len(p.Pool[imgtype]) diff --git a/utils/rule/extension.go b/utils/rule/extension.go new file mode 100644 index 00000000..5580d649 --- /dev/null +++ b/utils/rule/extension.go @@ -0,0 +1,16 @@ +package rule + +import zero "github.com/wdvxdr1123/ZeroBot" + +// FirstValueInList 判断正则匹配的第一个参数是否在列表中 +func FirstValueInList(list []string) zero.Rule { + return func(ctx *zero.Ctx) bool { + first := ctx.State["regex_matched"].([]string)[1] + for i := range list { + if first == list[i] { + return true + } + } + return false + } +} diff --git a/utils/sql/sqlite.go b/utils/sql/sqlite.go index 54746484..ab44e059 100644 --- a/utils/sql/sqlite.go +++ b/utils/sql/sqlite.go @@ -143,6 +143,11 @@ func (db *Sqlite) Find(table string, objptr interface{}, condition string) error return err } +// Pick 从 table 随机一行 +func (db *Sqlite) Pick(table string, objptr interface{}) error { + return db.Find(table, objptr, "ORDER BY RANDOM() limit 1") +} + // ListTables 列出所有表名 // 返回所有表名+错误 func (db *Sqlite) ListTables() (s []string, err error) {