From 275bc0cb535d151488c247bd7f6cfb411b5026e4 Mon Sep 17 00:00:00 2001 From: suyiiyii Date: Sun, 8 Sep 2024 18:21:57 +0800 Subject: [PATCH] =?UTF-8?q?:recycles:=20=E6=B3=A8=E9=87=8A=E6=8E=89cookie?= =?UTF-8?q?=E7=9B=B8=E5=85=B3=E4=BB=A3=E7=A0=81=EF=BC=8C=E4=BD=BF=E5=BE=97?= =?UTF-8?q?bison=E5=8F=AF=E4=BB=A5=E6=AD=A3=E5=B8=B8=E8=BF=90=E8=A1=8C?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- nonebot_bison/config/db_config.py | 248 +++++++++++++------------- nonebot_bison/sub_manager/__init__.py | 78 ++++---- nonebot_bison/utils/site.py | 219 ++++++++++++----------- 3 files changed, 275 insertions(+), 270 deletions(-) diff --git a/nonebot_bison/config/db_config.py b/nonebot_bison/config/db_config.py index 5cb0152..4fdea77 100644 --- a/nonebot_bison/config/db_config.py +++ b/nonebot_bison/config/db_config.py @@ -259,130 +259,130 @@ class DBConfig: ) return res - async def get_cookie(self, platform_name: str = None, target: T_Target = None) -> Sequence[Cookie]: - """根据平台名和订阅名获取 cookie，不会返回匿名cookie""" - async with create_session() as sess: - query = select(Cookie).distinct().where(Cookie.is_universal == False) # noqa: E712 - if platform_name: - query = query.where(Cookie.platform_name == platform_name) - query = query.outerjoin(CookieTarget).options(selectinload(Cookie.targets)) - res = (await sess.scalars(query)).all() - if target: - query = select(CookieTarget.cookie_id).join(Target).where(Target.target == target) - ids = set((await sess.scalars(query)).all()) - res = [cookie for cookie in res if cookie.id in ids] - return res - - async def get_unviersal_cookie(self, platform_name: str = None) -> Sequence[Cookie]: - async with create_session() as sess: - query = select(Cookie).distinct().where(Cookie.is_universal == True) # noqa: E712 - if platform_name: - query = query.where(Cookie.platform_name == platform_name) - res = (await sess.scalars(query)).all() - return res - - async def add_cookie_with_content(self, platform_name: str, content: str) -> int: - async with create_session() as sess: - cookie = Cookie(platform_name=platform_name, content=content) - sess.add(cookie) - await sess.commit() - await sess.refresh(cookie) - return cookie.id - - async def add_cookie(self, cookie: Cookie) -> int: - async with create_session() as sess: - sess.add(cookie) - await sess.commit() - await sess.refresh(cookie) - return cookie.id - - async def update_cookie(self, cookie: Cookie): - async with create_session() as sess: - cookie_in_db: Cookie | None = await sess.scalar(select(Cookie).where(Cookie.id == cookie.id)) - if not cookie_in_db: - return - cookie_in_db.content = cookie.content - cookie_in_db.last_usage = cookie.last_usage - cookie_in_db.status = cookie.status - cookie_in_db.tags = cookie.tags - await sess.commit() - - async def delete_cookie_by_id(self, cookie_id: int): - async with create_session() as sess: - cookie = await sess.scalar( - select(Cookie) - .where(Cookie.id == cookie_id) - .outerjoin(CookieTarget) - .options(selectinload(Cookie.targets)) - ) - if len(cookie.targets) > 0: - raise Exception(f"cookie {cookie.id} in use") - await sess.execute(delete(Cookie).where(Cookie.id == cookie_id)) - await sess.commit() - - async def get_cookie_by_target(self, target: T_Target, platform_name: str) -> Sequence[Cookie]: - async with create_session() as sess: - query = ( - select(Cookie) - .join(CookieTarget) - .join(Target) - .where(Target.platform_name == platform_name, Target.target == target) - ) - return (await sess.scalars(query)).all() - - async def get_universal_cookie(self, platform_name: str) -> Sequence[Cookie]: - async with create_session() as sess: - query = ( - select(Cookie) - .where(Cookie.platform_name == platform_name) - .where(Cookie.is_universal == True) # noqa: E712 - ) - return (await sess.scalars(query)).all() - - async def add_cookie_target(self, target: T_Target, platform_name: str, cookie_id: int): - async with create_session() as sess: - target_obj = await sess.scalar( - select(Target).where(Target.platform_name == platform_name, Target.target == target) - ) - # check if relation exists - cookie_target = await sess.scalar( - select(CookieTarget).where(CookieTarget.target == target_obj, CookieTarget.cookie_id == cookie_id) - ) - if cookie_target: - raise DuplicateCookieTargetException() - cookie_obj = await sess.scalar(select(Cookie).where(Cookie.id == cookie_id)) - cookie_target = CookieTarget(target=target_obj, cookie=cookie_obj) - sess.add(cookie_target) - await sess.commit() - - async def delete_cookie_target(self, target: T_Target, platform_name: str, cookie_id: int): - async with create_session() as sess: - target_obj = await sess.scalar( - select(Target).where(Target.platform_name == platform_name, Target.target == target) - ) - cookie_obj = await sess.scalar(select(Cookie).where(Cookie.id == cookie_id)) - await sess.execute( - delete(CookieTarget).where(CookieTarget.target == target_obj, CookieTarget.cookie == cookie_obj) - ) - await sess.commit() - - async def delete_cookie_target_by_id(self, cookie_target_id: int): - async with create_session() as sess: - await sess.execute(delete(CookieTarget).where(CookieTarget.id == cookie_target_id)) - await sess.commit() - - async def get_cookie_target(self) -> list[CookieTarget]: - async with create_session() as sess: - query = ( - select(CookieTarget) - .outerjoin(Target) - .options(selectinload(CookieTarget.target)) - .outerjoin(Cookie) - .options(selectinload(CookieTarget.cookie)) - ) - res = list((await sess.scalars(query)).all()) - res.sort(key=lambda x: (x.target.platform_name, x.cookie_id, x.target_id)) - return res + # async def get_cookie(self, platform_name: str = None, target: T_Target = None) -> Sequence[Cookie]: + # """根据平台名和订阅名获取 cookie，不会返回匿名cookie""" + # async with create_session() as sess: + # query = select(Cookie).distinct().where(Cookie.is_universal == False) # noqa: E712 + # if platform_name: + # query = query.where(Cookie.platform_name == platform_name) + # query = query.outerjoin(CookieTarget).options(selectinload(Cookie.targets)) + # res = (await sess.scalars(query)).all() + # if target: + # query = select(CookieTarget.cookie_id).join(Target).where(Target.target == target) + # ids = set((await sess.scalars(query)).all()) + # res = [cookie for cookie in res if cookie.id in ids] + # return res + # + # async def get_unviersal_cookie(self, platform_name: str = None) -> Sequence[Cookie]: + # async with create_session() as sess: + # query = select(Cookie).distinct().where(Cookie.is_universal == True) # noqa: E712 + # if platform_name: + # query = query.where(Cookie.platform_name == platform_name) + # res = (await sess.scalars(query)).all() + # return res + # + # async def add_cookie_with_content(self, platform_name: str, content: str) -> int: + # async with create_session() as sess: + # cookie = Cookie(platform_name=platform_name, content=content) + # sess.add(cookie) + # await sess.commit() + # await sess.refresh(cookie) + # return cookie.id + # + # async def add_cookie(self, cookie: Cookie) -> int: + # async with create_session() as sess: + # sess.add(cookie) + # await sess.commit() + # await sess.refresh(cookie) + # return cookie.id + # + # async def update_cookie(self, cookie: Cookie): + # async with create_session() as sess: + # cookie_in_db: Cookie | None = await sess.scalar(select(Cookie).where(Cookie.id == cookie.id)) + # if not cookie_in_db: + # return + # cookie_in_db.content = cookie.content + # cookie_in_db.last_usage = cookie.last_usage + # cookie_in_db.status = cookie.status + # cookie_in_db.tags = cookie.tags + # await sess.commit() + # + # async def delete_cookie_by_id(self, cookie_id: int): + # async with create_session() as sess: + # cookie = await sess.scalar( + # select(Cookie) + # .where(Cookie.id == cookie_id) + # .outerjoin(CookieTarget) + # .options(selectinload(Cookie.targets)) + # ) + # if len(cookie.targets) > 0: + # raise Exception(f"cookie {cookie.id} in use") + # await sess.execute(delete(Cookie).where(Cookie.id == cookie_id)) + # await sess.commit() + # + # async def get_cookie_by_target(self, target: T_Target, platform_name: str) -> Sequence[Cookie]: + # async with create_session() as sess: + # query = ( + # select(Cookie) + # .join(CookieTarget) + # .join(Target) + # .where(Target.platform_name == platform_name, Target.target == target) + # ) + # return (await sess.scalars(query)).all() + # + # async def get_universal_cookie(self, platform_name: str) -> Sequence[Cookie]: + # async with create_session() as sess: + # query = ( + # select(Cookie) + # .where(Cookie.platform_name == platform_name) + # .where(Cookie.is_universal == True) # noqa: E712 + # ) + # return (await sess.scalars(query)).all() + # + # async def add_cookie_target(self, target: T_Target, platform_name: str, cookie_id: int): + # async with create_session() as sess: + # target_obj = await sess.scalar( + # select(Target).where(Target.platform_name == platform_name, Target.target == target) + # ) + # # check if relation exists + # cookie_target = await sess.scalar( + # select(CookieTarget).where(CookieTarget.target == target_obj, CookieTarget.cookie_id == cookie_id) + # ) + # if cookie_target: + # raise DuplicateCookieTargetException() + # cookie_obj = await sess.scalar(select(Cookie).where(Cookie.id == cookie_id)) + # cookie_target = CookieTarget(target=target_obj, cookie=cookie_obj) + # sess.add(cookie_target) + # await sess.commit() + # + # async def delete_cookie_target(self, target: T_Target, platform_name: str, cookie_id: int): + # async with create_session() as sess: + # target_obj = await sess.scalar( + # select(Target).where(Target.platform_name == platform_name, Target.target == target) + # ) + # cookie_obj = await sess.scalar(select(Cookie).where(Cookie.id == cookie_id)) + # await sess.execute( + # delete(CookieTarget).where(CookieTarget.target == target_obj, CookieTarget.cookie == cookie_obj) + # ) + # await sess.commit() + # + # async def delete_cookie_target_by_id(self, cookie_target_id: int): + # async with create_session() as sess: + # await sess.execute(delete(CookieTarget).where(CookieTarget.id == cookie_target_id)) + # await sess.commit() + # + # async def get_cookie_target(self) -> list[CookieTarget]: + # async with create_session() as sess: + # query = ( + # select(CookieTarget) + # .outerjoin(Target) + # .options(selectinload(CookieTarget.target)) + # .outerjoin(Cookie) + # .options(selectinload(CookieTarget.cookie)) + # ) + # res = list((await sess.scalars(query)).all()) + # res.sort(key=lambda x: (x.target.platform_name, x.cookie_id, x.target_id)) + # return res config = DBConfig() diff --git a/nonebot_bison/sub_manager/__init__.py b/nonebot_bison/sub_manager/__init__.py index 943efd1..ec58e30 100644 --- a/nonebot_bison/sub_manager/__init__.py +++ b/nonebot_bison/sub_manager/__init__.py @@ -14,10 +14,10 @@ from nonebot.adapters.onebot.v11.event import PrivateMessageEvent from .add_sub import do_add_sub from .del_sub import do_del_sub from .query_sub import do_query_sub -from .add_cookie import do_add_cookie -from .del_cookie import do_del_cookie -from .add_cookie_target import do_add_cookie_target -from .del_cookie_target import do_del_cookie_target +# from .add_cookie import do_add_cookie +# from .del_cookie import do_del_cookie +# from .add_cookie_target import do_add_cookie_target +# from .del_cookie_target import do_del_cookie_target from .utils import common_platform, admin_permission, gen_handle_cancel, configurable_to_me, set_target_user_info add_sub_matcher = on_command( @@ -44,41 +44,41 @@ del_sub_matcher = on_command( del_sub_matcher.handle()(set_target_user_info) do_del_sub(del_sub_matcher) -add_cookie_matcher = on_command( - "添加cookie", - rule=configurable_to_me, - permission=SUPERUSER, - priority=5, - block=True, -) -do_add_cookie(add_cookie_matcher) - -add_cookie_target_matcher = on_command( - "关联cookie", - rule=configurable_to_me, - permission=SUPERUSER, - priority=5, - block=True, -) -do_add_cookie_target(add_cookie_target_matcher) - -del_cookie_target_matcher = on_command( - "取消关联cookie", - rule=configurable_to_me, - permission=SUPERUSER, - priority=5, - block=True, -) -do_del_cookie_target(del_cookie_target_matcher) - -del_cookie_matcher = on_command( - "删除cookie", - rule=configurable_to_me, - permission=SUPERUSER, - priority=5, - block=True, -) -do_del_cookie(del_cookie_matcher) +# add_cookie_matcher = on_command( +# "添加cookie", +# rule=configurable_to_me, +# permission=SUPERUSER, +# priority=5, +# block=True, +# ) +# do_add_cookie(add_cookie_matcher) +# +# add_cookie_target_matcher = on_command( +# "关联cookie", +# rule=configurable_to_me, +# permission=SUPERUSER, +# priority=5, +# block=True, +# ) +# do_add_cookie_target(add_cookie_target_matcher) +# +# del_cookie_target_matcher = on_command( +# "取消关联cookie", +# rule=configurable_to_me, +# permission=SUPERUSER, +# priority=5, +# block=True, +# ) +# do_del_cookie_target(del_cookie_target_matcher) +# +# del_cookie_matcher = on_command( +# "删除cookie", +# rule=configurable_to_me, +# permission=SUPERUSER, +# priority=5, +# block=True, +# ) +# do_del_cookie(del_cookie_matcher) group_manage_matcher = on_command("群管理", rule=to_me(), permission=SUPERUSER, priority=4, block=True) diff --git a/nonebot_bison/utils/site.py b/nonebot_bison/utils/site.py index b53aab2..b93ff56 100644 --- a/nonebot_bison/utils/site.py +++ b/nonebot_bison/utils/site.py @@ -45,117 +45,122 @@ class DefaultClientManager(ClientManager): def is_cookie_client_manager(manger: type[ClientManager]) -> bool: return hasattr(manger, "_cookie_client_manger_") - +# TODO: Implement CookieClientManager class CookieClientManager(ClientManager): - _cookie_client_manger_ = True +# _cookie_client_manger_ = True _platform_name: str - _cookie_cd: int = 10 +# _cookie_cd: int = 10 +# +# @classmethod +# async def init_universal_cookie(cls): +# """移除已有的匿名cookie，添加一个新的匿名cookie""" +# universal_cookies = await config.get_unviersal_cookie(cls._platform_name) +# universal_cookie = Cookie( +# platform_name=cls._platform_name, content="{}", is_universal=True, tags={"temporary": True} +# ) +# for cookie in universal_cookies: +# if not cookie.tags.get("temporary"): +# continue +# await config.delete_cookie_by_id(cookie.id) +# universal_cookie.id = cookie.id # 保持原有的id +# await config.add_cookie(universal_cookie) +# +# @classmethod +# async def init_cookie(cls, cookie: Cookie) -> Cookie: +# """初始化 cookie，添加用户 cookie 时使用""" +# cookie.cd = cls._cookie_cd +# cookie.last_usage = datetime.now() # 使得优先使用用户 cookie +# return cookie +# +# @classmethod +# async def valid_cookie(cls, content: str) -> bool: +# """验证 cookie 内容是否有效，添加 cookie 时用，可根据平台的具体情况进行重写""" +# try: +# data = json.loads(content) +# if not isinstance(data, dict): +# return False +# except JSONDecodeError: +# return False +# return True +# +# @classmethod +# async def get_cookie_friendly_name(cls, cookie: Cookie) -> str: +# """获取 cookie 的友好名字，用于展示""" +# from . import text_fletten +# +# return text_fletten(f"{cookie.platform_name} [{cookie.content[:10]}]") +# +# def _generate_hook(self, cookie: Cookie) -> callable: +# """hook 函数生成器，用于回写请求状态到数据库""" +# +# async def _response_hook(resp: httpx.Response): +# if resp.status_code == 200: +# logger.debug(f"请求成功: {cookie.id} {resp.request.url}") +# cookie.status = "success" +# else: +# logger.error(f"请求失败:{cookie.id} {resp.request.url}, 状态码: {resp.status_code}") +# cookie.status = "failed" +# cookie.last_usage = datetime.now() +# await config.update_cookie(cookie) +# +# return _response_hook +# +# async def _choose_cookie(self, target: Target) -> Cookie: +# """选择 cookie 的具体算法""" +# cookies = await config.get_universal_cookie(self._platform_name) +# if target: +# cookies += await config.get_cookie(self._platform_name, target) +# cookies = (cookie for cookie in cookies if cookie.last_usage + timedelta(seconds=cookie.cd) < datetime.now()) +# cookie = min(cookies, key=lambda x: x.last_usage) +# return cookie +# +# async def _check_cookie(self, cookie: Cookie) -> Cookie: +# """检查Cookie，可以做一些自定义的逻辑，比如说Site的统一风控""" +# return cookie +# +# async def get_client(self, target: Target | None) -> AsyncClient: +# """获取 client，根据 target 选择 cookie""" +# client = http_client() +# cookie = await self._choose_cookie(target) +# if cookie.is_universal: +# logger.debug(f"平台 {self._platform_name} 未获取到用户cookie, 使用匿名cookie") +# else: +# logger.debug(f"平台 {self._platform_name} 获取到用户cookie: {cookie.id}") +# +# return await self._assemble_client(client, cookie) +# +# async def _assemble_client(self, client, cookie) -> AsyncClient: +# """组装 client，可以自定义 cookie 对象的 content 装配到 client 中的方式""" +# cookies = httpx.Cookies() +# if cookie: +# cookies.update(json.loads(cookie.content)) +# client.cookies = cookies +# client._bison_cookie = cookie +# client.event_hooks = {"response": [self._generate_hook(cookie)]} +# return client +# +# async def get_client_for_static(self) -> AsyncClient: +# return http_client() +# +# async def get_query_name_client(self) -> AsyncClient: +# return http_client() +# +# async def refresh_client(self): +# pass +# +# +# def create_cookie_client_manager(platform_name: str) -> type[CookieClientManager]: +# """创建一个平台特化的 CookieClientManger""" +# return type( +# "CookieClientManager", +# (CookieClientManager,), +# {"_platform_name": platform_name}, +# ) - @classmethod - async def init_universal_cookie(cls): - """移除已有的匿名cookie，添加一个新的匿名cookie""" - universal_cookies = await config.get_unviersal_cookie(cls._platform_name) - universal_cookie = Cookie( - platform_name=cls._platform_name, content="{}", is_universal=True, tags={"temporary": True} - ) - for cookie in universal_cookies: - if not cookie.tags.get("temporary"): - continue - await config.delete_cookie_by_id(cookie.id) - universal_cookie.id = cookie.id # 保持原有的id - await config.add_cookie(universal_cookie) - - @classmethod - async def init_cookie(cls, cookie: Cookie) -> Cookie: - """初始化 cookie，添加用户 cookie 时使用""" - cookie.cd = cls._cookie_cd - cookie.last_usage = datetime.now() # 使得优先使用用户 cookie - return cookie - - @classmethod - async def valid_cookie(cls, content: str) -> bool: - """验证 cookie 内容是否有效，添加 cookie 时用，可根据平台的具体情况进行重写""" - try: - data = json.loads(content) - if not isinstance(data, dict): - return False - except JSONDecodeError: - return False - return True - - @classmethod - async def get_cookie_friendly_name(cls, cookie: Cookie) -> str: - """获取 cookie 的友好名字，用于展示""" - from . import text_fletten - - return text_fletten(f"{cookie.platform_name} [{cookie.content[:10]}]") - - def _generate_hook(self, cookie: Cookie) -> callable: - """hook 函数生成器，用于回写请求状态到数据库""" - - async def _response_hook(resp: httpx.Response): - if resp.status_code == 200: - logger.debug(f"请求成功: {cookie.id} {resp.request.url}") - cookie.status = "success" - else: - logger.error(f"请求失败:{cookie.id} {resp.request.url}, 状态码: {resp.status_code}") - cookie.status = "failed" - cookie.last_usage = datetime.now() - await config.update_cookie(cookie) - - return _response_hook - - async def _choose_cookie(self, target: Target) -> Cookie: - """选择 cookie 的具体算法""" - cookies = await config.get_universal_cookie(self._platform_name) - if target: - cookies += await config.get_cookie(self._platform_name, target) - cookies = (cookie for cookie in cookies if cookie.last_usage + timedelta(seconds=cookie.cd) < datetime.now()) - cookie = min(cookies, key=lambda x: x.last_usage) - return cookie - - async def _check_cookie(self, cookie: Cookie) -> Cookie: - """检查Cookie，可以做一些自定义的逻辑，比如说Site的统一风控""" - return cookie - - async def get_client(self, target: Target | None) -> AsyncClient: - """获取 client，根据 target 选择 cookie""" - client = http_client() - cookie = await self._choose_cookie(target) - if cookie.is_universal: - logger.debug(f"平台 {self._platform_name} 未获取到用户cookie, 使用匿名cookie") - else: - logger.debug(f"平台 {self._platform_name} 获取到用户cookie: {cookie.id}") - - return await self._assemble_client(client, cookie) - - async def _assemble_client(self, client, cookie) -> AsyncClient: - """组装 client，可以自定义 cookie 对象的 content 装配到 client 中的方式""" - cookies = httpx.Cookies() - if cookie: - cookies.update(json.loads(cookie.content)) - client.cookies = cookies - client._bison_cookie = cookie - client.event_hooks = {"response": [self._generate_hook(cookie)]} - return client - - async def get_client_for_static(self) -> AsyncClient: - return http_client() - - async def get_query_name_client(self) -> AsyncClient: - return http_client() - - async def refresh_client(self): - pass - - -def create_cookie_client_manager(platform_name: str) -> type[CookieClientManager]: +def create_cookie_client_manager(platform_name: str): """创建一个平台特化的 CookieClientManger""" - return type( - "CookieClientManager", - (CookieClientManager,), - {"_platform_name": platform_name}, - ) + return DefaultClientManager + class Site: