✏️ 优化 vtbquotation

plugin_vtb_quotation/cron.go

@@ -1,51 +1,37 @@
-package plugin_vtb_quotation
+package vtbquotation
 
 import (
 	"github.com/fumiama/cron"
 	log "github.com/sirupsen/logrus"
-	zero "github.com/wdvxdr1123/ZeroBot"
 
-	"github.com/FloatTech/ZeroBot-Plugin/plugin_vtb_quotation/firstVtb"
 	"github.com/FloatTech/ZeroBot-Plugin/plugin_vtb_quotation/model"
-	"github.com/FloatTech/ZeroBot-Plugin/plugin_vtb_quotation/secondVtb"
-)
-
-var (
-	AtriRule = true
 )
 
 func init() {
-	engine.OnMessage(atriRule).SetBlock(false).Handle(func(ctx *zero.Ctx) {
-		log.Println("定时任务只创建一次")
-		AtriRule = false
-		log.Println("开启vtb数据库日常更新")
-		vtbDaily()
-	})
+	log.Println("[vtb/cron] 开启vtb数据库日常更新")
+	vtbDaily()
 }
 
 func vtbDaily() {
-	log.Println("创建vtb数据库定时任务")
 	c := cron.New()
 	_, err := c.AddFunc("0 4 * * *", func() { vtbData() })
 	if err != nil {
-		log.Println("定时任务有错误:", err)
+		log.Errorln("定时任务有错误:", err)
 	} else {
 		log.Println("开启vtb数据库定时任务")
 		c.Start()
 	}
 }
-func vtbData() {
-	model.Init()
-	vtbListStr := firstVtb.GetVtbListStr()
-	uidList := firstVtb.DealVtbListStr(vtbListStr)
-	log.Println(uidList)
-	for _, v := range uidList {
-		vtbStr := secondVtb.GetVtbStr(v)
-		secondVtb.DealVtbStr(vtbStr, v)
-	}
-	model.Db.Close()
-}
 
-func atriRule(ctx *zero.Ctx) bool {
-	return AtriRule
+func vtbData() {
+	db := model.Init(dbpath)
+	if db != nil {
+		for _, v := range db.GetVtbList() {
+			db.StoreVtb(v)
+		}
+		err := db.Close()
+		if err != nil {
+			log.Errorln("[vtb/cron]", err)
+		}
+	}
 }

plugin_vtb_quotation/data.go

@@ -0,0 +1,36 @@
+package vtbquotation
+
+import (
+	"io"
+	"log"
+	"net/http"
+	"os"
+
+	"github.com/FloatTech/ZeroBot-Plugin/utils/file"
+)
+
+const pburl = "https://codechina.csdn.net/u011570312/ZeroBot-Plugin/-/raw/master/" + dbpath
+
+// 加载数据库
+func init() {
+	if !file.IsExist(dbpath) { // 如果没有数据库，则从 url 下载
+		f, err := os.Create(dbpath)
+		if err != nil {
+			panic(err)
+		}
+		defer f.Close()
+		resp, err := http.Get(pburl)
+		if err == nil {
+			defer resp.Body.Close()
+			if resp.ContentLength > 0 {
+				log.Printf("[vtb]从镜像下载数据库%d字节...", resp.ContentLength)
+				data, err := io.ReadAll(resp.Body)
+				if err == nil && len(data) > 0 {
+					_, _ = f.Write(data)
+				}
+				panic(err)
+			}
+		}
+		panic(err)
+	}
+}

plugin_vtb_quotation/firstVtb/firstVtb.go

@@ -1,79 +0,0 @@
-package firstVtb
-
-import (
-	"io/ioutil"
-	"net/http"
-	"strconv"
-	"strings"
-
-	"github.com/jinzhu/gorm"
-	log "github.com/sirupsen/logrus"
-	"github.com/tidwall/gjson"
-
-	"github.com/FloatTech/ZeroBot-Plugin/plugin_vtb_quotation/model"
-	"github.com/FloatTech/ZeroBot-Plugin/plugin_vtb_quotation/utils"
-)
-
-var vtbUrl = "https://vtbkeyboard.moe/api/get_vtb_list"
-
-func GetVtbListStr() string {
-	client := &http.Client{}
-	req, err := http.NewRequest("GET", vtbUrl, nil)
-	if err != nil {
-		log.Println(err)
-	}
-	// 自定义Header
-	req.Header.Set("User-Agent", utils.GetAgent())
-	resp, err := client.Do(req)
-	if err != nil {
-		log.Println(err)
-	}
-
-	defer resp.Body.Close()
-	bytes, err := ioutil.ReadAll(resp.Body)
-	if err != nil {
-		log.Println(err)
-	}
-	//log.Println(string(bytes))
-	vtbListStr, err := strconv.Unquote(strings.Replace(strconv.Quote(string(bytes)), `\\u`, `\u`, -1))
-	if err != nil {
-		log.Println(err)
-	}
-	log.Println(vtbListStr)
-	return vtbListStr
-}
-func DealVtbListStr(vtbListStr string) []string {
-	uidList := make([]string, 0)
-	count := gjson.Get(vtbListStr, "#").Int()
-	for i := int64(0); i < count; i++ {
-		item := gjson.Get(vtbListStr, strconv.FormatInt(i, 10))
-		log.Println(item)
-		fc := model.FirstCategory{
-			FirstCategoryIndex:       i,
-			FirstCategoryName:        item.Get("name").String(),
-			FirstCategoryDescription: item.Get("description").String(),
-			FirstCategoryIconPath:    item.Get("icon_path").String(),
-			FirstCategoryUid:         item.Get("uid").String(),
-		}
-		log.Println(fc)
-		//model.Db.Model(FirstCategory{}).Where("first_category_uid = ?", fc.FirstCategoryUid).FirstOrCreate(&fc)
-		if err := model.Db.Debug().Model(&model.FirstCategory{}).Where("first_category_uid = ?", fc.FirstCategoryUid).First(&fc).Error; err != nil {
-			// error handling...
-			if gorm.IsRecordNotFoundError(err) {
-				model.Db.Debug().Model(&model.FirstCategory{}).Create(&fc) // newUser not user
-			}
-		} else {
-			model.Db.Debug().Model(&model.FirstCategory{}).Where("first_category_uid = ?", fc.FirstCategoryUid).Update(
-				map[string]interface{}{
-					"first_category_index":       i,
-					"first_category_name":        item.Get("name").String(),
-					"first_category_description": item.Get("description").String(),
-					"first_category_icon_path":   item.Get("icon_path").String(),
-				})
-		}
-		uidList = append(uidList, fc.FirstCategoryUid)
-	}
-
-	log.Println(uidList)
-	return uidList
-}

plugin_vtb_quotation/model/model.go

@@ -1,39 +1,51 @@
 package model
 
 import (
-	"github.com/jinzhu/gorm"
-	_ "github.com/logoove/sqlite"
-	log "github.com/sirupsen/logrus"
-
+	"io/ioutil"
 	"math/rand"
+	"net/http"
 	"os"
 	"strconv"
+	"strings"
 	"time"
+	"unsafe"
+
+	"github.com/jinzhu/gorm"
+	_ "github.com/logoove/sqlite"
+	"github.com/sirupsen/logrus"
+	"github.com/tidwall/gjson"
 )
 
-var (
-	Db   *gorm.DB
-	path = "data/VtbQuotation/vtb.db"
-)
+type VtbDB gorm.DB
 
-func Init() {
+func Init(dbpath string) *VtbDB {
 	var err error
-	if _, err = os.Stat(path); err != nil || os.IsNotExist(err) {
+	if _, err = os.Stat(dbpath); err != nil || os.IsNotExist(err) {
 		// 生成文件
-		f, err := os.Create(path)
+		f, err := os.Create(dbpath)
 		if err != nil {
-			return
+			return nil
 		}
 		defer f.Close()
 	}
-	Db, err = gorm.Open("sqlite3", path)
+	gdb, err := gorm.Open("sqlite3", dbpath)
 	if err != nil {
-		panic("failed to connect database")
+		panic(err)
 	}
-	Db.AutoMigrate(FirstCategory{}).AutoMigrate(SecondCategory{}).AutoMigrate(ThirdCategory{})
+	gdb.AutoMigrate(FirstCategory{}).AutoMigrate(SecondCategory{}).AutoMigrate(ThirdCategory{})
+	return (*VtbDB)(unsafe.Pointer(gdb))
 }
 
-//第一品类
+func Open(dbpath string) (*VtbDB, error) {
+	db, err := gorm.Open("sqlite3", dbpath)
+	if err != nil {
+		return nil, err
+	} else {
+		return (*VtbDB)(unsafe.Pointer(db)), nil
+	}
+}
+
+// FirstCategory 第一品类
 type FirstCategory struct {
 	gorm.Model
 	FirstCategoryIndex       int64  `gorm:"column:first_category_index"`
@@ -47,7 +59,7 @@ func (FirstCategory) TableName() string {
 	return "first_category"
 }
 
-//第二品类
+// SecondCategory 第二品类
 type SecondCategory struct {
 	gorm.Model
 	SecondCategoryIndex       int64  `gorm:"column:second_category_index"`
@@ -61,7 +73,7 @@ func (SecondCategory) TableName() string {
 	return "second_category"
 }
 
-//第三品类
+// ThirdCategory 第三品类
 type ThirdCategory struct {
 	gorm.Model
 	ThirdCategoryIndex       int64  `gorm:"column:third_category_index"`
@@ -77,27 +89,29 @@ func (ThirdCategory) TableName() string {
 	return "third_category"
 }
 
-//取出所有vtb
-func GetAllFirstCategoryMessage(db *gorm.DB) string {
+// GetAllFirstCategoryMessage 取出所有vtb
+func (vdb *VtbDB) GetAllFirstCategoryMessage() string {
+	db := (*gorm.DB)(unsafe.Pointer(vdb))
 	firstStepMessage := "请选择一个vtb并发送序号:\n"
 	var fc FirstCategory
 	rows, err := db.Model(&FirstCategory{}).Rows()
 	if err != nil {
-		log.Println("数据库读取错误", err)
+		logrus.Errorln("[vtb/model]数据库读取错误", err)
 	}
 	if rows == nil {
 		return ""
 	}
 	for rows.Next() {
 		db.ScanRows(rows, &fc)
-		log.Println(fc)
+		// logrus.Println(fc)
 		firstStepMessage = firstStepMessage + strconv.FormatInt(fc.FirstCategoryIndex, 10) + ". " + fc.FirstCategoryName + "\n"
 	}
 	return firstStepMessage
 }
 
-//取得同一个vtb所有语录类别
-func GetAllSecondCategoryMessageByFirstIndex(db *gorm.DB, firstIndex int) string {
+// GetAllSecondCategoryMessageByFirstIndex 取得同一个vtb所有语录类别
+func (vdb *VtbDB) GetAllSecondCategoryMessageByFirstIndex(firstIndex int) string {
+	db := (*gorm.DB)(unsafe.Pointer(vdb))
 	SecondStepMessage := "请选择一个语录类别并发送序号:\n"
 	var sc SecondCategory
 	var count int
@@ -109,19 +123,20 @@ func GetAllSecondCategoryMessageByFirstIndex(db *gorm.DB, firstIndex int) string
 	}
 	rows, err := db.Model(&SecondCategory{}).Where("first_category_uid = ?", fc.FirstCategoryUid).Rows()
 	if err != nil {
-		log.Println("数据库读取错误", err)
+		logrus.Errorln("[vtb/model]数据库读取错误", err)
 	}
 
 	for rows.Next() {
 		db.ScanRows(rows, &sc)
-		log.Println(sc)
+		// logrus.Println(sc)
 		SecondStepMessage = SecondStepMessage + strconv.FormatInt(sc.SecondCategoryIndex, 10) + ". " + sc.SecondCategoryName + "\n"
 	}
 	return SecondStepMessage
 }
 
-//取得同一个vtb同个类别的所有语录
-func GetAllThirdCategoryMessageByFirstIndexAndSecondIndex(db *gorm.DB, firstIndex, secondIndex int) string {
+// GetAllThirdCategoryMessageByFirstIndexAndSecondIndex 取得同一个vtb同个类别的所有语录
+func (vdb *VtbDB) GetAllThirdCategoryMessageByFirstIndexAndSecondIndex(firstIndex, secondIndex int) string {
+	db := (*gorm.DB)(unsafe.Pointer(vdb))
 	ThirdStepMessage := "请选择一个语录并发送序号:\n"
 	var fc FirstCategory
 	db.Model(FirstCategory{}).Where("first_category_index = ?", firstIndex).First(&fc)
@@ -133,16 +148,19 @@ func GetAllThirdCategoryMessageByFirstIndexAndSecondIndex(db *gorm.DB, firstInde
 	var tc ThirdCategory
 	rows, err := db.Model(&ThirdCategory{}).Where("first_category_uid = ? and second_category_index = ?", fc.FirstCategoryUid, secondIndex).Rows()
 	if err != nil {
-		log.Println("数据库读取错误", err)
+		logrus.Errorln("[vtb/model]数据库读取错误", err)
 	}
 	for rows.Next() {
 		db.ScanRows(rows, &tc)
-		log.Println(tc)
+		// logrus.Println(tc)
 		ThirdStepMessage = ThirdStepMessage + strconv.FormatInt(tc.ThirdCategoryIndex, 10) + ". " + tc.ThirdCategoryName + "\n"
 	}
 	return ThirdStepMessage
 }
-func GetThirdCategory(db *gorm.DB, firstIndex, secondIndex, thirdIndex int) ThirdCategory {
+
+// GetThirdCategory
+func (vdb *VtbDB) GetThirdCategory(firstIndex, secondIndex, thirdIndex int) ThirdCategory {
+	db := (*gorm.DB)(unsafe.Pointer(vdb))
 	var fc FirstCategory
 	db.Model(FirstCategory{}).Where("first_category_index = ?", firstIndex).First(&fc)
 	var tc ThirdCategory
@@ -150,20 +168,195 @@ func GetThirdCategory(db *gorm.DB, firstIndex, secondIndex, thirdIndex int) Thir
 	return tc
 }
 
-func RandomVtb(db *gorm.DB) ThirdCategory {
+func (vdb *VtbDB) RandomVtb() ThirdCategory {
+	db := (*gorm.DB)(unsafe.Pointer(vdb))
 	rand.Seed(time.Now().UnixNano())
 	var count int
 	db.Model(&ThirdCategory{}).Count(&count)
-	log.Info("一共有", count, "个")
+	// logrus.Info("一共有", count, "个")
 	var tc ThirdCategory
 	db.Model(&ThirdCategory{}).Offset(rand.Intn(count)).Take(&tc)
-	log.Info(tc)
+	// logrus.Info(tc)
 	return tc
 }
 
-func GetFirstCategoryByFirstUid(db *gorm.DB, firstUid string) FirstCategory {
+func (vdb *VtbDB) GetFirstCategoryByFirstUid(firstUid string) FirstCategory {
+	db := (*gorm.DB)(unsafe.Pointer(vdb))
 	var fc FirstCategory
 	db.Model(FirstCategory{}).Where("first_category_uid = ?", firstUid).Take(&fc)
-	log.Info(fc)
+	// logrus.Info(fc)
 	return fc
 }
+
+func (vdb *VtbDB) Close() error {
+	db := (*gorm.DB)(unsafe.Pointer(vdb))
+	return db.Close()
+}
+
+const vtbUrl = "https://vtbkeyboard.moe/api/get_vtb_list"
+
+func (vdb *VtbDB) GetVtbList() []string {
+	db := (*gorm.DB)(unsafe.Pointer(vdb))
+	client := &http.Client{}
+	req, err := http.NewRequest("GET", vtbUrl, nil)
+	if err != nil {
+		logrus.Errorln(err)
+	}
+	// 自定义Header
+	req.Header.Set("User-Agent", randua())
+	resp, err := client.Do(req)
+	if err != nil {
+		logrus.Errorln(err)
+	}
+
+	defer resp.Body.Close()
+	bytes, err := ioutil.ReadAll(resp.Body)
+	if err != nil {
+		logrus.Errorln(err)
+	}
+	// logrus.Println(string(bytes))
+	vtbListStr, err := strconv.Unquote(strings.Replace(strconv.Quote(string(bytes)), `\\u`, `\u`, -1))
+	if err != nil {
+		logrus.Errorln(err)
+	}
+	// logrus.Println(vtbListStr)
+	uidList := make([]string, 0)
+	count := gjson.Get(vtbListStr, "#").Int()
+	for i := int64(0); i < count; i++ {
+		item := gjson.Get(vtbListStr, strconv.FormatInt(i, 10))
+		logrus.Println(item)
+		fc := FirstCategory{
+			FirstCategoryIndex:       i,
+			FirstCategoryName:        item.Get("name").String(),
+			FirstCategoryDescription: item.Get("description").String(),
+			FirstCategoryIconPath:    item.Get("icon_path").String(),
+			FirstCategoryUid:         item.Get("uid").String(),
+		}
+		logrus.Println(fc)
+		//db.Model(FirstCategory{}).Where("first_category_uid = ?", fc.FirstCategoryUid).FirstOrCreate(&fc)
+		if err := db.Debug().Model(&FirstCategory{}).Where("first_category_uid = ?", fc.FirstCategoryUid).First(&fc).Error; err != nil {
+			// error handling...
+			if gorm.IsRecordNotFoundError(err) {
+				db.Debug().Model(&FirstCategory{}).Create(&fc) // newUser not user
+			}
+		} else {
+			db.Debug().Model(&FirstCategory{}).Where("first_category_uid = ?", fc.FirstCategoryUid).Update(
+				map[string]interface{}{
+					"first_category_index":       i,
+					"first_category_name":        item.Get("name").String(),
+					"first_category_description": item.Get("description").String(),
+					"first_category_icon_path":   item.Get("icon_path").String(),
+				})
+		}
+		uidList = append(uidList, fc.FirstCategoryUid)
+	}
+
+	// logrus.Println(uidList)
+	return uidList
+}
+
+func (vdb *VtbDB) StoreVtb(uid string) {
+	db := (*gorm.DB)(unsafe.Pointer(vdb))
+	vtbUrl := "https://vtbkeyboard.moe/api/get_vtb_page?uid=" + uid
+	client := &http.Client{}
+	req, err := http.NewRequest("GET", vtbUrl, nil)
+	if err != nil {
+		logrus.Errorln(err)
+	}
+	// 自定义Header
+	req.Header.Set("User-Agent", randua())
+	resp, err := client.Do(req)
+	if err != nil {
+		logrus.Errorln(err)
+	}
+
+	defer resp.Body.Close()
+	bytes, err := ioutil.ReadAll(resp.Body)
+	if err != nil {
+		logrus.Errorln(err)
+	}
+	//logrus.Println(string(bytes))
+	vtbStr, err := strconv.Unquote(strings.Replace(strconv.Quote(string(bytes)), `\\u`, `\u`, -1))
+	if err != nil {
+		logrus.Errorln(err)
+	}
+	// logrus.Println(vtbListStr)
+	secondCount := gjson.Get(vtbStr, "data.voices.#").Int()
+	logrus.Println("二级品类一共有", secondCount)
+	for secondIndex := int64(0); secondIndex < secondCount; secondIndex++ {
+		secondItem := gjson.Get(vtbStr, "data.voices."+strconv.FormatInt(secondIndex, 10))
+		logrus.Println(secondItem)
+		sc := SecondCategory{
+			SecondCategoryName:        secondItem.Get("categoryName").String(),
+			SecondCategoryIndex:       secondIndex,
+			SecondCategoryAuthor:      secondItem.Get("author").String(),
+			SecondCategoryDescription: secondItem.Get("categoryDescription.zh-CN").String(),
+			FirstCategoryUid:          uid,
+		}
+		// logrus.Println(sc)
+		// db.Model(SecondCategory{}).Where("first_category_uid = ? and second_category_index = ?", uid, secondIndex).FirstOrCreate(&sc)
+		if err := db.Debug().Model(&SecondCategory{}).Where("first_category_uid = ? and second_category_index = ?", uid, secondIndex).First(&sc).Error; err != nil {
+			// error handling...
+			if gorm.IsRecordNotFoundError(err) {
+				db.Debug().Model(&SecondCategory{}).Create(&sc) // newUser not user
+			}
+		} else {
+			db.Debug().Model(&SecondCategory{}).Where("first_category_uid = ? and second_category_index = ?", uid, secondIndex).Update(
+				map[string]interface{}{
+					"second_category_name":        secondItem.Get("categoryName").String(),
+					"second_category_author":      secondItem.Get("author").String(),
+					"second_category_description": secondItem.Get("categoryDescription.zh-CN").String(),
+				})
+		}
+		thirdCount := secondItem.Get("voiceList.#").Int()
+		logrus.Println("三级品类一共有", thirdCount)
+		for thirdIndex := int64(0); thirdIndex < thirdCount; thirdIndex++ {
+			thirdItem := secondItem.Get("voiceList." + strconv.FormatInt(thirdIndex, 10))
+			logrus.Println(thirdItem)
+			tc := ThirdCategory{
+				ThirdCategoryName:        thirdItem.Get("name").String(),
+				ThirdCategoryIndex:       thirdIndex,
+				ThirdCategoryDescription: thirdItem.Get("description.zh-CN").String(),
+				FirstCategoryUid:         uid,
+				SecondCategoryIndex:      secondIndex,
+				ThirdCategoryPath:        thirdItem.Get("path").String(),
+				ThirdCategoryAuthor:      thirdItem.Get("author").String(),
+			}
+			logrus.Println(tc)
+			//db.Model(ThirdCategory{}).Where("first_category_uid = ? and second_category_index = ? and third_category_index = ?",
+			//	uid, secondIndex, thirdIndex).FirstOrCreate(&tc)
+			if err := db.Debug().Model(&ThirdCategory{}).Where("first_category_uid = ? and second_category_index = ? and third_category_index = ?",
+				uid, secondIndex, thirdIndex).First(&tc).Error; err != nil {
+				// error handling...
+				if gorm.IsRecordNotFoundError(err) {
+					db.Debug().Model(&ThirdCategory{}).Create(&tc) // newUser not user
+				}
+			} else {
+				db.Debug().Model(&ThirdCategory{}).Where("first_category_uid = ? and second_category_index = ? and third_category_index = ?",
+					uid, secondIndex, thirdIndex).Update(
+					map[string]interface{}{
+						"third_category_name":        thirdItem.Get("name").String(),
+						"third_category_description": thirdItem.Get("description.zh-CN").String(),
+						"third_category_path":        thirdItem.Get("path").String(),
+						"third_category_author":      thirdItem.Get("author").String(),
+					})
+			}
+		}
+	}
+}
+
+var agent = [...]string{
+	"Mozilla/5.0 (Windows NT 6.1; Win64; x64; rv:50.0) Gecko/20100101 Firefox/50.0",
+	"Opera/9.80 (Macintosh; Intel Mac OS X 10.6.8; U; en) Presto/2.8.131 Version/11.11",
+	"Opera/9.80 (Windows NT 6.1; U; en) Presto/2.8.131 Version/11.11",
+	"Mozilla/4.0 (compatible; MSIE 7.0; Windows NT 5.1; 360SE)",
+	"Mozilla/5.0 (Windows NT 6.1; rv:2.0.1) Gecko/20100101 Firefox/4.0.1",
+	"Mozilla/4.0 (compatible; MSIE 7.0; Windows NT 5.1; The World)",
+	"User-Agent,Mozilla/5.0 (Macintosh; U; Intel Mac OS X 10_6_8; en-us) AppleWebKit/534.50 (KHTML, like Gecko) Version/5.1 Safari/534.50",
+	"User-Agent, Mozilla/4.0 (compatible; MSIE 7.0; Windows NT 5.1; Maxthon 2.0)",
+	"User-Agent,Mozilla/5.0 (Windows; U; Windows NT 6.1; en-us) AppleWebKit/534.50 (KHTML, like Gecko) Version/5.1 Safari/534.50",
+}
+
+func randua() string {
+	return agent[rand.New(rand.NewSource(time.Now().UnixNano())).Intn(len(agent))]
+}

plugin_vtb_quotation/secondVtb/secondVtb.go

@@ -1,108 +0,0 @@
-package secondVtb
-
-import (
-	"io/ioutil"
-	"net/http"
-	"strconv"
-	"strings"
-
-	"github.com/jinzhu/gorm"
-	log "github.com/sirupsen/logrus"
-	"github.com/tidwall/gjson"
-
-	"github.com/FloatTech/ZeroBot-Plugin/plugin_vtb_quotation/model"
-	"github.com/FloatTech/ZeroBot-Plugin/plugin_vtb_quotation/utils"
-)
-
-func GetVtbStr(uid string) string {
-	vtbUrl := "https://vtbkeyboard.moe/api/get_vtb_page?uid=" + uid
-	client := &http.Client{}
-	req, err := http.NewRequest("GET", vtbUrl, nil)
-	if err != nil {
-		log.Println(err)
-	}
-	// 自定义Header
-	req.Header.Set("User-Agent", utils.GetAgent())
-	resp, err := client.Do(req)
-	if err != nil {
-		log.Println(err)
-	}
-
-	defer resp.Body.Close()
-	bytes, err := ioutil.ReadAll(resp.Body)
-	if err != nil {
-		log.Println(err)
-	}
-	//log.Println(string(bytes))
-	vtbListStr, err := strconv.Unquote(strings.Replace(strconv.Quote(string(bytes)), `\\u`, `\u`, -1))
-	if err != nil {
-		log.Println(err)
-	}
-	log.Println(vtbListStr)
-	return vtbListStr
-}
-
-func DealVtbStr(vtbStr, uid string) {
-	secondCount := gjson.Get(vtbStr, "data.voices.#").Int()
-	log.Println("二级品类一共有", secondCount)
-	for secondIndex := int64(0); secondIndex < secondCount; secondIndex++ {
-		secondItem := gjson.Get(vtbStr, "data.voices."+strconv.FormatInt(secondIndex, 10))
-		log.Println(secondItem)
-		sc := model.SecondCategory{
-			SecondCategoryName:        secondItem.Get("categoryName").String(),
-			SecondCategoryIndex:       secondIndex,
-			SecondCategoryAuthor:      secondItem.Get("author").String(),
-			SecondCategoryDescription: secondItem.Get("categoryDescription.zh-CN").String(),
-			FirstCategoryUid:          uid,
-		}
-		log.Println(sc)
-		//model.Db.Model(SecondCategory{}).Where("first_category_uid = ? and second_category_index = ?", uid, secondIndex).FirstOrCreate(&sc)
-		if err := model.Db.Debug().Model(&model.SecondCategory{}).Where("first_category_uid = ? and second_category_index = ?", uid, secondIndex).First(&sc).Error; err != nil {
-			// error handling...
-			if gorm.IsRecordNotFoundError(err) {
-				model.Db.Debug().Model(&model.SecondCategory{}).Create(&sc) // newUser not user
-			}
-		} else {
-			model.Db.Debug().Model(&model.SecondCategory{}).Where("first_category_uid = ? and second_category_index = ?", uid, secondIndex).Update(
-				map[string]interface{}{
-					"second_category_name":        secondItem.Get("categoryName").String(),
-					"second_category_author":      secondItem.Get("author").String(),
-					"second_category_description": secondItem.Get("categoryDescription.zh-CN").String(),
-				})
-		}
-		thirdCount := secondItem.Get("voiceList.#").Int()
-		log.Println("三级品类一共有", thirdCount)
-		for thirdIndex := int64(0); thirdIndex < thirdCount; thirdIndex++ {
-			thirdItem := secondItem.Get("voiceList." + strconv.FormatInt(thirdIndex, 10))
-			log.Println(thirdItem)
-			tc := model.ThirdCategory{
-				ThirdCategoryName:        thirdItem.Get("name").String(),
-				ThirdCategoryIndex:       thirdIndex,
-				ThirdCategoryDescription: thirdItem.Get("description.zh-CN").String(),
-				FirstCategoryUid:         uid,
-				SecondCategoryIndex:      secondIndex,
-				ThirdCategoryPath:        thirdItem.Get("path").String(),
-				ThirdCategoryAuthor:      thirdItem.Get("author").String(),
-			}
-			log.Println(tc)
-			//model.Db.Model(ThirdCategory{}).Where("first_category_uid = ? and second_category_index = ? and third_category_index = ?",
-			//	uid, secondIndex, thirdIndex).FirstOrCreate(&tc)
-			if err := model.Db.Debug().Model(&model.ThirdCategory{}).Where("first_category_uid = ? and second_category_index = ? and third_category_index = ?",
-				uid, secondIndex, thirdIndex).First(&tc).Error; err != nil {
-				// error handling...
-				if gorm.IsRecordNotFoundError(err) {
-					model.Db.Debug().Model(&model.ThirdCategory{}).Create(&tc) // newUser not user
-				}
-			} else {
-				model.Db.Debug().Model(&model.ThirdCategory{}).Where("first_category_uid = ? and second_category_index = ? and third_category_index = ?",
-					uid, secondIndex, thirdIndex).Update(
-					map[string]interface{}{
-						"third_category_name":        thirdItem.Get("name").String(),
-						"third_category_description": thirdItem.Get("description.zh-CN").String(),
-						"third_category_path":        thirdItem.Get("path").String(),
-						"third_category_author":      thirdItem.Get("author").String(),
-					})
-			}
-		}
-	}
-}

plugin_vtb_quotation/utils/utils.go

@@ -1,24 +0,0 @@
-package utils
-
-import (
-	"math/rand"
-	"time"
-)
-
-var agent = [...]string{
-	"Mozilla/5.0 (Windows NT 6.1; Win64; x64; rv:50.0) Gecko/20100101 Firefox/50.0",
-	"Opera/9.80 (Macintosh; Intel Mac OS X 10.6.8; U; en) Presto/2.8.131 Version/11.11",
-	"Opera/9.80 (Windows NT 6.1; U; en) Presto/2.8.131 Version/11.11",
-	"Mozilla/4.0 (compatible; MSIE 7.0; Windows NT 5.1; 360SE)",
-	"Mozilla/5.0 (Windows NT 6.1; rv:2.0.1) Gecko/20100101 Firefox/4.0.1",
-	"Mozilla/4.0 (compatible; MSIE 7.0; Windows NT 5.1; The World)",
-	"User-Agent,Mozilla/5.0 (Macintosh; U; Intel Mac OS X 10_6_8; en-us) AppleWebKit/534.50 (KHTML, like Gecko) Version/5.1 Safari/534.50",
-	"User-Agent, Mozilla/4.0 (compatible; MSIE 7.0; Windows NT 5.1; Maxthon 2.0)",
-	"User-Agent,Mozilla/5.0 (Windows; U; Windows NT 6.1; en-us) AppleWebKit/534.50 (KHTML, like Gecko) Version/5.1 Safari/534.50",
-}
-
-func GetAgent() string {
-	r := rand.New(rand.NewSource(time.Now().UnixNano()))
-	len1 := len(agent)
-	return agent[r.Intn(len1)]
-}

plugin_vtb_quotation/vtb_quotation.go

@@ -1,16 +1,14 @@
-package plugin_vtb_quotation
+package vtbquotation
 
 import (
-	"fmt"
 	"net/url"
 	"regexp"
 	"strconv"
 	"strings"
 	"time"
 
-	"github.com/jinzhu/gorm"
 	_ "github.com/logoove/sqlite"
-	log "github.com/sirupsen/logrus"
+	"github.com/sirupsen/logrus"
 	zero "github.com/wdvxdr1123/ZeroBot"
 	"github.com/wdvxdr1123/ZeroBot/message"
 
@@ -18,14 +16,12 @@ import (
 	"github.com/FloatTech/ZeroBot-Plugin/plugin_vtb_quotation/model"
 )
 
-var (
-	regStr = ".*/(.*)"
-)
+const regStr = ".*/(.*)"
+const dbpath = "data/VtbQuotation/vtb.db"
 
 var engine = control.Register("vtbquotation", &control.Options{
 	DisableOnDefault: false,
-	Help: "vtb语录\n" +
-		"随机vtb\n",
+	Help:             "vtbquotation\n- vtb语录\n- 随机vtb\n",
 })
 
 func init() {
@@ -37,97 +33,99 @@ func init() {
 			echo, cancel := ctx.FutureEvent("message",
 				ctx.CheckSession()). // 只复读开启复读模式的人的消息
 				Repeat()             // 不断监听复读
-			db, err := gorm.Open("sqlite3", "data/VtbQuotation/vtb.db")
+			db, err := model.Open(dbpath)
 			if err != nil {
-				panic("failed to connect database")
+				logrus.Errorln(err)
+				return
 			}
 			defer db.Close()
-			firstStepMessage := model.GetAllFirstCategoryMessage(db)
+			defer cancel()
+			firstStepMessage := db.GetAllFirstCategoryMessage()
 			ctx.SendChain(message.Reply(ctx.Event.MessageID), message.Text(firstStepMessage))
-			//步骤1，2，3，依次选择3个类别
-			step := 1
-			//错误次数
+			// 步骤0，1，2，依次选择3个类别
+			step := 0
+			// 错误次数
 			errorCount := 0
 			for {
 				select {
 				case e := <-echo: // 接收到需要复读的消息
-					//错误次数达到3次，结束命令
-					if errorCount == 3 {
+					// 错误次数达到3次，结束命令
+					if errorCount >= 3 {
 						ctx.SendChain(message.Reply(e.MessageID), message.Text("输入错误太多,请重新发指令"))
-						cancel()
 						return
 					}
-					if step == 1 {
+					switch step {
+					case 0:
 						firstIndex, err = strconv.Atoi(e.RawMessage)
-						log.Println(fmt.Sprintf("当前在第%d步", step))
-						log.Println(fmt.Sprintf("firstIndex:%d,secondIndex:%d,thirdIndex:%d", firstIndex, secondIndex, thirdIndex))
+						// log.Println(fmt.Sprintf("当前在第%d步", step))
+						// log.Println(fmt.Sprintf("firstIndex:%d,secondIndex:%d,thirdIndex:%d", firstIndex, secondIndex, thirdIndex))
 						if err != nil {
 							ctx.SendChain(message.Reply(e.MessageID), message.Text("请输入正确的序号,三次输入错误，指令可退出重输"))
 							errorCount++
 						} else {
-							SecondStepMessage := model.GetAllSecondCategoryMessageByFirstIndex(db, firstIndex)
-							log.Println(SecondStepMessage)
+							SecondStepMessage := db.GetAllSecondCategoryMessageByFirstIndex(firstIndex)
+							// log.Println(SecondStepMessage)
 							if SecondStepMessage == "" {
 								ctx.SendChain(message.Reply(e.MessageID), message.Text("你选择的序号没有内容，请重新选择，三次输入错误，指令可退出重输"))
-								ctx.SendChain(message.Reply(e.MessageID), message.Text(model.GetAllFirstCategoryMessage(db)))
+								ctx.SendChain(message.Reply(e.MessageID), message.Text(db.GetAllFirstCategoryMessage()))
 								errorCount++
 							} else {
 								ctx.SendChain(message.Reply(e.MessageID), message.Text(SecondStepMessage))
 								step++
 							}
 						}
-					} else if step == 2 {
+					case 1:
 						secondIndex, err = strconv.Atoi(e.RawMessage)
-						log.Println(fmt.Sprintf("当前在第%d步", step))
-						log.Println(fmt.Sprintf("firstIndex:%d,secondIndex:%d,thirdIndex:%d", firstIndex, secondIndex, thirdIndex))
+						// log.Println(fmt.Sprintf("当前在第%d步", step))
+						// log.Println(fmt.Sprintf("firstIndex:%d,secondIndex:%d,thirdIndex:%d", firstIndex, secondIndex, thirdIndex))
 						if err != nil {
 							ctx.SendChain(message.Reply(e.MessageID), message.Text("请输入正确的序号，三次输入错误，指令可退出重输"))
 							errorCount++
 						} else {
-							ThirdStepMessage := model.GetAllThirdCategoryMessageByFirstIndexAndSecondIndex(db, firstIndex, secondIndex)
-							log.Println(ThirdStepMessage)
+							ThirdStepMessage := db.GetAllThirdCategoryMessageByFirstIndexAndSecondIndex(firstIndex, secondIndex)
+							// log.Println(ThirdStepMessage)
 							if ThirdStepMessage == "" {
 								ctx.SendChain(message.Reply(e.MessageID), message.Text("你选择的序号没有内容，请重新选择，三次输入错误，指令可退出重输"))
-								ctx.SendChain(message.Reply(e.MessageID), message.Text(model.GetAllSecondCategoryMessageByFirstIndex(db, firstIndex)))
+								ctx.SendChain(message.Reply(e.MessageID), message.Text(db.GetAllSecondCategoryMessageByFirstIndex(firstIndex)))
 								errorCount++
 							} else {
 								ctx.SendChain(message.Reply(e.MessageID), message.Text(ThirdStepMessage))
 								step++
 							}
 						}
-					} else if step == 3 {
+					case 2:
 						thirdIndex, err = strconv.Atoi(e.RawMessage)
-						log.Println(fmt.Sprintf("当前在第%d步", step))
-						log.Println(fmt.Sprintf("firstIndex:%d,secondIndex:%d,thirdIndex:%d", firstIndex, secondIndex, thirdIndex))
+						// log.Println(fmt.Sprintf("当前在第%d步", step))
+						// log.Println(fmt.Sprintf("firstIndex:%d,secondIndex:%d,thirdIndex:%d", firstIndex, secondIndex, thirdIndex))
 						if err != nil {
 							ctx.SendChain(message.Reply(e.MessageID), message.Text("请输入正确的序号，三次输入错误，指令可退出重输"))
 							errorCount++
 						} else {
-							tc := model.GetThirdCategory(db, firstIndex, secondIndex, thirdIndex)
+							tc := db.GetThirdCategory(firstIndex, secondIndex, thirdIndex)
 							reg := regexp.MustCompile(regStr)
 							recordUrl := tc.ThirdCategoryPath
 							if recordUrl == "" {
 								ctx.SendChain(message.Reply(e.MessageID), message.Text("没有内容请重新选择，三次输入错误，指令可退出重输"))
-								ctx.SendChain(message.Reply(e.MessageID), message.Text(model.GetAllFirstCategoryMessage(db)))
+								ctx.SendChain(message.Reply(e.MessageID), message.Text(db.GetAllFirstCategoryMessage()))
 								errorCount++
 								step = 1
 							} else {
 								if reg.MatchString(recordUrl) {
-									log.Println(reg.FindStringSubmatch(recordUrl)[1])
-									log.Println(url.QueryEscape(reg.FindStringSubmatch(recordUrl)[1]))
+									// log.Println(reg.FindStringSubmatch(recordUrl)[1])
+									// log.Println(url.QueryEscape(reg.FindStringSubmatch(recordUrl)[1]))
 									recordUrl = strings.Replace(recordUrl, reg.FindStringSubmatch(recordUrl)[1], url.QueryEscape(reg.FindStringSubmatch(recordUrl)[1]), -1)
 									recordUrl = strings.Replace(recordUrl, "+", "%20", -1)
-									log.Println(recordUrl)
+									// log.Println(recordUrl)
 								}
 								ctx.SendChain(message.Reply(e.MessageID), message.Text("请欣赏《"+tc.ThirdCategoryName+"》"))
 								ctx.SendChain(message.Record(recordUrl))
-								cancel()
 								return
 							}
 						}
+					default:
+						return
 					}
 				case <-time.After(time.Second * 60):
-					cancel()
 					ctx.SendChain(message.Reply(ctx.Event.MessageID), message.Text("vtb语录指令过期"))
 					return
 				}
@@ -135,25 +133,26 @@ func init() {
 		})
 	engine.OnFullMatch("随机vtb").SetBlock(true).
 		Handle(func(ctx *zero.Ctx) {
-			db, err := gorm.Open("sqlite3", "data/VtbQuotation/vtb.db")
+			db, err := model.Open(dbpath)
 			if err != nil {
-				panic("failed to connect database")
+				logrus.Errorln(err)
+				return
 			}
-			defer db.Close()
-			tc := model.RandomVtb(db)
-			fc := model.GetFirstCategoryByFirstUid(db, tc.FirstCategoryUid)
+			tc := db.RandomVtb()
+			fc := db.GetFirstCategoryByFirstUid(tc.FirstCategoryUid)
 			if (tc != model.ThirdCategory{}) && (fc != model.FirstCategory{}) {
 				reg := regexp.MustCompile(regStr)
 				recordUrl := tc.ThirdCategoryPath
 				if reg.MatchString(recordUrl) {
-					log.Println(reg.FindStringSubmatch(recordUrl)[1])
-					log.Println(url.QueryEscape(reg.FindStringSubmatch(recordUrl)[1]))
+					// log.Println(reg.FindStringSubmatch(recordUrl)[1])
+					// log.Println(url.QueryEscape(reg.FindStringSubmatch(recordUrl)[1]))
 					recordUrl = strings.Replace(recordUrl, reg.FindStringSubmatch(recordUrl)[1], url.QueryEscape(reg.FindStringSubmatch(recordUrl)[1]), -1)
 					recordUrl = strings.Replace(recordUrl, "+", "%20", -1)
-					log.Println(recordUrl)
+					// log.Println(recordUrl)
 				}
 				ctx.SendChain(message.Reply(ctx.Event.MessageID), message.Text("请欣赏"+fc.FirstCategoryName+"的《"+tc.ThirdCategoryName+"》"))
 				ctx.SendChain(message.Record(recordUrl))
 			}
+			db.Close()
 		})
 }