mirror of
https://github.com/FloatTech/ZeroBot-Plugin.git
synced 2026-02-12 02:00:24 +00:00
✨ 新增 本地 setu 插件
This commit is contained in:
fumiama
96
plugin_nativesetu/data.go
Normal file
96
plugin_nativesetu/data.go
Normal file
@@ -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
|
||||
})
|
||||
}
|
||||
87
plugin_nativesetu/main.go
Normal file
87
plugin_nativesetu/main.go
Normal file
@@ -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()
|
||||
})
|
||||
}
|
||||
Reference in New Issue
Block a user