⚡️ 更好的识图触发模式

kanrichan · kanrichan · commit a960b7c732fe · 2021-04-16T19:18:29.000+08:00
diff --git 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))
+					ctx.SendChain(message.Text("ERROR: ", 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
+}
+
+// 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
 					}
 				}
 			}
-			if len(pics) == 0 {
-				ctx.Send("没有收到图片，搜图结束......")
-				return
+		}
+		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"])
 			}
-			// 开始搜索图片
-			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))
-				}
+		}
+		// 如果有图片就返回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
-		})
+			return false
+		case <-time.After(time.Second * 60):
+			return false
+		}
+	}
 }
diff --git a/setutime/utils/saucenao.go b/setutime/utils/saucenao.go
@@ -1,84 +1,77 @@
 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"
 
 		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
 }
