:recycles: 注释掉cookie相关代码,使得bison可以正常运行

This commit is contained in:
suyiiyii 2024-09-08 18:21:57 +08:00
parent ce1f1bbedb
commit 275bc0cb53
3 changed files with 275 additions and 270 deletions

View File

@ -259,130 +259,130 @@ class DBConfig:
) )
return res return res
async def get_cookie(self, platform_name: str = None, target: T_Target = None) -> Sequence[Cookie]: # async def get_cookie(self, platform_name: str = None, target: T_Target = None) -> Sequence[Cookie]:
"""根据平台名和订阅名获取 cookie不会返回匿名cookie""" # """根据平台名和订阅名获取 cookie不会返回匿名cookie"""
async with create_session() as sess: # async with create_session() as sess:
query = select(Cookie).distinct().where(Cookie.is_universal == False) # noqa: E712 # query = select(Cookie).distinct().where(Cookie.is_universal == False) # noqa: E712
if platform_name: # if platform_name:
query = query.where(Cookie.platform_name == platform_name) # query = query.where(Cookie.platform_name == platform_name)
query = query.outerjoin(CookieTarget).options(selectinload(Cookie.targets)) # query = query.outerjoin(CookieTarget).options(selectinload(Cookie.targets))
res = (await sess.scalars(query)).all() # res = (await sess.scalars(query)).all()
if target: # if target:
query = select(CookieTarget.cookie_id).join(Target).where(Target.target == target) # query = select(CookieTarget.cookie_id).join(Target).where(Target.target == target)
ids = set((await sess.scalars(query)).all()) # ids = set((await sess.scalars(query)).all())
res = [cookie for cookie in res if cookie.id in ids] # res = [cookie for cookie in res if cookie.id in ids]
return res # return res
#
async def get_unviersal_cookie(self, platform_name: str = None) -> Sequence[Cookie]: # async def get_unviersal_cookie(self, platform_name: str = None) -> Sequence[Cookie]:
async with create_session() as sess: # async with create_session() as sess:
query = select(Cookie).distinct().where(Cookie.is_universal == True) # noqa: E712 # query = select(Cookie).distinct().where(Cookie.is_universal == True) # noqa: E712
if platform_name: # if platform_name:
query = query.where(Cookie.platform_name == platform_name) # query = query.where(Cookie.platform_name == platform_name)
res = (await sess.scalars(query)).all() # res = (await sess.scalars(query)).all()
return res # return res
#
async def add_cookie_with_content(self, platform_name: str, content: str) -> int: # async def add_cookie_with_content(self, platform_name: str, content: str) -> int:
async with create_session() as sess: # async with create_session() as sess:
cookie = Cookie(platform_name=platform_name, content=content) # cookie = Cookie(platform_name=platform_name, content=content)
sess.add(cookie) # sess.add(cookie)
await sess.commit() # await sess.commit()
await sess.refresh(cookie) # await sess.refresh(cookie)
return cookie.id # return cookie.id
#
async def add_cookie(self, cookie: Cookie) -> int: # async def add_cookie(self, cookie: Cookie) -> int:
async with create_session() as sess: # async with create_session() as sess:
sess.add(cookie) # sess.add(cookie)
await sess.commit() # await sess.commit()
await sess.refresh(cookie) # await sess.refresh(cookie)
return cookie.id # return cookie.id
#
async def update_cookie(self, cookie: Cookie): # async def update_cookie(self, cookie: Cookie):
async with create_session() as sess: # async with create_session() as sess:
cookie_in_db: Cookie | None = await sess.scalar(select(Cookie).where(Cookie.id == cookie.id)) # cookie_in_db: Cookie | None = await sess.scalar(select(Cookie).where(Cookie.id == cookie.id))
if not cookie_in_db: # if not cookie_in_db:
return # return
cookie_in_db.content = cookie.content # cookie_in_db.content = cookie.content
cookie_in_db.last_usage = cookie.last_usage # cookie_in_db.last_usage = cookie.last_usage
cookie_in_db.status = cookie.status # cookie_in_db.status = cookie.status
cookie_in_db.tags = cookie.tags # cookie_in_db.tags = cookie.tags
await sess.commit() # await sess.commit()
#
async def delete_cookie_by_id(self, cookie_id: int): # async def delete_cookie_by_id(self, cookie_id: int):
async with create_session() as sess: # async with create_session() as sess:
cookie = await sess.scalar( # cookie = await sess.scalar(
select(Cookie) # select(Cookie)
.where(Cookie.id == cookie_id) # .where(Cookie.id == cookie_id)
.outerjoin(CookieTarget) # .outerjoin(CookieTarget)
.options(selectinload(Cookie.targets)) # .options(selectinload(Cookie.targets))
) # )
if len(cookie.targets) > 0: # if len(cookie.targets) > 0:
raise Exception(f"cookie {cookie.id} in use") # raise Exception(f"cookie {cookie.id} in use")
await sess.execute(delete(Cookie).where(Cookie.id == cookie_id)) # await sess.execute(delete(Cookie).where(Cookie.id == cookie_id))
await sess.commit() # await sess.commit()
#
async def get_cookie_by_target(self, target: T_Target, platform_name: str) -> Sequence[Cookie]: # async def get_cookie_by_target(self, target: T_Target, platform_name: str) -> Sequence[Cookie]:
async with create_session() as sess: # async with create_session() as sess:
query = ( # query = (
select(Cookie) # select(Cookie)
.join(CookieTarget) # .join(CookieTarget)
.join(Target) # .join(Target)
.where(Target.platform_name == platform_name, Target.target == target) # .where(Target.platform_name == platform_name, Target.target == target)
) # )
return (await sess.scalars(query)).all() # return (await sess.scalars(query)).all()
#
async def get_universal_cookie(self, platform_name: str) -> Sequence[Cookie]: # async def get_universal_cookie(self, platform_name: str) -> Sequence[Cookie]:
async with create_session() as sess: # async with create_session() as sess:
query = ( # query = (
select(Cookie) # select(Cookie)
.where(Cookie.platform_name == platform_name) # .where(Cookie.platform_name == platform_name)
.where(Cookie.is_universal == True) # noqa: E712 # .where(Cookie.is_universal == True) # noqa: E712
) # )
return (await sess.scalars(query)).all() # return (await sess.scalars(query)).all()
#
async def add_cookie_target(self, target: T_Target, platform_name: str, cookie_id: int): # async def add_cookie_target(self, target: T_Target, platform_name: str, cookie_id: int):
async with create_session() as sess: # async with create_session() as sess:
target_obj = await sess.scalar( # target_obj = await sess.scalar(
select(Target).where(Target.platform_name == platform_name, Target.target == target) # select(Target).where(Target.platform_name == platform_name, Target.target == target)
) # )
# check if relation exists # # check if relation exists
cookie_target = await sess.scalar( # cookie_target = await sess.scalar(
select(CookieTarget).where(CookieTarget.target == target_obj, CookieTarget.cookie_id == cookie_id) # select(CookieTarget).where(CookieTarget.target == target_obj, CookieTarget.cookie_id == cookie_id)
) # )
if cookie_target: # if cookie_target:
raise DuplicateCookieTargetException() # raise DuplicateCookieTargetException()
cookie_obj = await sess.scalar(select(Cookie).where(Cookie.id == cookie_id)) # cookie_obj = await sess.scalar(select(Cookie).where(Cookie.id == cookie_id))
cookie_target = CookieTarget(target=target_obj, cookie=cookie_obj) # cookie_target = CookieTarget(target=target_obj, cookie=cookie_obj)
sess.add(cookie_target) # sess.add(cookie_target)
await sess.commit() # await sess.commit()
#
async def delete_cookie_target(self, target: T_Target, platform_name: str, cookie_id: int): # async def delete_cookie_target(self, target: T_Target, platform_name: str, cookie_id: int):
async with create_session() as sess: # async with create_session() as sess:
target_obj = await sess.scalar( # target_obj = await sess.scalar(
select(Target).where(Target.platform_name == platform_name, Target.target == target) # select(Target).where(Target.platform_name == platform_name, Target.target == target)
) # )
cookie_obj = await sess.scalar(select(Cookie).where(Cookie.id == cookie_id)) # cookie_obj = await sess.scalar(select(Cookie).where(Cookie.id == cookie_id))
await sess.execute( # await sess.execute(
delete(CookieTarget).where(CookieTarget.target == target_obj, CookieTarget.cookie == cookie_obj) # delete(CookieTarget).where(CookieTarget.target == target_obj, CookieTarget.cookie == cookie_obj)
) # )
await sess.commit() # await sess.commit()
#
async def delete_cookie_target_by_id(self, cookie_target_id: int): # async def delete_cookie_target_by_id(self, cookie_target_id: int):
async with create_session() as sess: # async with create_session() as sess:
await sess.execute(delete(CookieTarget).where(CookieTarget.id == cookie_target_id)) # await sess.execute(delete(CookieTarget).where(CookieTarget.id == cookie_target_id))
await sess.commit() # await sess.commit()
#
async def get_cookie_target(self) -> list[CookieTarget]: # async def get_cookie_target(self) -> list[CookieTarget]:
async with create_session() as sess: # async with create_session() as sess:
query = ( # query = (
select(CookieTarget) # select(CookieTarget)
.outerjoin(Target) # .outerjoin(Target)
.options(selectinload(CookieTarget.target)) # .options(selectinload(CookieTarget.target))
.outerjoin(Cookie) # .outerjoin(Cookie)
.options(selectinload(CookieTarget.cookie)) # .options(selectinload(CookieTarget.cookie))
) # )
res = list((await sess.scalars(query)).all()) # res = list((await sess.scalars(query)).all())
res.sort(key=lambda x: (x.target.platform_name, x.cookie_id, x.target_id)) # res.sort(key=lambda x: (x.target.platform_name, x.cookie_id, x.target_id))
return res # return res
config = DBConfig() config = DBConfig()

View File

@ -14,10 +14,10 @@ from nonebot.adapters.onebot.v11.event import PrivateMessageEvent
from .add_sub import do_add_sub from .add_sub import do_add_sub
from .del_sub import do_del_sub from .del_sub import do_del_sub
from .query_sub import do_query_sub from .query_sub import do_query_sub
from .add_cookie import do_add_cookie # from .add_cookie import do_add_cookie
from .del_cookie import do_del_cookie # from .del_cookie import do_del_cookie
from .add_cookie_target import do_add_cookie_target # from .add_cookie_target import do_add_cookie_target
from .del_cookie_target import do_del_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 from .utils import common_platform, admin_permission, gen_handle_cancel, configurable_to_me, set_target_user_info
add_sub_matcher = on_command( add_sub_matcher = on_command(
@ -44,41 +44,41 @@ del_sub_matcher = on_command(
del_sub_matcher.handle()(set_target_user_info) del_sub_matcher.handle()(set_target_user_info)
do_del_sub(del_sub_matcher) do_del_sub(del_sub_matcher)
add_cookie_matcher = on_command( # add_cookie_matcher = on_command(
"添加cookie", # "添加cookie",
rule=configurable_to_me, # rule=configurable_to_me,
permission=SUPERUSER, # permission=SUPERUSER,
priority=5, # priority=5,
block=True, # block=True,
) # )
do_add_cookie(add_cookie_matcher) # do_add_cookie(add_cookie_matcher)
#
add_cookie_target_matcher = on_command( # add_cookie_target_matcher = on_command(
"关联cookie", # "关联cookie",
rule=configurable_to_me, # rule=configurable_to_me,
permission=SUPERUSER, # permission=SUPERUSER,
priority=5, # priority=5,
block=True, # block=True,
) # )
do_add_cookie_target(add_cookie_target_matcher) # do_add_cookie_target(add_cookie_target_matcher)
#
del_cookie_target_matcher = on_command( # del_cookie_target_matcher = on_command(
"取消关联cookie", # "取消关联cookie",
rule=configurable_to_me, # rule=configurable_to_me,
permission=SUPERUSER, # permission=SUPERUSER,
priority=5, # priority=5,
block=True, # block=True,
) # )
do_del_cookie_target(del_cookie_target_matcher) # do_del_cookie_target(del_cookie_target_matcher)
#
del_cookie_matcher = on_command( # del_cookie_matcher = on_command(
"删除cookie", # "删除cookie",
rule=configurable_to_me, # rule=configurable_to_me,
permission=SUPERUSER, # permission=SUPERUSER,
priority=5, # priority=5,
block=True, # block=True,
) # )
do_del_cookie(del_cookie_matcher) # do_del_cookie(del_cookie_matcher)
group_manage_matcher = on_command("群管理", rule=to_me(), permission=SUPERUSER, priority=4, block=True) group_manage_matcher = on_command("群管理", rule=to_me(), permission=SUPERUSER, priority=4, block=True)

View File

@ -45,117 +45,122 @@ class DefaultClientManager(ClientManager):
def is_cookie_client_manager(manger: type[ClientManager]) -> bool: def is_cookie_client_manager(manger: type[ClientManager]) -> bool:
return hasattr(manger, "_cookie_client_manger_") return hasattr(manger, "_cookie_client_manger_")
# TODO: Implement CookieClientManager
class CookieClientManager(ClientManager): class CookieClientManager(ClientManager):
_cookie_client_manger_ = True # _cookie_client_manger_ = True
_platform_name: str _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 def create_cookie_client_manager(platform_name: str):
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""" """创建一个平台特化的 CookieClientManger"""
return type( return DefaultClientManager
"CookieClientManager",
(CookieClientManager,),
{"_platform_name": platform_name},
)
class Site: class Site: