342 lines
14 KiB
Python
342 lines
14 KiB
Python
"""登录服务：复用 core/ 核心模块，在线程池中并发执行登录并推送日志。"""
|
|
|
|
import asyncio
|
|
import threading
|
|
import time
|
|
import uuid
|
|
from types import SimpleNamespace
|
|
from concurrent.futures import ThreadPoolExecutor, as_completed
|
|
from datetime import datetime, timezone
|
|
from typing import Optional
|
|
|
|
from sqlalchemy.orm import Session
|
|
|
|
from core.douyu import DouyuLogin
|
|
from core.douyu.proxy import resolve_working_proxy, get_proxy_manager
|
|
from ..models import Account as AccountModel, LoginTask, ProxyConfig as ProxyConfigModel
|
|
|
|
|
|
class LoginBatchRunner:
|
|
"""批量登录执行器，在线程中运行，通过 ThreadPoolExecutor 并发登录多个账号。"""
|
|
|
|
# 代理池耗尽时等待恢复的最大秒数
|
|
_PROXY_WAIT_MAX = 60
|
|
|
|
def __init__(
|
|
self,
|
|
db: Session,
|
|
account_ids: list[int],
|
|
created_by: int,
|
|
creator_permissions: list[str],
|
|
max_geetest_retries: int = 5,
|
|
max_proxy_retries: int = 0,
|
|
max_login_retries: int = 3,
|
|
max_total_time: float = 300,
|
|
proxy_config: Optional[ProxyConfigModel] = None,
|
|
log_queue: Optional[asyncio.Queue] = None,
|
|
loop: Optional[asyncio.AbstractEventLoop] = None,
|
|
concurrency: int = 3,
|
|
):
|
|
self.db = db
|
|
self.account_ids = account_ids
|
|
self.created_by = created_by
|
|
self.creator_permissions = creator_permissions
|
|
self.max_geetest_retries = max_geetest_retries
|
|
self.max_proxy_retries = max_proxy_retries
|
|
self.max_login_retries = max_login_retries
|
|
self.max_total_time = max_total_time
|
|
self.proxy_config = proxy_config
|
|
self.log_queue = log_queue
|
|
self.loop = loop
|
|
self.batch_id = uuid.uuid4().hex[:12]
|
|
self.concurrency = max(1, min(concurrency, 10)) # 限制 1-10
|
|
self._stop = threading.Event()
|
|
self._counter_lock = threading.Lock()
|
|
self._completed = 0
|
|
|
|
# 共享代理管理器（带锁，避免并发白名单限流；极验失败时可刷新代理）
|
|
self._shared_proxy_manager = None
|
|
if proxy_config and proxy_config.enabled and proxy_config.api_url:
|
|
wl_uid = proxy_config.whitelist_uid or "" if proxy_config.whitelist_enabled else ""
|
|
wl_ukey = proxy_config.whitelist_ukey or "" if proxy_config.whitelist_enabled else ""
|
|
self._shared_proxy_manager = get_proxy_manager(
|
|
proxy_config.api_url,
|
|
whitelist_uid=wl_uid,
|
|
whitelist_ukey=wl_ukey,
|
|
)
|
|
|
|
def stop(self):
|
|
self._stop.set()
|
|
|
|
def _push_log(self, level: str, message: str):
|
|
if self.log_queue and self.loop:
|
|
asyncio.run_coroutine_threadsafe(
|
|
self.log_queue.put({"level": level, "message": message}),
|
|
self.loop,
|
|
)
|
|
|
|
def _resolve_static_proxy(self) -> tuple[Optional[dict], str]:
|
|
"""
|
|
解析静态代理配置（仅处理无代理和静态代理场景）。
|
|
|
|
API代理由 DouyuLogin 通过 proxy_manager 内部管理，
|
|
不在此处预先获取——登录过程中的代理切换（极验失败、整体重试）
|
|
都在 DouyuLogin 内部自治完成。
|
|
"""
|
|
if not self.proxy_config or not self.proxy_config.enabled:
|
|
return None, ''
|
|
|
|
# 静态代理
|
|
if self.proxy_config.http or self.proxy_config.https:
|
|
proxy_url = self.proxy_config.http or self.proxy_config.https
|
|
return {'http': proxy_url, 'https': proxy_url}, f'使用静态代理: {proxy_url}'
|
|
|
|
# API代理：不在此处获取，由 DouyuLogin 通过 proxy_manager 内部管理
|
|
return None, ''
|
|
|
|
def _execute_one(self, task_id: int, acc_info: dict, total: int):
|
|
"""在独立线程中执行单个账号登录，使用独立的 DB 会话和代理。"""
|
|
if self._stop.is_set():
|
|
self._push_log("warning", f"任务已停止，跳过: {acc_info['username']}")
|
|
return
|
|
|
|
worker_db = SessionLocal()
|
|
try:
|
|
task = worker_db.query(LoginTask).filter(LoginTask.id == task_id).first()
|
|
if not task:
|
|
return
|
|
|
|
task.status = "running"
|
|
worker_db.commit()
|
|
|
|
with self._counter_lock:
|
|
self._completed += 1
|
|
current = self._completed
|
|
|
|
self._push_log("info", f"[{current}/{total}] 开始登录: {acc_info['username']}")
|
|
|
|
# 解析代理配置
|
|
proxy_dict, proxy_msg = self._resolve_static_proxy()
|
|
if proxy_msg:
|
|
self._push_log("info", f"[{current}] {proxy_msg}")
|
|
|
|
# API代理模式下，验证代理池是否有可用代理；池空时等待恢复（最多 _PROXY_WAIT_MAX 秒）
|
|
# 注意：不预取代理传给 DouyuLogin，让 DouyuLogin 通过 proxy_manager 内部自治管理
|
|
is_api_proxy = (
|
|
self.proxy_config
|
|
and self.proxy_config.enabled
|
|
and not (self.proxy_config.http or self.proxy_config.https)
|
|
and self._shared_proxy_manager
|
|
)
|
|
if is_api_proxy:
|
|
# 先验证代理池是否有可用代理（获取后立即归还，不占用）
|
|
test_proxy = self._shared_proxy_manager.get_proxy()
|
|
if test_proxy:
|
|
self._shared_proxy_manager.release_proxy(test_proxy)
|
|
self._push_log("info", f"[{current}] 代理池可用，由 DouyuLogin 内部管理代理获取与切换")
|
|
else:
|
|
# 代理池暂时耗尽，等待冷却代理恢复或新代理入池
|
|
self._push_log("warning", f"[{current}] 代理池暂时耗尽，等待恢复...")
|
|
pool_available = False
|
|
for wait_sec in range(0, self._PROXY_WAIT_MAX, 10):
|
|
if self._stop.is_set():
|
|
task.status = "error"
|
|
task.message = "任务已停止"
|
|
task.finished_at = datetime.now(timezone.utc)
|
|
worker_db.commit()
|
|
return
|
|
time.sleep(10)
|
|
test_proxy = self._shared_proxy_manager.get_proxy()
|
|
if test_proxy:
|
|
self._shared_proxy_manager.release_proxy(test_proxy)
|
|
pool_available = True
|
|
self._push_log("info", f"[{current}] 代理池恢复，由 DouyuLogin 内部管理代理")
|
|
break
|
|
remaining = self._PROXY_WAIT_MAX - wait_sec - 10
|
|
self._push_log("info", f"[{current}] 代理池仍为空，继续等待... (剩余 {remaining}s)")
|
|
|
|
if not pool_available:
|
|
task.status = "error"
|
|
task.message = f"代理池耗尽，等待 {self._PROXY_WAIT_MAX}s 后仍无可用代理"
|
|
task.finished_at = datetime.now(timezone.utc)
|
|
worker_db.commit()
|
|
self._push_log("error", f"[{current}] {acc_info['username']} 代理池耗尽，放弃")
|
|
return
|
|
|
|
# 静态代理启用但配置为空 → 不可用
|
|
if self.proxy_config and self.proxy_config.enabled and not (self.proxy_config.http or self.proxy_config.https) and not self._shared_proxy_manager and not proxy_dict:
|
|
task.status = "error"
|
|
task.message = "代理不可用: 未配置代理"
|
|
task.finished_at = datetime.now(timezone.utc)
|
|
worker_db.commit()
|
|
self._push_log("error", f"[{current}] {acc_info['username']} 代理不可用")
|
|
return
|
|
|
|
try:
|
|
account = SimpleNamespace(
|
|
username=acc_info["username"],
|
|
password=acc_info["password"],
|
|
email=acc_info["email"],
|
|
email_password=acc_info["email_password"],
|
|
email_imap_server=acc_info["email_imap_server"] or "",
|
|
email_imap_port=acc_info["email_imap_port"] or 993,
|
|
email_imap_ssl=acc_info["email_imap_ssl"],
|
|
)
|
|
|
|
loginer = DouyuLogin(
|
|
account,
|
|
proxy=proxy_dict,
|
|
max_geetest_retries=self.max_geetest_retries,
|
|
max_proxy_retries=self.max_proxy_retries,
|
|
max_login_retries=self.max_login_retries,
|
|
max_total_time=self.max_total_time,
|
|
proxy_manager=self._shared_proxy_manager,
|
|
)
|
|
result = loginer.login()
|
|
|
|
if result.success:
|
|
task.status = "success"
|
|
task.cookie = result.cookie
|
|
task.message = result.message or "登录成功"
|
|
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']} 登录失败: {result.message}")
|
|
|
|
except Exception as e:
|
|
task.status = "error"
|
|
task.message = str(e)
|
|
self._push_log("error", f"[{current}] {acc_info['username']} 登录异常: {e}")
|
|
|
|
task.finished_at = datetime.now(timezone.utc)
|
|
worker_db.commit()
|
|
|
|
finally:
|
|
worker_db.close()
|
|
|
|
def run(self):
|
|
"""在线程中执行批量登录。"""
|
|
batch_id = self.batch_id
|
|
concurrency = self.concurrency
|
|
self._push_log("info", f"批量登录任务 {batch_id} 开始，共 {len(self.account_ids)} 个账号，并发数: {concurrency}")
|
|
|
|
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}")
|
|
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,
|
|
)
|
|
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,
|
|
},
|
|
})
|
|
|
|
self.db.commit()
|
|
total = len(task_infos)
|
|
if total == 0:
|
|
self._push_log("warning", "没有可执行的账号")
|
|
self._push_log("result", "")
|
|
return
|
|
|
|
# 并发执行登录，每个账号独立获取代理
|
|
with ThreadPoolExecutor(max_workers=concurrency) as executor:
|
|
futures = []
|
|
for item in task_infos:
|
|
if self._stop.is_set():
|
|
self._push_log("warning", "任务已停止，跳过剩余账号")
|
|
break
|
|
future = executor.submit(
|
|
self._execute_one,
|
|
item["task_id"],
|
|
item["acc_info"],
|
|
total,
|
|
)
|
|
futures.append(future)
|
|
|
|
# 等待所有任务完成
|
|
for future in as_completed(futures):
|
|
try:
|
|
future.result()
|
|
except Exception as e:
|
|
self._push_log("error", f"Worker 异常: {e}")
|
|
|
|
self._push_log("info", f"批量登录任务 {batch_id} 完成")
|
|
self._push_log("result", "")
|
|
finally:
|
|
# 确保 DB Session 被关闭，避免连接泄漏
|
|
self.db.close()
|
|
|
|
|
|
class BatchRegistry:
|
|
"""管理运行中的登录批次状态。"""
|
|
|
|
def __init__(self):
|
|
self._batches: dict[str, dict] = {}
|
|
self._lock = threading.Lock()
|
|
|
|
def register(self, batch_id: str, log_queue: asyncio.Queue,
|
|
loop: asyncio.AbstractEventLoop, runner: LoginBatchRunner):
|
|
with self._lock:
|
|
self._batches[batch_id] = {
|
|
"log_queue": log_queue,
|
|
"loop": loop,
|
|
"runner": runner,
|
|
}
|
|
|
|
def get(self, batch_id: str):
|
|
with self._lock:
|
|
return self._batches.get(batch_id)
|
|
|
|
def pop(self, batch_id: str):
|
|
with self._lock:
|
|
return self._batches.pop(batch_id, None)
|
|
|
|
|
|
# 模块级单例
|
|
batch_registry = BatchRegistry()
|
|
|
|
|
|
# 在模块末尾导入 SessionLocal（避免循环导入）
|
|
from ..database import SessionLocal
|