From 66ff92e0645864a1fcb888e1d4da9b6562b15610 Mon Sep 17 00:00:00 2001 From: Yiwen-Chan Date: Fri, 16 Apr 2021 19:18:29 +0800 Subject: [PATCH] =?UTF-8?q?=E2=9A=A1=EF=B8=8F=20=E6=9B=B4=E5=A5=BD?= =?UTF-8?q?=E7=9A=84=E8=AF=86=E5=9B=BE=E8=A7=A6=E5=8F=91=E6=A8=A1=E5=BC=8F?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- setutime/pic_searcher.go | 180 ++++++++++++++++++------------------- setutime/utils/saucenao.go | 73 +++++++-------- 2 files changed, 118 insertions(+), 135 deletions(-) diff --git a/setutime/pic_searcher.go b/setutime/pic_searcher.go index 259aca44..3309b35b 100644 --- a/setutime/pic_searcher.go +++ b/setutime/pic_searcher.go @@ -2,9 +2,13 @@ package setutime import ( "fmt" + "strconv" + "strings" + "time" utils "github.com/Yiwen-Chan/ZeroBot-Plugin/setutime/utils" zero "github.com/wdvxdr1123/ZeroBot" + "github.com/wdvxdr1123/ZeroBot/message" ) func init() { // 插件主体 @@ -28,111 +32,97 @@ func init() { // 插件主体 ctx.Send(illust.DetailPic) return }) - // 通过回复以图搜图 - zero.OnRegex(`\[CQ:reply,id=(.*?)\](.*)搜索图片`).SetBlock(true).SetPriority(32). + // 以图搜图 + zero.OnMessage(FullMatchText("以图搜图", "搜索图片", "以图识图"), MustHasPicture()).SetBlock(true).SetPriority(999). Handle(func(ctx *zero.Ctx) { - var pics []string // 图片搜索池子 - // 获取回复的上文图片链接 - id := utils.Str2Int(ctx.State["regex_matched"].([]string)[1]) - for _, elem := range ctx.GetMessage(id).Elements { - if elem.Type == "image" { - pics = append(pics, elem.Data["url"]) - } - } - // 没有收到图片则向用户索取 - if len(pics) == 0 { - ctx.Send("请发送多张图片！") - next := ctx.FutureEvent("message", ctx.CheckSession()) - recv, cancel := next.Repeat() - for e := range recv { // 循环获取channel发来的信息 - if len(e.Message) == 1 && e.Message[0].Type == "text" { - cancel() // 如果是纯文本则退出索取 - break - } - for _, elem := range e.Message { - if elem.Type == "image" { // 将信息中的图片添加到搜索池子 - pics = append(pics, elem.Data["url"]) - } - } - if len(pics) >= 5 { - cancel() // 如果是图片数量大于等于5则退出索取 - break - } - } - } - if len(pics) == 0 { - ctx.Send("没有收到图片，搜图结束......") - return - } // 开始搜索图片 ctx.Send("少女祈祷中......") - for _, pic := range pics { - if text, err := utils.SauceNaoSearch(pic); err == nil { - ctx.Send(text) // 返回SauceNAO的结果 + for _, pic := range ctx.State["image_url"].([]string) { + fmt.Println(pic) + if m, err := utils.SauceNaoSearch(pic); err == nil { + ctx.SendChain(m...) // 返回SauceNAO的结果 continue } else { - ctx.Send(fmt.Sprintf("ERROR: %v", err)) + ctx.SendChain(message.Text("ERROR: ", err)) } - if text, err := utils.Ascii2dSearch(pic); err == nil { - ctx.Send(text) // 返回Ascii2d的结果 + if m, err := utils.Ascii2dSearch(pic); err == nil { + ctx.SendChain(m...) // 返回Ascii2d的结果 continue } else { - ctx.Send(fmt.Sprintf("ERROR: %v", err)) - } - } - return - }) - // 通过命令以图搜图 - zero.OnKeywordGroup([]string{"以图识图", "以图搜图", "搜索图片"}).SetBlock(true).SetPriority(33). - Handle(func(ctx *zero.Ctx) { - var pics []string // 图片搜索池子 - // 获取信息中图片链接 - for _, elem := range ctx.Event.Message { - if elem.Type == "image" { - pics = append(pics, elem.Data["url"]) - } - } - // 没有收到图片则向用户索取 - if len(pics) == 0 { - ctx.Send("请发送多张图片！") - next := ctx.FutureEvent("message", zero.CheckUser(ctx.Event.UserID)) - recv, cancel := next.Repeat() - for e := range recv { // 循环获取channel发来的信息 - if len(e.Message) == 1 && e.Message[0].Type == "text" { - cancel() // 如果是纯文本则退出索取 - break - } - for _, elem := range e.Message { - if elem.Type == "image" { // 将信息中的图片添加到搜索池子 - pics = append(pics, elem.Data["url"]) - } - } - if len(pics) >= 5 { - cancel() // 如果是图片数量大于等于5则退出索取 - break - } - } - } - if len(pics) == 0 { - ctx.Send("没有收到图片，搜图结束......") - return - } - // 开始搜索图片 - ctx.Send("少女祈祷中......") - for _, pic := range pics { - if text, err := utils.SauceNaoSearch(pic); err == nil { - ctx.Send(text) // 返回SauceNAO的结果 - continue - } else { - ctx.Send(fmt.Sprintf("ERROR: %v", err)) - } - if text, err := utils.Ascii2dSearch(pic); err == nil { - ctx.Send(text) // 返回Ascii2d的结果 - continue - } else { - ctx.Send(fmt.Sprintf("ERROR: %v", err)) + ctx.SendChain(message.Text("ERROR: ", err)) } } return }) } + +// FullMatchText 如果信息中文本完全匹配则返回 true +func FullMatchText(src ...string) zero.Rule { + return func(ctx *zero.Ctx) bool { + msg := ctx.Event.Message + for _, elem := range msg { + if elem.Type == "text" { + text := elem.Data["text"] + text = strings.ReplaceAll(text, " ", "") + text = strings.ReplaceAll(text, "\r", "") + text = strings.ReplaceAll(text, "\n", "") + for _, s := range src { + if text == s { + return true + } + } + } + } + return false + } +} + +// HasPicture 消息含有图片返回 true +func HasPicture() zero.Rule { + return func(ctx *zero.Ctx) bool { + msg := ctx.Event.Message + url := []string{} + // 如果是回复信息则将信息替换成被回复的那条 + if msg[0].Type == "reply" { + id, _ := strconv.Atoi(msg[0].Data["id"]) + msg = ctx.GetMessage(int64(id)).Elements + } + // 遍历信息中所有图片 + for _, elem := range msg { + if elem.Type == "image" { + url = append(url, elem.Data["url"]) + } + } + // 如果有图片就返回true + if len(url) > 0 { + ctx.State["image_url"] = url + return true + } + return false + } +} + +// MustHasPicture 消息不存在图片阻塞60秒至有图片，超时返回 false +func MustHasPicture() zero.Rule { + return func(ctx *zero.Ctx) bool { + if HasPicture()(ctx) { + return true + } + // 没有图片就索取 + ctx.Send("请发送一张图片") + next := zero.NewFutureEvent("message", 999, false, zero.CheckUser(ctx.Event.UserID), HasPicture()) + recv, cancel := next.Repeat() + select { + case e := <-recv: + cancel() + newCtx := &zero.Ctx{Event: e, State: zero.State{}} + if HasPicture()(newCtx) { + ctx.State["image_url"] = newCtx.State["image_url"] + return true + } + return false + case <-time.After(time.Second * 60): + return false + } + } +} diff --git a/setutime/utils/saucenao.go b/setutime/utils/saucenao.go index 5584e0a7..8071cba7 100644 --- a/setutime/utils/saucenao.go +++ b/setutime/utils/saucenao.go @@ -1,18 +1,17 @@ package utils import ( - "errors" "fmt" "io/ioutil" "net/http" "net/url" - "strings" "github.com/tidwall/gjson" + "github.com/wdvxdr1123/ZeroBot/message" ) // SauceNaoSearch SauceNao 以图搜图 需要链接 返回错误和信息 -func SauceNaoSearch(pic string) (text string, err error) { +func SauceNaoSearch(pic string) (message.Message, error) { var ( api = "https://saucenao.com/search.php" apiKey = "2cc2772ca550dbacb4c35731a79d341d1a143cb5" @@ -20,65 +19,59 @@ func SauceNaoSearch(pic string) (text string, err error) { minSimilarity = 70.0 // 返回图片结果的最小相似度 ) - transport := http.Transport{ - DisableKeepAlives: true, - } - client := &http.Client{ - Transport: &transport, - } - // 包装请求参数 - data := url.Values{} - data.Set("url", pic) // 图片链接 - data.Set("api_key", apiKey) // api_key - data.Set("db", "5") // 只搜索Pixiv - data.Set("numres", "1") // 返回一个结果 - data.Set("output_type", "2") // 返回JSON格式数据 - fromData := strings.NewReader(data.Encode()) + link, _ := url.Parse(api) + link.RawQuery = url.Values{ + "url": []string{pic}, + "api_key": []string{apiKey}, + "db": []string{"5"}, + "numres": []string{"1"}, + "output_type": []string{"2"}, + }.Encode() // 网络请求 - req, err := http.NewRequest("POST", api, fromData) + client := &http.Client{} + req, err := http.NewRequest("GET", link.String(), nil) if err != nil { - return "", err + return nil, err } req.Header.Set("Content-Type", "application/x-www-form-urlencoded") req.Header.Set("Accept", "Mozilla/5.0 (Windows NT 6.1; WOW64; rv:6.0) Gecko/20100101 Firefox/6.0") resp, err := client.Do(req) if err != nil { - return "", err + return nil, err } defer resp.Body.Close() body, err := ioutil.ReadAll(resp.Body) if err != nil { - return "", err + return nil, err } - if code := resp.StatusCode; code != 200 { + if resp.StatusCode != http.StatusOK { // 如果返回不是200则立刻抛出错误 - return "", errors.New(fmt.Sprintf("SauceNAO not found, code %d", code)) + return nil, fmt.Errorf("SauceNAO not found, code %d", resp.StatusCode) } content := gjson.ParseBytes(body) if status := content.Get("header.status").Int(); status != 0 { // 如果json信息返回status不为0则立刻抛出错误 - return "", errors.New(fmt.Sprintf("SauceNAO not found, status %d", status)) + return nil, fmt.Errorf("SauceNAO not found, status %d", status) } if content.Get("results.0.header.similarity").Float() < minSimilarity { - return "", errors.New("SauceNAO not found") + return nil, fmt.Errorf("SauceNAO not found") } + result := content.Get("results.0") // 正常发送 - return fmt.Sprintf( - `[SetuTime] 我有把握是这个！[CQ:image,file=%s]相似度：%s%% -标题：%s -插画ID：%d -画师：%s -画师ID：%d -直链：https://pixivel.moe/detail?id=%d`, - content.Get("results.0.header.thumbnail").Str, - content.Get("results.0.header.similarity").Str, - content.Get("results.0.data.title").Str, - content.Get("results.0.data.pixiv_id").Int(), - content.Get("results.0.data.member_name").Str, - content.Get("results.0.data.member_id").Int(), - content.Get("results.0.data.pixiv_id").Int(), - ), nil + return message.Message{ + message.Text("[SetuTime] 我有把握是这个！"), + message.Image(result.Get("header.thumbnail").Str), + message.Text( + "\n", + "相似度：", result.Get("header.similarity").Str, "\n", + "标题：", result.Get("data.title").Str, "\n", + "插画ID：", result.Get("data.pixiv_id").Int(), "\n", + "画师：", result.Get("data.member_name").Str, "\n", + "画师ID：", result.Get("data.member_id").Int(), "\n", + "直链：", "https://pixivel.moe/detail?id=", result.Get("data.pixiv_id").Int(), + ), + }, nil }