diff --git a/plugin_setutime/setu_geter.go b/plugin_setutime/setu_geter.go index 2561d61d..25132ce3 100644 --- a/plugin_setutime/setu_geter.go +++ b/plugin_setutime/setu_geter.go @@ -2,6 +2,7 @@ package setutime import ( + "errors" "fmt" "strconv" "strings" @@ -26,43 +27,24 @@ import ( // Pools 图片缓冲池 type imgpool struct { - Lock sync.Mutex - DB *sql.Sqlite - Path string - Group int64 - List []string - Max int - Pool map[string][]*pixiv.Illust - Form int64 + lock sync.Mutex + db *sql.Sqlite + path string + max int + pool map[string][]*pixiv.Illust } -// NewPoolsCache 返回一个缓冲池对象 -func newPools() *imgpool { - cache := &imgpool{ - DB: &sql.Sqlite{DBPath: "data/SetuTime/SetuTime.db"}, - Path: pixiv.CacheDir, - Group: 0, - List: []string{"涩图", "二次元", "风景", "车万"}, // 可以自己加类别,得自己加图片进数据库 - Max: 10, - Pool: map[string][]*pixiv.Illust{}, - Form: 0, +func (i *imgpool) list() (l []string) { + var err error + l, err = i.db.ListTables() + if err != nil { + l = []string{"涩图", "二次元", "风景", "车万"} } - // 如果数据库不存在则下载 - _, _ = fileutil.GetLazyData(cache.DB.DBPath, false, false) - for i := range cache.List { - if err := cache.DB.Create(cache.List[i], &pixiv.Illust{}); err != nil { - panic(err) - } - } - return cache + return l } -var ( - pool *imgpool - limit = rate.NewManager(time.Minute*1, 5) -) - func init() { // 插件主体 + limit := rate.NewManager(time.Minute*1, 5) engine := control.Register("setutime", order.PrioSetuTime, &control.Options{ DisableOnDefault: false, Help: "涩图\n" + @@ -72,8 +54,27 @@ func init() { // 插件主体 "- >setu status", }) process.SleepAbout1sTo2s() - pool = newPools() - engine.OnRegex(`^来份(.*)$`, rule.FirstValueInList(pool.List)).SetBlock(true). + pool := func() *imgpool { + cache := &imgpool{ + db: &sql.Sqlite{DBPath: "data/SetuTime/SetuTime.db"}, + path: pixiv.CacheDir, + max: 10, + pool: map[string][]*pixiv.Illust{}, + } + // 如果数据库不存在则下载 + _, _ = fileutil.GetLazyData(cache.db.DBPath, false, false) + err := cache.db.Open() + if err != nil { + panic(err) + } + for _, imgtype := range cache.list() { + if err := cache.db.Create(imgtype, &pixiv.Illust{}); err != nil { + panic(err) + } + } + return cache + }() + engine.OnRegex(`^来份(.*)$`, rule.FirstValueInList(pool.list())).SetBlock(true). Handle(func(ctx *zero.Ctx) { if !limit.Load(ctx.Event.UserID).Acquire() { ctx.SendChain(message.Text("请稍后重试0x0...")) @@ -81,20 +82,7 @@ func init() { // 插件主体 } var imgtype = ctx.State["regex_matched"].([]string)[1] // 补充池子 - go func() { - times := math.Min(pool.Max-pool.size(imgtype), 2) - for i := 0; i < times; i++ { - illust := &pixiv.Illust{} - // 查询出一张图片 - if err := pool.DB.Pick(imgtype, illust); err != nil { - ctx.SendChain(message.Text("ERROR: ", err)) - continue - } - // 向缓冲池添加一张图片 - pool.push(ctx, imgtype, illust) - process.SleepAbout1sTo2s() - } - }() + go pool.fill(ctx, imgtype) // 如果没有缓存,阻塞10秒 if pool.size(imgtype) == 0 { ctx.SendChain(message.Text("INFO: 正在填充弹药......")) @@ -105,50 +93,33 @@ func init() { // 插件主体 } } // 从缓冲池里抽一张 - if id := ctx.SendChain(message.Image(file(pool.pop(imgtype)))); id.ID() == 0 { + if id := ctx.SendChain(message.Image(pool.popfile(imgtype))); id.ID() == 0 { ctx.SendChain(message.Text("ERROR: 可能被风控了")) } }) - engine.OnRegex(`^添加(.*?)(\d+)$`, rule.FirstValueInList(pool.List), zero.SuperUserPermission).SetBlock(true). + engine.OnRegex(`^添加(.*?)(\d+)$`, zero.SuperUserPermission).SetBlock(true). Handle(func(ctx *zero.Ctx) { var ( imgtype = ctx.State["regex_matched"].([]string)[1] id, _ = strconv.ParseInt(ctx.State["regex_matched"].([]string)[2], 10, 64) ) - ctx.SendChain(message.Text("少女祈祷中......")) - // 查询P站插图信息 - illust, err := pixiv.Works(id) + err := pool.add(ctx, imgtype, id) if err != nil { - ctx.SendChain(message.Text("ERROR: ", err)) + ctx.SendChain(message.Text("ERROR:", err)) return } - // 下载插画 - if _, err := illust.DownloadToCache(0, strconv.FormatInt(id, 10)+"_p0"); err != nil { - ctx.SendChain(message.Text("ERROR: ", err)) - return - } - // 发送到发送者 - if id := ctx.SendChain(message.Image(file(illust))); id.ID() == 0 { - ctx.SendChain(message.Text("ERROR: 可能被风控,发送失败")) - return - } - // 添加插画到对应的数据库table - if err := pool.DB.Insert(imgtype, illust); err != nil { - ctx.SendChain(message.Text("ERROR: ", err)) - return - } - ctx.SendChain(message.Text("添加成功")) + ctx.SendChain(message.Text("成功向分类", imgtype, "添加图片", id)) }) - engine.OnRegex(`^删除(.*?)(\d+)$`, rule.FirstValueInList(pool.List), zero.SuperUserPermission).SetBlock(true). + engine.OnRegex(`^删除(.*?)(\d+)$`, rule.FirstValueInList(pool.list()), zero.SuperUserPermission).SetBlock(true). Handle(func(ctx *zero.Ctx) { var ( imgtype = ctx.State["regex_matched"].([]string)[1] id, _ = strconv.ParseInt(ctx.State["regex_matched"].([]string)[2], 10, 64) ) // 查询数据库 - if err := pool.DB.Del(imgtype, fmt.Sprintf("WHERE pid=%d", id)); err != nil { + if err := pool.remove(imgtype, id); err != nil { ctx.SendChain(message.Text("ERROR: ", err)) return } @@ -159,13 +130,13 @@ func init() { // 插件主体 engine.OnFullMatchGroup([]string{">setu status"}).SetBlock(true). Handle(func(ctx *zero.Ctx) { state := []string{"[SetuTime]"} - for i := range pool.List { - num, err := pool.DB.Count(pool.List[i]) + for _, imgtype := range pool.list() { + num, err := pool.db.Count(imgtype) if err != nil { num = 0 } state = append(state, "\n") - state = append(state, pool.List[i]) + state = append(state, imgtype) state = append(state, ": ") state = append(state, fmt.Sprintf("%d", num)) } @@ -175,16 +146,9 @@ func init() { // 插件主体 // size 返回缓冲池指定类型的现有大小 func (p *imgpool) size(imgtype string) int { - return len(p.Pool[imgtype]) + return len(p.pool[imgtype]) } -/* -// isFull 返回缓冲池指定类型是否已满 -func (p *imgpool) isFull(imgtype string) bool { - return len(p.Pool[imgtype]) >= p.Max -}*/ - -// push 向缓冲池插入一张图片 func (p *imgpool) push(ctx *zero.Ctx, imgtype string, illust *pixiv.Illust) { u := illust.ImageUrls[0] n := u[strings.LastIndex(u, "/")+1 : len(u)-4] @@ -199,31 +163,30 @@ func (p *imgpool) push(ctx *zero.Ctx, imgtype string, illust *pixiv.Illust) { m.SetFile(fileutil.BOTPATH + "/" + f) _, _ = m.Push(ctxext.SendToSelf(ctx), ctxext.GetMessage(ctx)) } - p.Lock.Lock() - p.Pool[imgtype] = append(p.Pool[imgtype], illust) - p.Lock.Unlock() + p.lock.Lock() + p.pool[imgtype] = append(p.pool[imgtype], illust) + p.lock.Unlock() } -// Push 在缓冲池拿出一张图片 func (p *imgpool) pop(imgtype string) (illust *pixiv.Illust) { - p.Lock.Lock() - defer p.Lock.Unlock() + p.lock.Lock() + defer p.lock.Unlock() if p.size(imgtype) == 0 { return } - illust = p.Pool[imgtype][0] - p.Pool[imgtype] = p.Pool[imgtype][1:] + illust = p.pool[imgtype][0] + p.pool[imgtype] = p.pool[imgtype][1:] return } -func file(i *pixiv.Illust) string { +func (p *imgpool) file(i *pixiv.Illust) string { u := i.ImageUrls[0] m, err := imagepool.GetImage(u[strings.LastIndex(u, "/")+1 : len(u)-4]) if err == nil { return m.String() } filename := fmt.Sprint(i.Pid) + "_p0" - filepath := fileutil.BOTPATH + `/` + pool.Path + filename + filepath := fileutil.BOTPATH + `/` + p.path + filename if fileutil.IsExist(filepath + ".jpg") { return `file:///` + filepath + ".jpg" } @@ -235,3 +198,52 @@ func file(i *pixiv.Illust) string { } return "" } + +func (p *imgpool) popfile(imgtype string) string { + return p.file(p.pop(imgtype)) +} + +// fill 补充池子 +func (p *imgpool) fill(ctx *zero.Ctx, imgtype string) { + times := math.Min(p.max-p.size(imgtype), 2) + for i := 0; i < times; i++ { + illust := &pixiv.Illust{} + // 查询出一张图片 + if err := p.db.Pick(imgtype, illust); err != nil { + ctx.SendChain(message.Text("ERROR: ", err)) + continue + } + // 向缓冲池添加一张图片 + p.push(ctx, imgtype, illust) + process.SleepAbout1sTo2s() + } +} + +func (p *imgpool) add(ctx *zero.Ctx, imgtype string, id int64) error { + if err := p.db.Create(imgtype, &pixiv.Illust{}); err != nil { + return err + } + ctx.SendChain(message.Text("少女祈祷中......")) + // 查询P站插图信息 + illust, err := pixiv.Works(id) + if err != nil { + return err + } + // 下载插画 + if _, err := illust.DownloadToCache(0, strconv.FormatInt(id, 10)+"_p0"); err != nil { + return err + } + // 发送到发送者 + if id := ctx.SendChain(message.Image(p.file(illust))); id.ID() == 0 { + return errors.New("可能被风控,发送失败") + } + // 添加插画到对应的数据库table + if err := p.db.Insert(imgtype, illust); err != nil { + return err + } + return nil +} + +func (p *imgpool) remove(imgtype string, id int64) error { + return p.db.Del(imgtype, fmt.Sprintf("WHERE pid=%d", id)) +}