// Package saucenao P站ID/saucenao/ascii2d搜图 package saucenao import ( "fmt" "net/http" "os" "reflect" "strconv" "github.com/sirupsen/logrus" zero "github.com/wdvxdr1123/ZeroBot" "github.com/wdvxdr1123/ZeroBot/message" "github.com/FloatTech/AnimeAPI/ascii2d" "github.com/FloatTech/AnimeAPI/pixiv" "github.com/jozsefsallai/gophersauce" "github.com/FloatTech/zbputils/binary" "github.com/FloatTech/zbputils/control" "github.com/FloatTech/zbputils/ctxext" "github.com/FloatTech/zbputils/file" "github.com/FloatTech/zbputils/img/pool" ) var ( saucenaocli *gophersauce.Client ) func init() { // 插件主体 engine := control.Register("saucenao", &control.Options{ DisableOnDefault: false, Help: "搜图\n" + "- 以图搜图 | 搜索图片 | 以图识图[图片]\n" + "- 搜图[P站图片ID]", PrivateDataFolder: "saucenao", }) apikeyfile := engine.DataFolder() + "apikey.txt" if file.IsExist(apikeyfile) { key, err := os.ReadFile(apikeyfile) if err != nil { panic(err) } saucenaocli, err = gophersauce.NewClient(&gophersauce.Settings{ MaxResults: 1, APIKey: binary.BytesToString(key), }) if err != nil { panic(err) } } // 根据 PID 搜图 engine.OnRegex(`^搜图(\d+)$`).SetBlock(true). Handle(func(ctx *zero.Ctx) { id, _ := strconv.ParseInt(ctx.State["regex_matched"].([]string)[1], 10, 64) ctx.SendChain(message.Text("少女祈祷中......")) // 获取P站插图信息 illust, err := pixiv.Works(id) if err != nil { ctx.SendChain(message.Text("ERROR:", err)) return } if illust.Pid > 0 { name := strconv.FormatInt(illust.Pid, 10) var imgs message.Message for i := range illust.ImageUrls { f := file.BOTPATH + "/" + illust.Path(i) n := name + "_p" + strconv.Itoa(i) var m *pool.Image if file.IsNotExist(f) { m, err = pool.GetImage(n) if err == nil { imgs = append(imgs, message.Image(m.String())) continue } logrus.Debugln("[sausenao]开始下载", n) logrus.Debugln("[sausenao]urls:", illust.ImageUrls) err1 := illust.DownloadToCache(i) if err1 == nil { m.SetFile(f) _, _ = m.Push(ctxext.SendToSelf(ctx), ctxext.GetMessage(ctx)) } if err1 != nil { logrus.Debugln("[sausenao]下载err:", err1) } } imgs = append(imgs, message.Image("file:///"+f)) } txt := message.Text( "标题: ", illust.Title, "\n", "插画ID: ", illust.Pid, "\n", "画师: ", illust.UserName, "\n", "画师ID: ", illust.UserId, "\n", "直链: ", "https://pixivel.moe/detail?id=", illust.Pid, ) if imgs != nil { // 发送搜索结果 ctx.Send(append(imgs, message.Text("\n"), txt)) } else { // 图片下载失败,仅发送文字结果 ctx.SendChain(txt) } } else { ctx.SendChain(message.Text("图片不存在!")) } }) // 以图搜图 engine.OnKeywordGroup([]string{"以图搜图", "搜索图片", "以图识图"}, zero.OnlyGroup, zero.MustProvidePicture).SetBlock(true). Handle(func(ctx *zero.Ctx) { // 开始搜索图片 ctx.SendChain(message.Text("少女祈祷中......")) for _, pic := range ctx.State["image_url"].([]string) { if saucenaocli != nil { resp, err := saucenaocli.FromURL(pic) if err == nil && resp.Count() > 0 { result := resp.First() // 返回SauceNAO的结果 source := "" switch { case result.IsPixiv(): source = "Pixiv" case result.IsAniDB(): source = "AniDB" case result.IsBcy(): source = "Bcy" case result.IsDanbooru(): source = "Danbooru" case result.IsDeviantArt(): source = "DeviantArt" case result.IsIMDb(): source = "IMDb" case result.IsPawoo(): source = "Pawoo" case result.IsSankaku(): source = "Sankaku" case result.IsSeiga(): source = "Seiga" } if source != "" { rr := reflect.ValueOf(&result).Elem() b := binary.NewWriterF(func(w *binary.Writer) { r := rr.Type() for i := 0; i < r.NumField(); i++ { if !rr.Field(i).IsZero() { w.WriteString("\n") w.WriteString(r.Field(i).Name) w.WriteString(": ") w.WriteString(fmt.Sprint(rr.Field(i).Interface())) } } }) resp, err := http.Head(result.Header.Thumbnail) if err == nil && resp.StatusCode == http.StatusOK { ctx.SendChain( message.Text("我有把握是这个!"), message.Image(result.Header.Thumbnail), message.Text("\n图源: ", source, binary.BytesToString(b)), ) } else { ctx.SendChain( message.Text("我有把握是这个!"), message.Image(pic), message.Text("\n图源: ", source, binary.BytesToString(b)), ) } continue } } ctx.SendChain(message.Text("ERROR:", err)) } else { ctx.SendChain(message.Text("请私聊发送 设置 saucenao api key [apikey] 以启用 saucenao 搜图, key 请前往 https://saucenao.com/user.php?page=search-api 获取")) } // ascii2d 搜索 if result, err := ascii2d.Ascii2d(pic); err != nil { ctx.SendChain(message.Text("ERROR:", err)) continue } else { msg := message.Message{ctxext.FakeSenderForwardNode(ctx, message.Text("ascii2d搜图结果"))} for i := 0; i < len(result) && i < 5; i++ { msg = append(msg, ctxext.FakeSenderForwardNode(ctx, message.Image(result[i].Thumb), message.Text(fmt.Sprintf( "标题: %s\n图源: %s\n画师: %s\n画师链接: %s\n图片链接: %s", result[i].Name, result[i].Type, result[i].AuthNm, result[i].Author, result[i].Link, ))), ) } if id := ctx.SendGroupForwardMessage( ctx.Event.GroupID, msg, ).Get("message_id").Int(); id == 0 { ctx.SendChain(message.Text("ERROR:可能被风控了")) } } } }) engine.OnRegex(`^设置\s?saucenao\s?api\s?key\s?([0-9a-f]{40})$`, zero.SuperUserPermission, zero.OnlyPrivate).SetBlock(true). Handle(func(ctx *zero.Ctx) { var err error saucenaocli, err = gophersauce.NewClient(&gophersauce.Settings{ MaxResults: 1, APIKey: ctx.State["regex_matched"].([]string)[1], }) if err != nil { ctx.SendChain(message.Text("ERROR:", err)) return } err = os.WriteFile(apikeyfile, binary.StringToBytes(saucenaocli.APIKey), 0644) if err != nil { ctx.SendChain(message.Text("ERROR:", err)) return } ctx.SendChain(message.Text("成功!")) }) }