✨ 新增 本地 setu 插件

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"`

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

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=

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"      // 搜番

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
+	})
+}

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()
+		})
+}

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])

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
+	}
+}

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) {