mirror of
https://github.com/FloatTech/ZeroBot-Plugin.git
synced 2025-12-20 06:20:08 +08:00
366 lines
12 KiB
Go
366 lines
12 KiB
Go
package model
|
|
|
|
import (
|
|
"io/ioutil"
|
|
"math/rand"
|
|
"net/http"
|
|
"os"
|
|
"strconv"
|
|
"strings"
|
|
"time"
|
|
|
|
"github.com/jinzhu/gorm"
|
|
_ "github.com/logoove/sqlite"
|
|
"github.com/sirupsen/logrus"
|
|
"github.com/tidwall/gjson"
|
|
)
|
|
|
|
type VtbDB gorm.DB
|
|
|
|
func Initialize(dbpath string) *VtbDB {
|
|
var err error
|
|
if _, err = os.Stat(dbpath); err != nil || os.IsNotExist(err) {
|
|
// 生成文件
|
|
f, err := os.Create(dbpath)
|
|
if err != nil {
|
|
return nil
|
|
}
|
|
defer f.Close()
|
|
}
|
|
gdb, err := gorm.Open("sqlite3", dbpath)
|
|
if err != nil {
|
|
panic(err)
|
|
}
|
|
gdb.AutoMigrate(FirstCategory{}).AutoMigrate(SecondCategory{}).AutoMigrate(ThirdCategory{})
|
|
return (*VtbDB)(gdb)
|
|
}
|
|
|
|
func Open(dbpath string) (*VtbDB, error) {
|
|
db, err := gorm.Open("sqlite3", dbpath)
|
|
if err != nil {
|
|
return nil, err
|
|
} else {
|
|
return (*VtbDB)(db), nil
|
|
}
|
|
}
|
|
|
|
// FirstCategory 第一品类
|
|
type FirstCategory struct {
|
|
gorm.Model
|
|
FirstCategoryIndex int64 `gorm:"column:first_category_index"`
|
|
FirstCategoryName string `gorm:"column:first_category_name"`
|
|
FirstCategoryUid string `gorm:"column:first_category_uid"`
|
|
FirstCategoryDescription string `gorm:"column:first_category_description;type:varchar(1024)"`
|
|
FirstCategoryIconPath string `gorm:"column:first_category_icon_path"`
|
|
}
|
|
|
|
func (FirstCategory) TableName() string {
|
|
return "first_category"
|
|
}
|
|
|
|
// SecondCategory 第二品类
|
|
type SecondCategory struct {
|
|
gorm.Model
|
|
SecondCategoryIndex int64 `gorm:"column:second_category_index"`
|
|
FirstCategoryUid string `gorm:"column:first_category_uid;association_foreignkey:first_category_uid"`
|
|
SecondCategoryName string `gorm:"column:second_category_name"`
|
|
SecondCategoryAuthor string `gorm:"column:second_category_author"`
|
|
SecondCategoryDescription string `gorm:"column:second_category_description"`
|
|
}
|
|
|
|
func (SecondCategory) TableName() string {
|
|
return "second_category"
|
|
}
|
|
|
|
// ThirdCategory 第三品类
|
|
type ThirdCategory struct {
|
|
gorm.Model
|
|
ThirdCategoryIndex int64 `gorm:"column:third_category_index"`
|
|
SecondCategoryIndex int64 `gorm:"column:second_category_index"`
|
|
FirstCategoryUid string `gorm:"column:first_category_uid"`
|
|
ThirdCategoryName string `gorm:"column:third_category_name"`
|
|
ThirdCategoryPath string `gorm:"column:third_category_path"`
|
|
ThirdCategoryAuthor string `gorm:"column:third_category_author"`
|
|
ThirdCategoryDescription string `gorm:"column:third_category_description"`
|
|
}
|
|
|
|
func (ThirdCategory) TableName() string {
|
|
return "third_category"
|
|
}
|
|
|
|
// GetAllFirstCategoryMessage 取出所有vtb
|
|
func (vdb *VtbDB) GetAllFirstCategoryMessage() string {
|
|
db := (*gorm.DB)(vdb)
|
|
firstStepMessage := "请选择一个vtb并发送序号:\n"
|
|
var fc FirstCategory
|
|
rows, err := db.Model(&FirstCategory{}).Rows()
|
|
if err != nil {
|
|
logrus.Errorln("[vtb/model]数据库读取错误", err)
|
|
}
|
|
if rows == nil {
|
|
return ""
|
|
}
|
|
for rows.Next() {
|
|
db.ScanRows(rows, &fc)
|
|
// logrus.Println(fc)
|
|
firstStepMessage = firstStepMessage + strconv.FormatInt(fc.FirstCategoryIndex, 10) + ". " + fc.FirstCategoryName + "\n"
|
|
}
|
|
return firstStepMessage
|
|
}
|
|
|
|
// GetAllSecondCategoryMessageByFirstIndex 取得同一个vtb所有语录类别
|
|
func (vdb *VtbDB) GetAllSecondCategoryMessageByFirstIndex(firstIndex int) string {
|
|
db := (*gorm.DB)(vdb)
|
|
SecondStepMessage := "请选择一个语录类别并发送序号:\n"
|
|
var sc SecondCategory
|
|
var count int
|
|
var fc FirstCategory
|
|
db.Model(FirstCategory{}).Where("first_category_index = ?", firstIndex).First(&fc)
|
|
db.Model(&SecondCategory{}).Where("first_category_uid = ?", fc.FirstCategoryUid).Count(&count)
|
|
if count == 0 {
|
|
return ""
|
|
}
|
|
rows, err := db.Model(&SecondCategory{}).Where("first_category_uid = ?", fc.FirstCategoryUid).Rows()
|
|
if err != nil {
|
|
logrus.Errorln("[vtb/model]数据库读取错误", err)
|
|
}
|
|
|
|
for rows.Next() {
|
|
db.ScanRows(rows, &sc)
|
|
// logrus.Println(sc)
|
|
SecondStepMessage = SecondStepMessage + strconv.FormatInt(sc.SecondCategoryIndex, 10) + ". " + sc.SecondCategoryName + "\n"
|
|
}
|
|
return SecondStepMessage
|
|
}
|
|
|
|
// GetAllThirdCategoryMessageByFirstIndexAndSecondIndex 取得同一个vtb同个类别的所有语录
|
|
func (vdb *VtbDB) GetAllThirdCategoryMessageByFirstIndexAndSecondIndex(firstIndex, secondIndex int) string {
|
|
db := (*gorm.DB)(vdb)
|
|
ThirdStepMessage := "请选择一个语录并发送序号:\n"
|
|
var fc FirstCategory
|
|
db.Model(FirstCategory{}).Where("first_category_index = ?", firstIndex).First(&fc)
|
|
var count int
|
|
db.Model(&ThirdCategory{}).Where("first_category_uid = ? and second_category_index = ?", fc.FirstCategoryUid, secondIndex).Count(&count)
|
|
if count == 0 {
|
|
return ""
|
|
}
|
|
var tc ThirdCategory
|
|
rows, err := db.Model(&ThirdCategory{}).Where("first_category_uid = ? and second_category_index = ?", fc.FirstCategoryUid, secondIndex).Rows()
|
|
if err != nil {
|
|
logrus.Errorln("[vtb/model]数据库读取错误", err)
|
|
}
|
|
for rows.Next() {
|
|
db.ScanRows(rows, &tc)
|
|
// logrus.Println(tc)
|
|
ThirdStepMessage = ThirdStepMessage + strconv.FormatInt(tc.ThirdCategoryIndex, 10) + ". " + tc.ThirdCategoryName + "\n"
|
|
}
|
|
return ThirdStepMessage
|
|
}
|
|
|
|
// GetThirdCategory
|
|
func (vdb *VtbDB) GetThirdCategory(firstIndex, secondIndex, thirdIndex int) ThirdCategory {
|
|
db := (*gorm.DB)(vdb)
|
|
var fc FirstCategory
|
|
db.Model(FirstCategory{}).Where("first_category_index = ?", firstIndex).First(&fc)
|
|
var tc ThirdCategory
|
|
db.Model(&ThirdCategory{}).Where("first_category_uid = ? and second_category_index = ? and third_category_index = ?", fc.FirstCategoryUid, secondIndex, thirdIndex).Take(&tc)
|
|
return tc
|
|
}
|
|
|
|
func (vdb *VtbDB) RandomVtb() ThirdCategory {
|
|
db := (*gorm.DB)(vdb)
|
|
rand.Seed(time.Now().UnixNano())
|
|
var count int
|
|
db.Model(&ThirdCategory{}).Count(&count)
|
|
// logrus.Info("一共有", count, "个")
|
|
var tc ThirdCategory
|
|
db.Model(&ThirdCategory{}).Offset(rand.Intn(count)).Take(&tc)
|
|
// logrus.Info(tc)
|
|
return tc
|
|
}
|
|
|
|
func (vdb *VtbDB) GetFirstCategoryByFirstUid(firstUid string) FirstCategory {
|
|
db := (*gorm.DB)(vdb)
|
|
var fc FirstCategory
|
|
db.Model(FirstCategory{}).Where("first_category_uid = ?", firstUid).Take(&fc)
|
|
// logrus.Info(fc)
|
|
return fc
|
|
}
|
|
|
|
func (vdb *VtbDB) Close() error {
|
|
db := (*gorm.DB)(vdb)
|
|
return db.Close()
|
|
}
|
|
|
|
const vtbUrl = "https://vtbkeyboard.moe/api/get_vtb_list"
|
|
|
|
func (vdb *VtbDB) GetVtbList() (uidList []string) {
|
|
db := (*gorm.DB)(vdb)
|
|
client := &http.Client{}
|
|
req, err := http.NewRequest("GET", vtbUrl, nil)
|
|
if err != nil {
|
|
logrus.Errorln(err)
|
|
return
|
|
}
|
|
// 自定义Header
|
|
req.Header.Set("User-Agent", randua())
|
|
resp, err := client.Do(req)
|
|
if err != nil {
|
|
logrus.Errorln(err)
|
|
return
|
|
}
|
|
|
|
defer resp.Body.Close()
|
|
bytes, err := ioutil.ReadAll(resp.Body)
|
|
if err != nil {
|
|
logrus.Errorln(err)
|
|
return
|
|
}
|
|
|
|
vtbListStr, err := strconv.Unquote(strings.Replace(strconv.Quote(string(bytes)), `\\u`, `\u`, -1))
|
|
if err != nil {
|
|
logrus.Errorln(err)
|
|
return
|
|
}
|
|
|
|
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)
|
|
|
|
if err := db.Debug().Model(&FirstCategory{}).Where("first_category_uid = ?", fc.FirstCategoryUid).First(&fc).Error; err != nil {
|
|
|
|
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)
|
|
}
|
|
|
|
return uidList
|
|
}
|
|
|
|
func (vdb *VtbDB) StoreVtb(uid string) {
|
|
db := (*gorm.DB)(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)
|
|
return
|
|
}
|
|
// 自定义Header
|
|
req.Header.Set("User-Agent", randua())
|
|
resp, err := client.Do(req)
|
|
if err != nil {
|
|
logrus.Errorln(err)
|
|
return
|
|
}
|
|
|
|
defer resp.Body.Close()
|
|
bytes, err := ioutil.ReadAll(resp.Body)
|
|
if err != nil {
|
|
logrus.Errorln(err)
|
|
return
|
|
}
|
|
|
|
vtbStr, err := strconv.Unquote(strings.Replace(strconv.Quote(string(bytes)), `\\u`, `\u`, -1))
|
|
if err != nil {
|
|
logrus.Errorln(err)
|
|
return
|
|
}
|
|
|
|
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,
|
|
}
|
|
|
|
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)
|
|
|
|
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 {
|
|
|
|
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))]
|
|
}
|