From 370fc250f03de5d7375eca338c2c99ba3973cfcb Mon Sep 17 00:00:00 2001
From: suyiiyii <suyiiyii@gmail.com>
Date: Thu, 5 Sep 2024 19:32:47 +0800
Subject: [PATCH] =?UTF-8?q?:sparkles:=20=E6=95=B0=E6=8D=AE=E5=BA=93Cookie?=
 =?UTF-8?q?=E8=A1=A8=E6=B7=BB=E5=8A=A0is=5Funiversal=E5=B1=9E=E6=80=A7?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 nonebot_bison/config/db_model.py              |  1 +
 .../migrations/fc2b8481bdde_add_cookie.py     | 61 +++++++++++++++++++
 nonebot_bison/utils/site.py                   |  9 ++-
 3 files changed, 69 insertions(+), 2 deletions(-)
 create mode 100644 nonebot_bison/config/migrations/fc2b8481bdde_add_cookie.py

diff --git a/nonebot_bison/config/db_model.py b/nonebot_bison/config/db_model.py
index d7c84f9..4362ee4 100644
--- a/nonebot_bison/config/db_model.py
+++ b/nonebot_bison/config/db_model.py
@@ -77,6 +77,7 @@ class Cookie(Model):
     last_usage: Mapped[datetime.datetime] = mapped_column(DateTime, default=datetime.datetime(1970, 1, 1))
     status: Mapped[str] = mapped_column(String(20), default="")
     cd: Mapped[int] = mapped_column(default=0)
+    is_universal: Mapped[bool] = mapped_column(default=False)
     tags: Mapped[dict[str, Any]] = mapped_column(JSON().with_variant(JSONB, "postgresql"), default={})
 
     targets: Mapped[list["CookieTarget"]] = relationship(back_populates="cookie")
diff --git a/nonebot_bison/config/migrations/fc2b8481bdde_add_cookie.py b/nonebot_bison/config/migrations/fc2b8481bdde_add_cookie.py
new file mode 100644
index 0000000..b296ecf
--- /dev/null
+++ b/nonebot_bison/config/migrations/fc2b8481bdde_add_cookie.py
@@ -0,0 +1,61 @@
+"""empty message
+
+Revision ID: fc2b8481bdde
+Revises: f9baef347cc8
+Create Date: 2024-09-05 19:31:59.366754
+
+"""
+
+import sqlalchemy as sa
+from alembic import op
+from sqlalchemy import Text
+from sqlalchemy.dialects import postgresql
+
+# revision identifiers, used by Alembic.
+revision = "fc2b8481bdde"
+down_revision = "f9baef347cc8"
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.create_table(
+        "nonebot_bison_cookie",
+        sa.Column("id", sa.Integer(), nullable=False),
+        sa.Column("platform_name", sa.String(length=20), nullable=False),
+        sa.Column("content", sa.String(length=1024), nullable=False),
+        sa.Column("last_usage", sa.DateTime(), nullable=False),
+        sa.Column("status", sa.String(length=20), nullable=False),
+        sa.Column("cd", sa.Integer(), nullable=False),
+        sa.Column("is_universal", sa.Boolean(), nullable=False),
+        sa.Column("tags", sa.JSON().with_variant(postgresql.JSONB(astext_type=Text()), "postgresql"), nullable=False),
+        sa.PrimaryKeyConstraint("id", name=op.f("pk_nonebot_bison_cookie")),
+    )
+    op.create_table(
+        "nonebot_bison_cookietarget",
+        sa.Column("id", sa.Integer(), nullable=False),
+        sa.Column("target_id", sa.Integer(), nullable=False),
+        sa.Column("cookie_id", sa.Integer(), nullable=False),
+        sa.ForeignKeyConstraint(
+            ["cookie_id"],
+            ["nonebot_bison_cookie.id"],
+            name=op.f("fk_nonebot_bison_cookietarget_cookie_id_nonebot_bison_cookie"),
+            ondelete="CASCADE",
+        ),
+        sa.ForeignKeyConstraint(
+            ["target_id"],
+            ["nonebot_bison_target.id"],
+            name=op.f("fk_nonebot_bison_cookietarget_target_id_nonebot_bison_target"),
+            ondelete="CASCADE",
+        ),
+        sa.PrimaryKeyConstraint("id", name=op.f("pk_nonebot_bison_cookietarget")),
+    )
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.drop_table("nonebot_bison_cookietarget")
+    op.drop_table("nonebot_bison_cookie")
+    # ### end Alembic commands ###
diff --git a/nonebot_bison/utils/site.py b/nonebot_bison/utils/site.py
index 913b815..ac3ec94 100644
--- a/nonebot_bison/utils/site.py
+++ b/nonebot_bison/utils/site.py
@@ -71,11 +71,16 @@ class CookieClientManager(ClientManager):
         cookies = (
             cookie for cookie in cookies if cookie.last_usage + timedelta(seconds=self._cookie_cd) < datetime.now()
         )
+        cookies = list(cookies)
         if not cookies:
             return Cookie(content="{}")
         cookie = max(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()
@@ -85,9 +90,9 @@ class CookieClientManager(ClientManager):
         else:
             logger.debug(f"平台 {self._platform_name} 未获取到用户cookie, 使用空cookie")
 
-        return await self.assemble_client(client, cookie)
+        return await self._assemble_client(client, cookie)
 
-    async def assemble_client(self, client, cookie):
+    async def _assemble_client(self, client, cookie):
         cookies = httpx.Cookies()
         if cookie:
             cookies.update(json.loads(cookie.content))