This commit is contained in:
suyiiyii 2024-09-13 10:11:54 +08:00
parent b25fcd9ac2
commit 4b8d6a9379
5 changed files with 23 additions and 27 deletions

View File

@ -1,20 +1,20 @@
import asyncio
from collections import defaultdict
from collections.abc import Callable, Sequence, Awaitable
from datetime import time, datetime
from collections.abc import Callable, Sequence, Awaitable
from nonebot.compat import model_dump
from nonebot_plugin_datastore import create_session
from nonebot_plugin_saa import PlatformTarget
from sqlalchemy import func, delete, select
from sqlalchemy.exc import IntegrityError
from sqlalchemy.orm import selectinload
from sqlalchemy.exc import IntegrityError
from sqlalchemy import func, delete, select
from nonebot_plugin_saa import PlatformTarget
from nonebot_plugin_datastore import create_session
from .db_model import User, Cookie, Target, Subscribe, CookieTarget, ScheduleTimeWeight
from .utils import NoSuchTargetException, DuplicateCookieTargetException
from ..types import Category, UserSubInfo, WeightConfig, TimeWeightConfig, PlatformWeightConfigResp
from ..types import Tag
from ..types import Target as T_Target
from .utils import NoSuchTargetException, DuplicateCookieTargetException
from .db_model import User, Cookie, Target, Subscribe, CookieTarget, ScheduleTimeWeight
from ..types import Category, UserSubInfo, WeightConfig, TimeWeightConfig, PlatformWeightConfigResp
def _get_time():
@ -332,18 +332,14 @@ class DBConfig:
async def get_universal_cookie(self, site_name: str) -> Sequence[Cookie]:
async with create_session() as sess:
query = (
select(Cookie)
.where(Cookie.site_name == site_name)
.where(Cookie.is_universal == True) # noqa: E712
)
query = select(Cookie).where(Cookie.site_name == site_name).where(Cookie.is_universal == True) # noqa: E712
return (await sess.scalars(query)).all()
async def add_cookie_target(self, target: T_Target, site_name: str, cookie_id: int):
async def add_cookie_target(self, target: T_Target, platform_name: str, cookie_id: int):
"""通过 cookie_id 可以唯一确定一个 Cookie通过 target 和 platform_name 可以唯一确定一个 Target"""
async with create_session() as sess:
target_obj = await sess.scalar(
select(Target).where(Target.target == target)
# TODO: 仅判断 target可能会有重名现象还要判断 platform_name
select(Target).where(Target.platform_name == platform_name, Target.target == target)
)
# check if relation exists
cookie_target = await sess.scalar(
@ -358,9 +354,7 @@ class DBConfig:
async def delete_cookie_target(self, target: T_Target, site_name: str, cookie_id: int):
async with create_session() as sess:
target_obj = await sess.scalar(
select(Target).where(Target.site_name == site_name, Target.target == target)
)
target_obj = await sess.scalar(select(Target).where(Target.site_name == site_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)

View File

@ -46,6 +46,7 @@ do_del_sub(del_sub_matcher)
add_cookie_matcher = on_command(
"添加cookie",
aliases={"添加Cookie"},
rule=configurable_to_me,
permission=SUPERUSER,
priority=5,
@ -55,6 +56,7 @@ do_add_cookie(add_cookie_matcher)
add_cookie_target_matcher = on_command(
"关联cookie",
aliases={"关联Cookie"},
rule=configurable_to_me,
permission=SUPERUSER,
priority=5,
@ -64,6 +66,7 @@ do_add_cookie_target(add_cookie_target_matcher)
del_cookie_target_matcher = on_command(
"取消关联cookie",
aliases={"取消关联Cookie"},
rule=configurable_to_me,
permission=SUPERUSER,
priority=5,
@ -73,6 +76,7 @@ do_del_cookie_target(del_cookie_target_matcher)
del_cookie_matcher = on_command(
"删除cookie",
aliases={"删除Cookie"},
rule=configurable_to_me,
permission=SUPERUSER,
priority=5,
@ -80,7 +84,6 @@ del_cookie_matcher = on_command(
)
do_del_cookie(del_cookie_matcher)
group_manage_matcher = on_command("群管理", rule=to_me(), permission=SUPERUSER, priority=4, block=True)
group_handle_cancel = gen_handle_cancel(group_manage_matcher, "已取消")

View File

@ -8,7 +8,7 @@ from nonebot.internal.adapter import MessageTemplate
from ..config import config
from ..utils import parse_text
from ..platform import platform_manager, site_manager
from ..platform import platform_manager
from ..utils.site import CookieClientManager
from .utils import gen_handle_cancel, generate_sub_list_text
@ -64,7 +64,7 @@ def do_add_cookie_target(add_cookie_target_matcher: type[Matcher]):
@add_cookie_target_matcher.handle()
async def add_cookie_target_process(state: T_State):
await config.add_cookie_target(state["target"]["target"], state["site"].name, state["cookie"].id)
await config.add_cookie_target(state["target"]["target"], state["target"]["platform_name"], state["cookie"].id)
cookie = state["cookie"]
client_mgr = cast(CookieClientManager, state["site"].client_mgr)
await add_cookie_target_matcher.finish(

View File

@ -13,8 +13,8 @@ from nonebot_plugin_saa import PlatformTarget, extract_target
from ..config import config
from ..types import Category
from ..platform import platform_manager, site_manager
from ..plugin_config import plugin_config
from ..platform import site_manager, platform_manager
from ..utils.site import CookieClientManager, is_cookie_client_manager
@ -75,6 +75,7 @@ async def generate_sub_list_text(
is_show_cookie=False,
is_hide_no_cookie_platfrom=False,
):
"""根据配置参数生产订阅列表文本同时将订阅信息存入state["sub_table"]"""
if user_info:
sub_list = await config.list_subscribe(user_info)
else:
@ -107,9 +108,7 @@ async def generate_sub_list_text(
if sub.tags:
res += " {}".format(", ".join(sub.tags)) + "\n"
if is_show_cookie:
target_cookies = await config.get_cookie(
target=sub.target.target, site_name=platform.site.name
)
target_cookies = await config.get_cookie(target=sub.target.target, site_name=platform.site.name)
if target_cookies:
res += " 关联的 Cookie\n"
for cookie in target_cookies:

View File

@ -64,7 +64,7 @@ class CookieClientManager(ClientManager):
"""添加用户 cookie"""
cookie = Cookie(site_name=cls._site_name, content=content)
cookie.cd = cls._default_cd
config.add_cookie(cookie)
await config.add_cookie(cookie)
@classmethod
async def validate_cookie(cls, content: str) -> bool: