zhenxun_bot/plugins/nonebot_plugin_picsearcher/__init__.py

# -*- coding: utf-8 -*-
import re
import traceback
import aiohttp
from typing import Dict

from aiohttp.client_exceptions import ClientError
from nonebot.plugin import on_command, on_message
from nonebot.adapters.cqhttp import Bot, MessageEvent, GroupMessageEvent
from nonebot.typing import T_State
from services.log import logger
from utils.utils import get_message_text, get_message_imgs
from configs.config import MAX_FIND_IMG_COUNT
from nonebot.rule import to_me

from .ex import get_des as get_des_ex
from .iqdb import get_des as get_des_iqdb
from .saucenao import get_des as get_des_sau
from .ascii2d import get_des as get_des_asc
from .trace import get_des as get_des_trace
from .yandex import get_des as get_des_yandex
from .google import Google
from nonebot.adapters.cqhttp import MessageSegment

__plugin_name__ = "识图"

__plugin_usage__ = f"""用法：识图 [参数](默认1) [图片]
    参数列表：
    1. google
    2. baidu(计划中)
    3. iqdb
    4. ex
    5. trace
    6. yandex
    7. asc
    8. sau"""


async def get_des(url: str, mode: str, user_id: int):
    """
    :param url: 图片链接
    :param mode: 图源
    :return:
    """
    if mode == "3":
        async for msg in get_des_iqdb(url):
            yield msg
    elif mode == "4":
        async for msg in get_des_ex(url):
            yield msg
    elif mode == "5":
        async for msg in get_des_trace(url):
            yield msg
    elif mode == "6":
        async for msg in get_des_yandex(url):
            yield msg
    elif mode.startswith("7"):
        async for msg in get_des_asc(url, user_id):
            yield msg
    elif mode.startswith("8"):
        async for msg in get_des_sau(url, user_id):
            yield msg
    else:
        google = Google()
        res = google.search(url)
        a = 0
        for i in range(len(res.raw)):
            if res.raw[i].thumbnail != 'No directable url' and res.raw[i].thumbnail != '':
                if 'google.com' in res.raw[i].thumbnail:
                    temp_url = re.search(r'(imgurl)[\=](.*)(&imgrefurl)', res.raw[i].thumbnail).group()
                    img_url = re.search(r'(http).*[^&imgrefurl]', temp_url).group()
                    msg = f"{res.raw[i].title}\n" + MessageSegment.image(img_url) + f"\n来源:{res.raw[i].url}"
                    logger.info(f"原图：{img_url}  谷歌图片：{res.raw[i].thumbnail}")
                    if not await check_status(img_url):
                        continue
                    else:
                        yield msg
                    a += 1
                else:
                    msg = f"{res.raw[i].title}\n" + MessageSegment.image(res.raw[i].thumbnail) + f"\n来源:{res.raw[i].url}"
                    logger.info(f"{res.raw[i].thumbnail}")
                    if not await check_status(res.raw[i].thumbnail):
                        continue
                    else:
                        yield msg
                    a += 1
            else:
                continue
            if a == 3:
                break
        if a == 0:
            await setu.finish('该图暂无法识别..')


async def check_status(url) -> str:
    async with aiohttp.ClientSession() as session:
        async with session.get(url, timeout=3) as r:
            if r.status != 200:
                return False
            else:
                return True

setu = on_command("识图", aliases={"search"}, block=True, priority=5)


@setu.handle()
async def handle_first_receive(bot: Bot, event: MessageEvent, state: T_State):
    msg = get_message_text(event.json())
    imgs = get_message_imgs(event.json())
    if msg in ["帮助"]:
        await setu.finish("示例：\n\t识图 (图片)\n\t识图 2 (图片)")
    if imgs:
        state["setu"] = imgs[0]
    if msg:
        state["mod"] = msg


# ex/nao/trace/iqdb/ascii2d
# @setu.got("mod", prompt="从哪里查找呢? ex/nao/trace/iqdb/ascii2d")
# async def get_func(bot: Bot, event: MessageEvent, state: dict):
#     pass


@setu.args_parser
async def get_setu(bot: Bot, event: MessageEvent, state: T_State):
    imgs = get_message_imgs(event.json())
    msg = get_message_text(event.json())
    if not imgs:
        await setu.reject()
    if msg:
        state["mod"] = msg
    state["setu"] = imgs[0]


@setu.got("setu", prompt="图呢？")
async def get_setu(bot: Bot, event: MessageEvent, state: T_State):
    """
    发现没有的时候要发问
    :return:
    """
    url: str = state["setu"]
    mod: str = state["mod"] if state.get("mod") else ""  # 模式
    try:
        await bot.send(event=event, message="正在处理图片")
        idx = 1
        async for msg in get_des(url, mod, event.user_id):
            await bot.send(event=event, message=msg)
            if idx == MAX_FIND_IMG_COUNT:
                break
            idx += 1
        logger.info(
            f"(USER {event.user_id}, GROUP "
            f"{event.group_id if event.message_type != 'private' else 'private'}) 识图:{url}"
        )
        # image_data: List[Tuple] = await get_pic_from_url(url)
        # await setu.finish("hso")
    except IndexError:
        # await bot.send(event, traceback.format_exc())
        await setu.finish("参数错误")
    except ClientError:
        await setu.finish("连接失败")


pic_map: Dict[str, str] = {}  # 保存这个群的其阿金一张色图 {"123456":http://xxx"}


async def check_pic(bot: Bot, event: MessageEvent, state: T_State) -> bool:
    if isinstance(event, MessageEvent):
        for msg in event.message:
            if msg.type == "image":
                url: str = msg.data["url"]
                state["url"] = url
                return True
        return False


notice_pic = on_message(check_pic, block=False, priority=1)


@notice_pic.handle()
async def handle_pic(bot: Bot, event: GroupMessageEvent, state: T_State):
    try:
        group_id: str = str(event.group_id)
        pic_map.update({group_id: state["url"]})
    except AttributeError:
        pass


previous = on_command("上一张图是什么", aliases={"上一张", "这是什么"}, rule=to_me(), block=True)


@previous.handle()
async def handle_previous(bot: Bot, event: GroupMessageEvent, state: T_State):
    await bot.send(event=event, message="processing...")
    try:
        url: str = pic_map[str(event.group_id)]
        idx = 1
        async for msg in get_des(url, "nao", event.user_id):
            await bot.send(event=event, message=msg)
            if idx == MAX_FIND_IMG_COUNT:
                break
            idx += 1
    except IndexError:
        await previous.finish("参数错误")
    except ClientError:
        await previous.finish("连接错误")
    except KeyError:
        await previous.finish("没有图啊QAQ")