feat(cookies): Cookie 失效后可账号重新登录替换, 并修复 Cookie 列空白过多

2026-08-08 01:39:02 +08:00
parent 64b017fd6a
commit 39ea0d2184
5 changed files with 269 additions and 60 deletions
@@ -57,6 +57,7 @@ class LoginBatchRunner:
        concurrency: int = 3,
        api_strategy: str = "wgapi",
        mode: str = "login",
+        relogin_task_ids: Optional[list[int]] = None,
    ):
        self.db = db
        self.account_ids = account_ids
@@ -72,7 +73,8 @@ class LoginBatchRunner:
        self.batch_id = uuid.uuid4().hex[:12]
        self.concurrency = max(1, min(concurrency, 10))  # 限制 1-10
        self.api_strategy = _create_api_strategy(api_strategy)
-        self.mode = "check" if mode == "check" else "login"
+        self.mode = mode if mode in ("login", "check", "relogin") else "login"
+        self.relogin_task_ids = list(relogin_task_ids) if relogin_task_ids else []
        self._stop = threading.Event()
        self._counter_lock = threading.Lock()
        self._completed = 0
@@ -133,6 +135,17 @@ class LoginBatchRunner:
        """在独立线程中执行单个账号登录，使用独立的 DB 会话。"""
        if self._stop.is_set():
            self._push_log("warning", f"任务已停止，跳过: {acc_info['username']}")
+            if self.mode == "relogin":
+                # 重新登录被停止时恢复原 Cookie 记录，避免列表行消失
+                worker_db = SessionLocal()
+                try:
+                    task = worker_db.query(LoginTask).filter(LoginTask.id == task_id).first()
+                    if task and task.status in ("pending", "running"):
+                        task.status = "success"
+                        task.finished_at = datetime.now(timezone.utc)
+                        worker_db.commit()
+                finally:
+                    worker_db.close()
            return

        worker_db = SessionLocal()
@@ -148,7 +161,7 @@ class LoginBatchRunner:
                self._completed += 1
                current = self._completed

-            action_name = "检测" if self.mode == "check" else "登录"
+            action_name = "检测" if self.mode == "check" else "重新登录" if self.mode == "relogin" else "登录"
            self._push_log("info", f"[{current}/{total}] 开始{action_name}: {acc_info['username']}")

            # 解析代理配置
@@ -198,16 +211,34 @@ class LoginBatchRunner:
                    task.status = "success"
                    task.cookie = result.cookie
                    task.message = result.message or "登录成功"
-                    self._push_log("success", f"[{current}] {acc_info['username']} {task.message}")
+                    if self.mode == "relogin":
+                        task.ck_check_status = ""
+                        task.ck_check_result = None
+                        task.ck_checked_at = None
+                        task.message = "重新登录成功"
+                        self._push_log("success", f"[{current}] {acc_info['username']} 重新登录成功，Cookie 已替换")
+                    else:
+                        self._push_log("success", f"[{current}] {acc_info['username']} {task.message}")
                else:
-                    task.status = "failed"
-                    task.message = result.message
-                    self._push_log("error", f"[{current}] {acc_info['username']} {action_name}失败: {result.message}")
+                    if self.mode == "relogin":
+                        # 重新登录失败时保留旧 Cookie 与成功状态，仅记录失败原因，行不消失
+                        task.status = "success"
+                        task.message = f"重新登录失败: {result.message}"
+                        self._push_log("error", f"[{current}] {acc_info['username']} 重新登录失败: {result.message}")
+                    else:
+                        task.status = "failed"
+                        task.message = result.message
+                        self._push_log("error", f"[{current}] {acc_info['username']} {action_name}失败: {result.message}")

            except Exception as e:
-                task.status = "error"
-                task.message = str(e)
-                self._push_log("error", f"[{current}] {acc_info['username']} {action_name}异常: {e}")
+                if self.mode == "relogin":
+                    task.status = "success"
+                    task.message = f"重新登录异常: {e}"
+                    self._push_log("error", f"[{current}] {acc_info['username']} 重新登录异常: {e}")
+                else:
+                    task.status = "error"
+                    task.message = str(e)
+                    self._push_log("error", f"[{current}] {acc_info['username']} {action_name}异常: {e}")

            task.finished_at = datetime.now(timezone.utc)
            worker_db.commit()
@@ -219,8 +250,8 @@ class LoginBatchRunner:
        """在线程中执行批量登录。"""
        batch_id = self.batch_id
        concurrency = self.concurrency
-        action_name = "账号检测" if self.mode == "check" else "登录"
-        self._push_log("info", f"批量{action_name}任务 {batch_id} 开始，共 {len(self.account_ids)} 个账号，并发数: {concurrency}")
+        action_name = "账号检测" if self.mode == "check" else "重新登录" if self.mode == "relogin" else "登录"
+        self._push_log("info", f"批量{action_name}任务 {batch_id} 开始，共 {len(self.account_ids or self.relogin_task_ids)} 个账号，并发数: {concurrency}")

        # 批次开始前同步一次出口 IP 到白名单，后续 fetch_new_proxy 不再主动同步
        if self._shared_proxy_fetcher:
@@ -228,56 +259,77 @@ class LoginBatchRunner:
            if msg != "无白名单凭据，跳过":
                self._push_log("info" if ok else "warning", f"白名单预热: {msg}")

+        def _append_task_info(task: LoginTask, acc: AccountModel):
+            task.batch_id = batch_id
+            task.status = "pending"
+            task.message = ""
+            task.finished_at = None
+            self.db.flush()
+            task_infos.append({
+                "task_id": task.id,
+                "acc_info": {
+                    "username": acc.username,
+                    "password": acc.password,
+                    "email": acc.email,
+                    "email_password": acc.email_password,
+                    "email_imap_server": acc.email_imap_server or "",
+                    "email_imap_port": acc.email_imap_port or 993,
+                    "email_imap_ssl": acc.email_imap_ssl if acc.email_imap_ssl is not None else True,
+                },
+            })
+
        try:
            # 创建或复用任务记录（顺序执行，线程安全）
            task_infos: list[dict] = []  # {task_id, acc_info}
-            for aid in self.account_ids:
-                acc = self.db.query(AccountModel).filter(AccountModel.id == aid).first()
-                if not acc:
-                    continue
-                # 权限检查：客服只能跑分配给自己的
-                if "login:view_all" not in self.creator_permissions:
-                    if acc.assigned_to != self.created_by:
-                        self._push_log("warning", f"跳过无权账号: {acc.username}")
+            if self.relogin_task_ids:
+                # 重新登录模式：复用指定 Cookie 记录，登录成功后原地替换 Cookie
+                for task_id in self.relogin_task_ids:
+                    task = self.db.query(LoginTask).filter(LoginTask.id == task_id).first()
+                    if not task:
                        continue
+                    acc = self.db.query(AccountModel).filter(AccountModel.id == task.account_id).first()
+                    if not acc:
+                        self._push_log("warning", f"跳过无账号的任务 #{task_id}")
+                        continue
+                    if "login:view_all" not in self.creator_permissions:
+                        if acc.assigned_to != self.created_by:
+                            self._push_log("warning", f"跳过无权账号: {acc.username}")
+                            continue
+                    task.ck_check_status = ""
+                    task.ck_check_result = None
+                    task.ck_checked_at = None
+                    _append_task_info(task, acc)
+            else:
+                for aid in self.account_ids:
+                    acc = self.db.query(AccountModel).filter(AccountModel.id == aid).first()
+                    if not acc:
+                        continue
+                    # 权限检查：客服只能跑分配给自己的
+                    if "login:view_all" not in self.creator_permissions:
+                        if acc.assigned_to != self.created_by:
+                            self._push_log("warning", f"跳过无权账号: {acc.username}")
+                            continue

-                # 复用该账号最近一条失败任务记录，避免重复产生多条
-                existing_task = (
-                    self.db.query(LoginTask)
-                    .filter(LoginTask.account_id == aid, LoginTask.status.in_(["failed", "error"]))
-                    .order_by(LoginTask.id.desc())
-                    .first()
-                )
-                if existing_task:
-                    existing_task.batch_id = batch_id
-                    existing_task.status = "pending"
-                    existing_task.cookie = ""
-                    existing_task.message = ""
-                    existing_task.finished_at = None
-                    task = existing_task
-                else:
-                    task = LoginTask(
-                        batch_id=batch_id,
-                        account_id=aid,
-                        status="pending",
-                        created_by=self.created_by,
+                    # 复用该账号最近一条失败任务记录，避免重复产生多条
+                    existing_task = (
+                        self.db.query(LoginTask)
+                        .filter(LoginTask.account_id == aid, LoginTask.status.in_(["failed", "error"]))
+                        .order_by(LoginTask.id.desc())
+                        .first()
                    )
-                    self.db.add(task)
+                    if existing_task:
+                        existing_task.cookie = ""
+                        task = existing_task
+                    else:
+                        task = LoginTask(
+                            batch_id=batch_id,
+                            account_id=aid,
+                            status="pending",
+                            created_by=self.created_by,
+                        )
+                        self.db.add(task)

-                self.db.flush()  # 获取 task.id
-
-                task_infos.append({
-                    "task_id": task.id,
-                    "acc_info": {
-                        "username": acc.username,
-                        "password": acc.password,
-                        "email": acc.email,
-                        "email_password": acc.email_password,
-                        "email_imap_server": acc.email_imap_server or "",
-                        "email_imap_port": acc.email_imap_port or 993,
-                        "email_imap_ssl": acc.email_imap_ssl if acc.email_imap_ssl is not None else True,
-                    },
-                })
+                    _append_task_info(task, acc)

            self.db.commit()
            total = len(task_infos)