diff --git a/README.md b/README.md index db5949a..d3b2c64 100644 --- a/README.md +++ b/README.md @@ -19,3 +19,60 @@ _✨ [Nonebot2](https://github.com/nonebot/nonebot2) 用户信息获取插件 + +多平台的用户信息获取插件，可以获取用户名、用户头像等信息 + + +### 安装 + +- 使用 nb-cli + +``` +nb plugin install nonebot_plugin_userinfo +``` + +- 使用 pip + +``` +pip install nonebot_plugin_userinfo +``` + + +### 使用 + +```python +from nonebot_plugin_userinfo import get_user_info + +@matcher.handle() +async def handle(bot: Bot, event: Event): + user_info = get_user_info(bot, event, event.get_user_id()) # 获取当前事件主体用户的信息 +``` + +可以用依赖注入的方式使用： + +```python +from nonebot_plugin_userinfo import EventUserInfo, UserInfo + +@matcher.handle() +async def handle(user_info: UserInfo = EventUserInfo()): # 获取当前事件主体用户的信息 + pass +``` + +```python +from nonebot_plugin_userinfo import BotUserInfo, UserInfo + +@matcher.handle() +async def handle(user_info: UserInfo = BotUserInfo()): # 获取Bot用户信息 + pass +``` + + +### 支持的 adapter + +- [ ] TODO + + +### 鸣谢 + +- [nonebot-plugin-send-anything-anywhere](https://github.com/felinae98/nonebot-plugin-send-anything-anywhere) 项目的灵感来源以及部分实现的参考 +- [uy/sun](https://github.com/he0119) 感谢歪日佬的技术支持 diff --git a/nonebot_plugin_userinfo/__init__.py b/nonebot_plugin_userinfo/__init__.py index c97e9a6..3ea9a05 100644 --- a/nonebot_plugin_userinfo/__init__.py +++ b/nonebot_plugin_userinfo/__init__.py @@ -1 +1,4 @@ from . import adapters +from .getter import BotUserInfo, EventUserInfo, get_user_info +from .image_source import ImageSource +from .user_info import ImageData, UserInfo diff --git a/nonebot_plugin_userinfo/adapters/onebot_v11.py b/nonebot_plugin_userinfo/adapters/onebot_v11.py index e69de29..a0dc854 100644 --- a/nonebot_plugin_userinfo/adapters/onebot_v11.py +++ b/nonebot_plugin_userinfo/adapters/onebot_v11.py @@ -0,0 +1,49 @@ +from typing import Optional + +from nonebot.exception import ActionFailed +from nonebot.log import logger +from nonebot_plugin_session import SessionLevel + +from ..getter import UserInfoGetter, register_user_info_getter +from ..image_source import QQAvatar +from ..user_info import UserInfo + +try: + from nonebot.adapters.onebot.v11 import Bot, Event + + @register_user_info_getter(Bot, Event) + class Getter(UserInfoGetter[Bot, Event]): + async def _get_info(self, user_id: str) -> Optional[UserInfo]: + info = None + + if self.session.level == SessionLevel.LEVEL2: + if self.session.id2: + try: + info = await self.bot.get_group_member_info( + group_id=int(self.session.id2), user_id=int(user_id) + ) + except ActionFailed as e: + logger.warning(f"Error calling get_group_member_info: {e}") + pass + + if not info: + try: + info = await self.bot.get_stranger_info(user_id=int(user_id)) + except ActionFailed as e: + logger.warning(f"Error calling get_stranger_info failed: {e}") + pass + + if info: + name = info.get("nickname", "") + card = info.get("card") + gender = info.get("sex", "unknown") + return UserInfo( + user_id=user_id, + user_name=name, + user_displayname=card, + user_avatar=QQAvatar(int(user_id)), + user_gender=gender, + ) + +except ImportError: + pass diff --git a/nonebot_plugin_userinfo/adapters/onebot_v12.py b/nonebot_plugin_userinfo/adapters/onebot_v12.py new file mode 100644 index 0000000..38c3119 --- /dev/null +++ b/nonebot_plugin_userinfo/adapters/onebot_v12.py @@ -0,0 +1,87 @@ +from typing import Optional + +from nonebot.exception import ActionFailed +from nonebot.log import logger +from nonebot_plugin_session import SessionLevel + +from ..getter import UserInfoGetter, register_user_info_getter +from ..image_source import QQAvatar +from ..user_info import UserInfo + +try: + from nonebot.adapters.onebot.v12 import Bot, Event + + @register_user_info_getter(Bot, Event) + class Getter(UserInfoGetter[Bot, Event]): + async def _get_info(self, user_id: str) -> Optional[UserInfo]: + avatar = None + if self.bot.platform == "qq": + avatar = QQAvatar(qq=int(user_id)) + + info = None + + if self.session.level == SessionLevel.LEVEL3: + if self.session.id3: + if self.session.id2: + try: + info = await self.bot.get_channel_member_info( + guild_id=self.session.id3, + channel_id=self.session.id2, + user_id=user_id, + ) + except ActionFailed as e: + logger.warning( + f"Error calling get_channel_member_info: {e}" + ) + pass + else: + try: + info = await self.bot.get_guild_member_info( + guild_id=self.session.id3, user_id=user_id + ) + except ActionFailed as e: + logger.warning(f"Error calling get_guild_member_info: {e}") + pass + + platform = self.bot.platform + impl = self.bot.impl + if platform == "qqguild" and impl == "nonebot-plugin-all4one": + # 先转成 dict，这样就算以后用扩展模型也不会出错 + event_dict = self.event.dict() + try: + if user_id == str(event_dict["qqguild"]["author"]["id"]): + avatar = str(event_dict["qqguild"]["author"]["avatar"]) + if not avatar and info: + avatar = str(info["qqguild"]["user"]["avatar"]) # type: ignore + except KeyError: + pass + + elif self.session.level == SessionLevel.LEVEL2: + if self.session.id2: + info = await self.bot.get_group_member_info( + group_id=self.session.id2, user_id=user_id + ) + + if not info: + try: + info = await self.bot.get_user_info(user_id=user_id) + except ActionFailed as e: + logger.warning(f"Error calling get_user_info: {e}") + pass + + if info: + user_name = info["user_name"] + user_displayname = info["user_displayname"] + user_remark = info.get("user_remark") + gender = "unknown" + return UserInfo( + user_id=user_id, + user_name=user_name, + user_displayname=user_displayname, + user_remark=user_remark, + user_avatar=avatar, + user_gender=gender, + ) + +except ImportError: + pass diff --git a/nonebot_plugin_userinfo/exception.py b/nonebot_plugin_userinfo/exception.py new file mode 100644 index 0000000..30babf1 --- /dev/null +++ b/nonebot_plugin_userinfo/exception.py @@ -0,0 +1,2 @@ +class NetworkError(Exception): + pass diff --git a/nonebot_plugin_userinfo/getter.py b/nonebot_plugin_userinfo/getter.py new file mode 100644 index 0000000..abd7ee9 --- /dev/null +++ b/nonebot_plugin_userinfo/getter.py @@ -0,0 +1,82 @@ +from typing import Generic, List, NamedTuple, Optional, Type, TypeVar + +from cachetools import TTLCache +from nonebot import require +from nonebot.adapters import Bot, Event +from nonebot.params import Depends + +require("nonebot_plugin_session") +from nonebot_plugin_session import SessionIdType, extract_session + +from .user_info import UserInfo + +_user_info_cache = TTLCache(maxsize=float("inf"), ttl=10) + +B = TypeVar("B", bound=Bot) +E = TypeVar("E", bound=Event) + + +class UserInfoGetter(Generic[B, E]): + def __init__(self, bot: B, event: E): + self.bot = bot + self.event = event + self.session = extract_session(bot, event) + + async def _get_info(self, user_id: str) -> Optional[UserInfo]: + raise NotImplementedError + + async def get_info( + self, user_id: str, use_cache: bool = True + ) -> Optional[UserInfo]: + if use_cache: + session_id = self.session.get_id(SessionIdType.GROUP_USER) + id = f"{session_id}_{user_id}" + if id in _user_info_cache: + return _user_info_cache[id] + info = await self._get_info(user_id) + if info: + _user_info_cache[id] = info + return info + else: + return await self._get_info(user_id) + + +class UserInfoGetterTuple(NamedTuple): + bot: Type[Bot] + event: Type[Event] + getter: Type[UserInfoGetter] + + +_user_info_getters: List[UserInfoGetterTuple] = [] + + +def register_user_info_getter(bot: Type[Bot], event: Type[Event]): + def wrapper(getter: Type[UserInfoGetter]): + _user_info_getters.append(UserInfoGetterTuple(bot, event, getter)) + return getter + + return wrapper + + +async def get_user_info( + bot: Bot, event: Event, user_id: str, use_cache: bool = True +) -> Optional[UserInfo]: + for getter_tuple in _user_info_getters: + if isinstance(bot, getter_tuple.bot) and isinstance(event, getter_tuple.event): + return await getter_tuple.getter(bot, event).get_info( + user_id, use_cache=use_cache + ) + + +def EventUserInfo(use_cache: bool = True): + async def dependency(bot: Bot, event: Event) -> Optional[UserInfo]: + return await get_user_info(bot, event, event.get_user_id(), use_cache=use_cache) + + return Depends(dependency) + + +def BotUserInfo(use_cache: bool = True): + async def dependency(bot: Bot, event: Event) -> Optional[UserInfo]: + return await get_user_info(bot, event, bot.self_id, use_cache=use_cache) + + return Depends(dependency) diff --git a/nonebot_plugin_userinfo/image_source.py b/nonebot_plugin_userinfo/image_source.py new file mode 100644 index 0000000..e4f6be9 --- /dev/null +++ b/nonebot_plugin_userinfo/image_source.py @@ -0,0 +1,22 @@ +import hashlib +from dataclasses import dataclass + +from .utils import download_url + + +class ImageSource: + async def get_image(self) -> bytes: + raise NotImplementedError + + +@dataclass +class QQAvatar(ImageSource): + qq: int + + async def get_image(self) -> bytes: + url = f"http://q1.qlogo.cn/g?b=qq&nk={self.qq}&s=640" + data = await download_url(url) + if hashlib.md5(data).hexdigest() == "acef72340ac0e914090bd35799f5594e": + url = f"http://q1.qlogo.cn/g?b=qq&nk={self.qq}&s=100" + data = await download_url(url) + return data diff --git a/nonebot_plugin_userinfo/user_info.py b/nonebot_plugin_userinfo/user_info.py index 8537aef..d132c50 100644 --- a/nonebot_plugin_userinfo/user_info.py +++ b/nonebot_plugin_userinfo/user_info.py @@ -4,7 +4,9 @@ from pydantic import BaseModel -ImageData = Union[str, bytes, Path, BytesIO] +from .image_source import ImageSource + +ImageData = Union[str, bytes, Path, BytesIO, ImageSource] class UserInfo(BaseModel): diff --git a/nonebot_plugin_userinfo/utils.py b/nonebot_plugin_userinfo/utils.py new file mode 100644 index 0000000..d17a310 --- /dev/null +++ b/nonebot_plugin_userinfo/utils.py @@ -0,0 +1,19 @@ +import asyncio + +import httpx +from nonebot.log import logger + +from .exception import NetworkError + + +async def download_url(url: str) -> bytes: + async with httpx.AsyncClient() as client: + for i in range(3): + try: + resp = await client.get(url, timeout=10) + resp.raise_for_status() + return resp.content + except Exception as e: + logger.warning(f"Error downloading {url}, retry {i}/3: {e}") + await asyncio.sleep(3) + raise NetworkError(f"{url} 下载失败！") diff --git a/pyproject.toml b/pyproject.toml index 21992f3..d65f6d1 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -11,6 +11,9 @@ repository = "https://github.com/noneplugin/nonebot-plugin-userinfo" [tool.poetry.dependencies] python = "^3.8" nonebot2 = { version = "^2.0.0", extras = ["fastapi"] } +nonebot-plugin-session = ">=0.0.5,<0.1.0" +httpx = ">=0.20.0,<1.0.0" +cachetools = "^5.0.0" [tool.poetry.group.dev.dependencies] black = "^22.8.0" diff --git a/tests/__init__.py b/tests/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 0000000..4e0ef5f --- /dev/null +++ b/tests/conftest.py @@ -0,0 +1,7 @@ +import pytest +from nonebug import App + + +@pytest.fixture +async def app(): + yield App() diff --git a/tests/utils.py b/tests/utils.py new file mode 100644 index 0000000..6bb0ad2 --- /dev/null +++ b/tests/utils.py @@ -0,0 +1,22 @@ +from typing import TYPE_CHECKING, Optional + +if TYPE_CHECKING: + from nonebot_plugin_userinfo import ImageData, UserInfo + + +def assert_user_info( + user_info: "UserInfo", + *, + user_id: str, + user_name: str, + user_displayname: Optional[str] = None, + user_remark: Optional[str] = None, + user_avatar: Optional[ImageData] = None, + user_gender: str = "unknown" +): + assert user_info.user_id == user_id + assert user_info.user_name == user_name + assert user_info.user_displayname == user_displayname + assert user_info.user_remark == user_remark + assert user_info.user_avatar == user_avatar + assert user_info.user_gender == user_gender