mirror of
https://github.com/FloatTech/ZeroBot-Plugin.git
synced 2026-02-12 10:10:25 +00:00
🎨⚡️✨ 优化代码
This commit is contained in:
Yiwen-Chan
259
plugin_setutime/setu_geter.go
Normal file
259
plugin_setutime/setu_geter.go
Normal file
@@ -0,0 +1,259 @@
|
||||
package plugin_setutime
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"os"
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
zero "github.com/wdvxdr1123/ZeroBot"
|
||||
"github.com/wdvxdr1123/ZeroBot/extension/rate"
|
||||
"github.com/wdvxdr1123/ZeroBot/message"
|
||||
|
||||
"github.com/FloatTech/AnimeAPI/pixiv"
|
||||
)
|
||||
|
||||
// Pools 图片缓冲池
|
||||
type Pool struct {
|
||||
Lock sync.Mutex
|
||||
DB *Sqlite
|
||||
Path string
|
||||
Group int64
|
||||
List []string
|
||||
Max int
|
||||
Pool map[string][]*pixiv.Illust
|
||||
Form int64
|
||||
}
|
||||
|
||||
// NewPoolsCache 返回一个缓冲池对象
|
||||
func NewPools() *Pool {
|
||||
cache := &Pool{
|
||||
DB: &Sqlite{DBPath: "data/SetuTime/SetuTime.db"},
|
||||
Path: "data/SetuTime/cache/",
|
||||
Group: 0,
|
||||
List: []string{"涩图", "二次元", "风景", "车万"}, // 可以自己加类别，得自己加图片进数据库
|
||||
Max: 10,
|
||||
Pool: map[string][]*pixiv.Illust{},
|
||||
Form: 0,
|
||||
}
|
||||
err := os.MkdirAll(cache.Path, 0644)
|
||||
if err != nil {
|
||||
panic(err)
|
||||
}
|
||||
for i := range cache.List {
|
||||
if err := cache.DB.Create(cache.List[i], &pixiv.Illust{}); err != nil {
|
||||
panic(err)
|
||||
}
|
||||
}
|
||||
return cache
|
||||
}
|
||||
|
||||
var (
|
||||
POOL = NewPools()
|
||||
limit = rate.NewManager(time.Minute*1, 5)
|
||||
)
|
||||
|
||||
func init() { // 插件主体
|
||||
zero.OnRegex(`^来份(.*)$`, FirstValueInList(POOL.List)).SetBlock(true).SetPriority(20).
|
||||
Handle(func(ctx *zero.Ctx) {
|
||||
if !limit.Load(ctx.Event.UserID).Acquire() {
|
||||
ctx.SendChain(message.Text("少女祈祷中......"))
|
||||
return
|
||||
}
|
||||
var type_ = ctx.State["regex_matched"].([]string)[1]
|
||||
// 补充池子
|
||||
go func() {
|
||||
times := Min(POOL.Max-POOL.Size(type_), 2)
|
||||
for i := 0; i < times; i++ {
|
||||
illust := &pixiv.Illust{}
|
||||
// 查询出一张图片
|
||||
if err := POOL.DB.Select(type_, illust, "ORDER BY RANDOM() limit 1"); err != nil {
|
||||
ctx.SendChain(message.Text("ERROR: ", err))
|
||||
continue
|
||||
}
|
||||
// 下载图片
|
||||
if _, err := download(illust, POOL.Path); err != nil {
|
||||
ctx.SendChain(message.Text("ERROR: ", err))
|
||||
continue
|
||||
}
|
||||
ctx.SendGroupMessage(POOL.Group, []message.MessageSegment{message.Image(file(illust))})
|
||||
// 向缓冲池添加一张图片
|
||||
POOL.Push(type_, illust)
|
||||
|
||||
time.Sleep(time.Second * 1)
|
||||
}
|
||||
}()
|
||||
// 如果没有缓存，阻塞5秒
|
||||
if POOL.Size(type_) == 0 {
|
||||
ctx.SendChain(message.Text("INFO: 正在填充弹药......"))
|
||||
<-time.After(time.Second * 5)
|
||||
if POOL.Size(type_) == 0 {
|
||||
ctx.SendChain(message.Text("ERROR: 等待填充，请稍后再试......"))
|
||||
return
|
||||
}
|
||||
}
|
||||
// 从缓冲池里抽一张
|
||||
if id := ctx.SendChain(message.Image(file(POOL.Pop(type_)))); id == 0 {
|
||||
ctx.SendChain(message.Text("ERROR: 可能被风控了"))
|
||||
}
|
||||
return
|
||||
})
|
||||
|
||||
zero.OnRegex(`^添加(.*?)(\d+)$`, FirstValueInList(POOL.List), zero.SuperUserPermission).SetBlock(true).SetPriority(21).
|
||||
Handle(func(ctx *zero.Ctx) {
|
||||
var (
|
||||
type_ = 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)
|
||||
if err != nil {
|
||||
ctx.SendChain(message.Text("ERROR: ", err))
|
||||
return
|
||||
}
|
||||
// 下载插画
|
||||
if _, err := download(illust, POOL.Path); err != nil {
|
||||
ctx.SendChain(message.Text("ERROR: ", err))
|
||||
return
|
||||
}
|
||||
// 发送到发送者
|
||||
if id := ctx.SendChain(message.Image(file(illust))); id == 0 {
|
||||
ctx.SendChain(message.Text("ERROR: 可能被风控，发送失败"))
|
||||
return
|
||||
}
|
||||
// 添加插画到对应的数据库table
|
||||
if err := POOL.DB.Insert(type_, illust); err != nil {
|
||||
ctx.SendChain(message.Text("ERROR: ", err))
|
||||
return
|
||||
}
|
||||
ctx.Send("添加成功")
|
||||
return
|
||||
})
|
||||
|
||||
zero.OnRegex(`^删除(.*?)(\d+)$`, FirstValueInList(POOL.List), zero.SuperUserPermission).SetBlock(true).SetPriority(22).
|
||||
Handle(func(ctx *zero.Ctx) {
|
||||
var (
|
||||
type_ = ctx.State["regex_matched"].([]string)[1]
|
||||
id, _ = strconv.ParseInt(ctx.State["regex_matched"].([]string)[2], 10, 64)
|
||||
)
|
||||
// 查询数据库
|
||||
if err := POOL.DB.Delete(type_, fmt.Sprintf("WHERE pid=%d", id)); err != nil {
|
||||
ctx.Send(fmt.Sprintf("ERROR: %v", err))
|
||||
return
|
||||
}
|
||||
ctx.Send("删除成功")
|
||||
return
|
||||
})
|
||||
|
||||
// 查询数据库涩图数量
|
||||
zero.OnFullMatchGroup([]string{">setu status"}).SetBlock(true).SetPriority(23).
|
||||
Handle(func(ctx *zero.Ctx) {
|
||||
state := []string{"[SetuTime]"}
|
||||
for i := range POOL.List {
|
||||
num, err := POOL.DB.Num(POOL.List[i])
|
||||
if err != nil {
|
||||
num = 0
|
||||
}
|
||||
state = append(state, "\n")
|
||||
state = append(state, POOL.List[i])
|
||||
state = append(state, ": ")
|
||||
state = append(state, fmt.Sprintf("%d", num))
|
||||
}
|
||||
ctx.Send(strings.Join(state, ""))
|
||||
return
|
||||
})
|
||||
}
|
||||
|
||||
// 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
|
||||
}
|
||||
}
|
||||
|
||||
// Min 返回两数最小值
|
||||
func Min(a, b int) int {
|
||||
switch {
|
||||
default:
|
||||
return a
|
||||
case a > b:
|
||||
return b
|
||||
case a < b:
|
||||
return a
|
||||
}
|
||||
}
|
||||
|
||||
// Size 返回缓冲池指定类型的现有大小
|
||||
func (p *Pool) Size(type_ string) int {
|
||||
return len(p.Pool[type_])
|
||||
}
|
||||
|
||||
// IsFull 返回缓冲池指定类型是否已满
|
||||
func (p *Pool) IsFull(type_ string) bool {
|
||||
return len(p.Pool[type_]) >= p.Max
|
||||
}
|
||||
|
||||
// Push 向缓冲池插入一张图片
|
||||
func (p *Pool) Push(type_ string, illust *pixiv.Illust) {
|
||||
p.Lock.Lock()
|
||||
defer p.Lock.Unlock()
|
||||
p.Pool[type_] = append(p.Pool[type_], illust)
|
||||
}
|
||||
|
||||
// Push 在缓冲池拿出一张图片
|
||||
func (p *Pool) Pop(type_ string) (illust *pixiv.Illust) {
|
||||
p.Lock.Lock()
|
||||
defer p.Lock.Unlock()
|
||||
if p.Size(type_) == 0 {
|
||||
return
|
||||
}
|
||||
illust = p.Pool[type_][0]
|
||||
p.Pool[type_] = p.Pool[type_][1:]
|
||||
return
|
||||
}
|
||||
|
||||
func file(i *pixiv.Illust) string {
|
||||
filename := fmt.Sprint(i.Pid)
|
||||
pwd, _ := os.Getwd()
|
||||
filepath := pwd + `/` + POOL.Path + filename
|
||||
if _, err := os.Stat(filepath + ".jpg"); err == nil || os.IsExist(err) {
|
||||
return `file:///` + filepath + ".jpg"
|
||||
}
|
||||
if _, err := os.Stat(filepath + ".png"); err == nil || os.IsExist(err) {
|
||||
return `file:///` + filepath + ".png"
|
||||
}
|
||||
if _, err := os.Stat(filepath + ".gif"); err == nil || os.IsExist(err) {
|
||||
return `file:///` + filepath + ".gif"
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
func download(i *pixiv.Illust, filedir string) (string, error) {
|
||||
filename := fmt.Sprint(i.Pid)
|
||||
filepath := filedir + filename
|
||||
if _, err := os.Stat(filepath + ".jpg"); err == nil || os.IsExist(err) {
|
||||
return filepath + ".jpg", nil
|
||||
}
|
||||
if _, err := os.Stat(filepath + ".png"); err == nil || os.IsExist(err) {
|
||||
return filepath + ".png", nil
|
||||
}
|
||||
if _, err := os.Stat(filepath + ".gif"); err == nil || os.IsExist(err) {
|
||||
return filepath + ".gif", nil
|
||||
}
|
||||
// 下载最大分辨率为 1200 的图片
|
||||
link := i.ImageUrls
|
||||
link = strings.ReplaceAll(link, "img-original", "img-master")
|
||||
link = strings.ReplaceAll(link, "_p0", "_p0_master1200")
|
||||
link = strings.ReplaceAll(link, ".png", ".jpg")
|
||||
// 下载
|
||||
return pixiv.Download(link, filedir, filename)
|
||||
}
|
||||
245
plugin_setutime/sqlite.go
Normal file
245
plugin_setutime/sqlite.go
Normal file
@@ -0,0 +1,245 @@
|
||||
package plugin_setutime
|
||||
|
||||
import (
|
||||
"database/sql"
|
||||
"errors"
|
||||
"reflect"
|
||||
"strings"
|
||||
|
||||
_ "modernc.org/sqlite"
|
||||
)
|
||||
|
||||
// Sqlite 数据库对象
|
||||
type Sqlite struct {
|
||||
DB *sql.DB
|
||||
DBPath string
|
||||
}
|
||||
|
||||
// Create 生成数据库
|
||||
// 默认结构体的第一个元素为主键
|
||||
// 返回错误
|
||||
func (db *Sqlite) Create(table string, objptr interface{}) (err error) {
|
||||
if db.DB == nil {
|
||||
database, err := sql.Open("sqlite", db.DBPath)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
db.DB = database
|
||||
}
|
||||
var (
|
||||
tags = tags(objptr)
|
||||
kinds = kinds(objptr)
|
||||
top = len(tags) - 1
|
||||
cmd = []string{}
|
||||
)
|
||||
cmd = append(cmd, "CREATE TABLE IF NOT EXISTS")
|
||||
cmd = append(cmd, table)
|
||||
cmd = append(cmd, "(")
|
||||
for i := range tags {
|
||||
cmd = append(cmd, tags[i])
|
||||
cmd = append(cmd, kinds[i])
|
||||
switch i {
|
||||
default:
|
||||
cmd = append(cmd, "NULL,")
|
||||
case 0:
|
||||
cmd = append(cmd, "PRIMARY KEY")
|
||||
cmd = append(cmd, "NOT NULL,")
|
||||
case top:
|
||||
cmd = append(cmd, "NULL);")
|
||||
}
|
||||
}
|
||||
if _, err := db.DB.Exec(strings.Join(cmd, " ")); err != nil {
|
||||
return err
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// Insert 插入数据集
|
||||
// 默认结构体的第一个元素为主键
|
||||
// 返回错误
|
||||
func (db *Sqlite) Insert(table string, objptr interface{}) (err error) {
|
||||
rows, err := db.DB.Query("SELECT * FROM " + table)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
tags, _ := rows.Columns()
|
||||
rows.Close()
|
||||
var (
|
||||
values = values(objptr)
|
||||
top = len(tags) - 1
|
||||
cmd = []string{}
|
||||
)
|
||||
cmd = append(cmd, "INSERT INTO")
|
||||
cmd = append(cmd, table)
|
||||
for i := range tags {
|
||||
switch i {
|
||||
default:
|
||||
cmd = append(cmd, tags[i])
|
||||
cmd = append(cmd, ",")
|
||||
case 0:
|
||||
cmd = append(cmd, "(")
|
||||
cmd = append(cmd, tags[i])
|
||||
cmd = append(cmd, ",")
|
||||
case top:
|
||||
cmd = append(cmd, tags[i])
|
||||
cmd = append(cmd, ")")
|
||||
}
|
||||
}
|
||||
for i := range tags {
|
||||
switch i {
|
||||
default:
|
||||
cmd = append(cmd, "?")
|
||||
cmd = append(cmd, ",")
|
||||
case 0:
|
||||
cmd = append(cmd, "VALUES (")
|
||||
cmd = append(cmd, "?")
|
||||
cmd = append(cmd, ",")
|
||||
case top:
|
||||
cmd = append(cmd, "?")
|
||||
cmd = append(cmd, ")")
|
||||
}
|
||||
}
|
||||
stmt, err := db.DB.Prepare(strings.Join(cmd, " "))
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
_, err = stmt.Exec(values...)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// Select 查询数据库
|
||||
// condition 可为"WHERE id = 0"
|
||||
// 默认字段与结构体元素顺序一致
|
||||
// 返回错误
|
||||
func (db *Sqlite) Select(table string, objptr interface{}, condition string) (err error) {
|
||||
var cmd = []string{}
|
||||
cmd = append(cmd, "SELECT * FROM ")
|
||||
cmd = append(cmd, table)
|
||||
cmd = append(cmd, condition)
|
||||
rows, err := db.DB.Query(strings.Join(cmd, " "))
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer rows.Close()
|
||||
|
||||
for rows.Next() {
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
err = rows.Scan(addrs(objptr)...)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
return nil
|
||||
}
|
||||
return errors.New("数据库无此条件项目")
|
||||
}
|
||||
|
||||
// Delete 删除数据库
|
||||
// condition 可为"WHERE id = 0"
|
||||
// 返回错误
|
||||
func (db *Sqlite) Delete(table string, condition string) (err error) {
|
||||
var cmd = []string{}
|
||||
cmd = append(cmd, "DELETE FROM")
|
||||
cmd = append(cmd, table)
|
||||
cmd = append(cmd, condition)
|
||||
stmt, err := db.DB.Prepare(strings.Join(cmd, " "))
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
_, err = stmt.Exec()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// Num 查询数据库行数
|
||||
// 返回行数以及错误
|
||||
func (db *Sqlite) Num(table string) (num int, err error) {
|
||||
var cmd = []string{}
|
||||
cmd = append(cmd, "SELECT * FROM")
|
||||
cmd = append(cmd, table)
|
||||
rows, err := db.DB.Query(strings.Join(cmd, " "))
|
||||
if err != nil {
|
||||
return num, err
|
||||
}
|
||||
defer rows.Close()
|
||||
for rows.Next() {
|
||||
num++
|
||||
}
|
||||
return num, nil
|
||||
}
|
||||
|
||||
// tags 反射 返回结构体对象的 tag 数组
|
||||
func tags(objptr interface{}) []string {
|
||||
var tags []string
|
||||
elem := reflect.ValueOf(objptr).Elem()
|
||||
// 判断第一个元素是否为匿名字段
|
||||
if elem.Type().Field(0).Anonymous {
|
||||
elem = elem.Field(0)
|
||||
}
|
||||
for i, flen := 0, elem.Type().NumField(); i < flen; i++ {
|
||||
tags = append(tags, elem.Type().Field(i).Tag.Get("db"))
|
||||
}
|
||||
return tags
|
||||
}
|
||||
|
||||
// kinds 反射 返回结构体对象的 kinds 数组
|
||||
func kinds(objptr interface{}) []string {
|
||||
var kinds []string
|
||||
elem := reflect.ValueOf(objptr).Elem()
|
||||
// 判断第一个元素是否为匿名字段
|
||||
if elem.Type().Field(0).Anonymous {
|
||||
elem = elem.Field(0)
|
||||
}
|
||||
for i, flen := 0, elem.Type().NumField(); i < flen; i++ {
|
||||
switch elem.Field(i).Type().String() {
|
||||
case "int64":
|
||||
kinds = append(kinds, "INT")
|
||||
case "string":
|
||||
kinds = append(kinds, "TEXT")
|
||||
default:
|
||||
kinds = append(kinds, "TEXT")
|
||||
}
|
||||
}
|
||||
return kinds
|
||||
}
|
||||
|
||||
// values 反射 返回结构体对象的 values 数组
|
||||
func values(objptr interface{}) []interface{} {
|
||||
var values []interface{}
|
||||
elem := reflect.ValueOf(objptr).Elem()
|
||||
// 判断第一个元素是否为匿名字段
|
||||
if elem.Type().Field(0).Anonymous {
|
||||
elem = elem.Field(0)
|
||||
}
|
||||
for i, flen := 0, elem.Type().NumField(); i < flen; i++ {
|
||||
switch elem.Field(i).Type().String() {
|
||||
case "int64":
|
||||
values = append(values, elem.Field(i).Int())
|
||||
case "string":
|
||||
values = append(values, elem.Field(i).String())
|
||||
default:
|
||||
values = append(values, elem.Field(i).String())
|
||||
}
|
||||
}
|
||||
return values
|
||||
}
|
||||
|
||||
// addrs 反射 返回结构体对象的 addrs 数组
|
||||
func addrs(objptr interface{}) []interface{} {
|
||||
var addrs []interface{}
|
||||
elem := reflect.ValueOf(objptr).Elem()
|
||||
// 判断第一个元素是否为匿名字段
|
||||
if elem.Type().Field(0).Anonymous {
|
||||
elem = elem.Field(0)
|
||||
}
|
||||
for i, flen := 0, elem.Type().NumField(); i < flen; i++ {
|
||||
addrs = append(addrs, elem.Field(i).Addr().Interface())
|
||||
}
|
||||
return addrs
|
||||
}
|
||||
Reference in New Issue
Block a user