108 lines
3.3 KiB
Python
108 lines
3.3 KiB
Python
"""
|
|
Utility functions for cookie management and Playwright setup
|
|
"""
|
|
from pathlib import Path
|
|
from playwright.async_api import async_playwright
|
|
import json
|
|
from loguru import logger
|
|
from app.core.config import settings
|
|
|
|
|
|
async def set_init_script(context):
|
|
"""
|
|
Add stealth script to prevent bot detection
|
|
|
|
Args:
|
|
context: Playwright browser context
|
|
|
|
Returns:
|
|
Modified context
|
|
"""
|
|
# Add stealth.js if available
|
|
stealth_js_path = settings.BASE_DIR / "app" / "services" / "uploader" / "stealth.min.js"
|
|
|
|
if stealth_js_path.exists():
|
|
await context.add_init_script(path=stealth_js_path)
|
|
|
|
# Grant geolocation permission
|
|
await context.grant_permissions(['geolocation'])
|
|
|
|
return context
|
|
|
|
|
|
async def generate_cookie_with_qr(platform: str, platform_url: str, account_file: str):
|
|
"""
|
|
Generate cookie by scanning QR code with Playwright
|
|
|
|
Args:
|
|
platform: Platform name (for logging)
|
|
platform_url: Platform login URL
|
|
account_file: Path to save cookies
|
|
|
|
Returns:
|
|
bool: Success status
|
|
"""
|
|
try:
|
|
logger.info(f"[{platform}] 开始自动生成 Cookie...")
|
|
|
|
async with async_playwright() as playwright:
|
|
browser = await playwright.chromium.launch(headless=False)
|
|
context = await browser.new_context()
|
|
|
|
# Add stealth script
|
|
context = await set_init_script(context)
|
|
|
|
page = await context.new_page()
|
|
await page.goto(platform_url)
|
|
|
|
logger.info(f"[{platform}] 请在浏览器中扫码登录...")
|
|
logger.info(f"[{platform}] 登录后点击 Playwright Inspector 的 '继续' 按钮")
|
|
|
|
# Pause for user to login
|
|
await page.pause()
|
|
|
|
# Save cookies
|
|
await context.storage_state(path=account_file)
|
|
|
|
await browser.close()
|
|
|
|
logger.success(f"[{platform}] Cookie 已保存到: {account_file}")
|
|
return True
|
|
|
|
except Exception as e:
|
|
logger.exception(f"[{platform}] Cookie 生成失败: {e}")
|
|
return False
|
|
|
|
|
|
async def extract_bilibili_cookies(account_file: str):
|
|
"""
|
|
Extract specific Bilibili cookies needed by biliup
|
|
|
|
Args:
|
|
account_file: Path to cookies file
|
|
|
|
Returns:
|
|
dict: Extracted cookies
|
|
"""
|
|
try:
|
|
# Read Playwright storage_state format
|
|
with open(account_file, 'r', encoding='utf-8') as f:
|
|
storage = json.load(f)
|
|
|
|
# Extract cookies
|
|
cookie_dict = {}
|
|
for cookie in storage.get('cookies', []):
|
|
if cookie['name'] in ['SESSDATA', 'bili_jct', 'DedeUserID', 'DedeUserID__ckMd5']:
|
|
cookie_dict[cookie['name']] = cookie['value']
|
|
|
|
# Save in biliup format
|
|
with open(account_file, 'w', encoding='utf-8') as f:
|
|
json.dump(cookie_dict, f, indent=2)
|
|
|
|
logger.info(f"[B站] Cookie 已转换为 biliup 格式")
|
|
return cookie_dict
|
|
|
|
except Exception as e:
|
|
logger.exception(f"[B站] Cookie 提取失败: {e}")
|
|
return {}
|