From 2349350ba8db8a572a49e931891eaa42a4207df4 Mon Sep 17 00:00:00 2001
From: Kyomotoi <1172294279@qq.com>
Date: Thu, 15 Apr 2021 14:15:25 +0800
Subject: =?UTF-8?q?=E2=9C=A8=F0=9F=90=9B=E2=9A=A1=EF=B8=8F=20=E4=B8=80?=
 =?UTF-8?q?=E4=BA=9B=E6=94=B9=E5=8A=A8?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

新增：以图搜图
挖坑：今天吃什么

修复：只响应私聊消息
修复：nsfw部分变量错误

优化：插件名
优化：部分插件代码结构
---
 ATRI/plugins/saucenao/__init__.py    | 70 ++++++++++++++++++++++++++++++++++++
 ATRI/plugins/saucenao/data_source.py |  6 ++--
 2 files changed, 72 insertions(+), 4 deletions(-)

(limited to 'ATRI/plugins/saucenao')

diff --git a/ATRI/plugins/saucenao/__init__.py b/ATRI/plugins/saucenao/__init__.py
index c36fde1..f158cf9 100644
--- a/ATRI/plugins/saucenao/__init__.py
+++ b/ATRI/plugins/saucenao/__init__.py
@@ -1,7 +1,17 @@
+import re
+import json
+from random import choice
+
 from nonebot.adapters.cqhttp import Bot, MessageEvent
+from nonebot.adapters.cqhttp.message import Message, MessageSegment
+from nonebot.typing import T_State
 
+from ATRI.config import Config
 from ATRI.service import Service as sv
 from ATRI.rule import is_in_service
+from ATRI.exceptions import RequestTimeOut
+
+from .data_source import SauceNao
 
 
 __doc__ = """
@@ -10,3 +20,63 @@ __doc__ = """
 用法：
   以图搜图 (pic)
 """
+
+saucenao = sv.on_command(
+  cmd='以图搜图',
+  docs=__doc__,
+  rule=is_in_service('以图搜图')
+)
+
+@saucenao.args_parser  # type: ignore
+async def _load_saucenao(bot: Bot, event: MessageEvent,
+                         state: T_State) -> None:
+    msg = str(event.message)
+    quit_list = ['算了', '罢了', '不搜了']
+    if msg in quit_list:
+        await saucenao.finish('好吧...')
+    
+    if not msg:
+        await saucenao.reject('图呢？')
+    else:
+        state['pic_sau'] = msg
+
+@saucenao.handle()
+async def _sauce_nao(bot: Bot, event: MessageEvent, state: T_State) -> None:
+    msg = str(event.message).strip()
+    if msg:
+        state['pic_sau'] = msg
+
+@saucenao.got('pic_sau', prompt='图呢？')
+async def _deal_saucenao(bot: Bot, event: MessageEvent, state: T_State) -> None:
+    msg = state['pic_sau']
+    img = re.findall(r'url=(.*?)]', msg)
+    if not img:
+        await saucenao.finish('请发送图片而不是其他东西！！')
+    
+    try:
+        task = SauceNao(api_key=Config.SauceNAO.key)
+        data = json.loads(await task.search(img[0]))
+    except RequestTimeOut:
+        raise RequestTimeOut('Request failed!')
+    
+    res = data['results']
+    result = list()
+    for i in range(0, 3):
+        data = res[i]
+        
+        _result = dict()
+        _result['similarity'] = data['header']['similarity']
+        _result['index_name'] = data['header']['index_name']
+        _result['url'] = choice(data['data'].get('ext_urls', ['None']))
+        result.append(_result)
+    
+    msg0 = f"> {MessageSegment.at(event.user_id)}"
+    for i in result:
+        msg0 = msg0 + (
+            "\n——————————\n"
+            f"Similarity: {i['similarity']}\n"
+            f"Name: {i['index_name']}\n"
+            f"URL: {i['url'].replace('https://', '')}"
+        )
+    
+    await saucenao.finish(Message(msg0))
diff --git a/ATRI/plugins/saucenao/data_source.py b/ATRI/plugins/saucenao/data_source.py
index 636c540..cd88554 100644
--- a/ATRI/plugins/saucenao/data_source.py
+++ b/ATRI/plugins/saucenao/data_source.py
@@ -8,16 +8,14 @@ class SauceNao:
     def __init__(self,
                  api_key: str,
                  output_type=2,
-                 testmode=0,
-                 dbmask=None,
+                 testmode=1,
                  dbmaski=32768,
                  db=5,
-                 numres=1) -> None:
+                 numres=5) -> None:
         params = dict()
         params['api_key'] = api_key
         params['output_type'] = output_type
         params['testmode'] = testmode
-        params['dbmask'] = dbmask
         params['dbmaski'] = dbmaski
         params['db'] = db
         params['numres'] = numres
-- 
cgit v1.2.3