DouBao2Api/app/browser_manager.py

import asyncio
import json
from typing import Optional

from loguru import logger
from playwright.async_api import (
    Browser,
    BrowserContext,
    Page,
    async_playwright,
)

from .config import settings

STEALTH_JS = """
Object.defineProperty(navigator, 'webdriver', {get: () => undefined});
Object.defineProperty(navigator, 'plugins', {
    get: () => [1, 2, 3, 4, 5],
});
Object.defineProperty(navigator, 'languages', {
    get: () => ['zh-CN', 'zh', 'en'],
});
Object.defineProperty(navigator, 'platform', {
    get: () => 'Win32',
});
window.chrome = { runtime: {} };
const originalQuery = window.navigator.permissions.query;
window.navigator.permissions.query = (parameters) =>
    parameters.name === 'notifications'
        ? Promise.resolve({ state: Notification.permission })
        : originalQuery(parameters);
"""


class BrowserManager:
    _instance: Optional["BrowserManager"] = None
    _lock: asyncio.Lock = asyncio.Lock()

    def __new__(cls) -> "BrowserManager":
        if cls._instance is None:
            cls._instance = super().__new__(cls)
        return cls._instance

    def __init__(self):
        if not hasattr(self, "_initialized"):
            self._initialized = True
            self.playwright = None
            self.browser: Optional[Browser] = None
            self.context: Optional[BrowserContext] = None
            self.page: Optional[Page] = None
            self._ready = False

    @property
    def ready(self) -> bool:
        return self._ready and self.page is not None and not self.page.is_closed()

    async def start(self) -> Page:
        async with self._lock:
            if self.ready:
                return self.page

            logger.info("Starting Playwright browser...")
            self.playwright = await async_playwright().start()

            self.browser = await self.playwright.chromium.launch(
                headless=settings.browser_headless,
                args=[
                    "--no-sandbox",
                    "--disable-setuid-sandbox",
                    "--disable-dev-shm-usage",
                    "--disable-blink-features=AutomationControlled",
                    "--disable-features=IsolateOrigins,site-per-process",
                    "--window-size=1280,720",
                ],
            )

            storage_state = (
                str(settings.state_file) if settings.state_file.exists() else None
            )

            self.context = await self.browser.new_context(
                viewport={
                    "width": settings.viewport_width,
                    "height": settings.viewport_height,
                },
                user_agent=(
                    "Mozilla/5.0 (Windows NT 10.0; Win64; x64) "
                    "AppleWebKit/537.36 (KHTML, like Gecko) "
                    "Chrome/131.0.0.0 Safari/537.36"
                ),
                locale="zh-CN",
                timezone_id="Asia/Shanghai",
                storage_state=storage_state,
            )

            await self.context.add_init_script(STEALTH_JS)
            self.context.set_default_timeout(settings.page_load_timeout)
            self.context.set_default_navigation_timeout(60000)

            self.page = await self.context.new_page()

            await self._navigate_to_doubao()
            self._ready = True
            logger.info("Browser started and navigated to Doubao")
            return self.page

    async def _navigate_to_doubao(self):
        if not self.page:
            return
        try:
            await self.page.goto(settings.doubao_url, wait_until="domcontentloaded")
            await self.page.wait_for_timeout(3000)
        except Exception as e:
            logger.warning(f"Navigation warning: {e}")

    async def save_state(self):
        if self.context:
            await self.context.storage_state(path=str(settings.state_file))
            cookies = await self.context.cookies()
            with open(settings.cookies_file, "w", encoding="utf-8") as f:
                json.dump(cookies, f, ensure_ascii=False, indent=2)
            logger.info(f"Saved browser state ({len(cookies)} cookies)")

    async def check_login(self) -> bool:
        """Check if the user is logged in by looking for login indicators."""
        if not self.ready:
            return False
        try:
            login_indicators = await self.page.evaluate(
                """
                () => {
                    const body = document.body ? document.body.innerText : '';
                    const url = window.location.href;
                    if (url.includes('login') || url.includes('passport')) return false;

                    const loginBtn = document.querySelector(
                        '[class*="login"], [class*="sign-in"], [data-testid*="login"]'
                    );
                    if (loginBtn && loginBtn.offsetParent !== null) return false;

                    const inputArea = document.querySelector(
                        'textarea, [contenteditable="true"], [role="textbox"]'
                    );
                    return !!inputArea;
                }
                """
            )
            return login_indicators
        except Exception as e:
            logger.warning(f"Login check failed: {e}")
            return False

    async def get_page(self) -> Page:
        if not self.ready:
            await self.start()
        return self.page

    async def reload(self):
        if self.page:
            await self.page.reload(wait_until="domcontentloaded")
            await self.page.wait_for_timeout(2000)

    async def restart(self):
        logger.info("Restarting browser session...")
        await self.stop()
        await asyncio.sleep(1)
        await self.start()

    async def stop(self):
        self._ready = False
        try:
            if self.context:
                await self.save_state()
        except Exception as e:
            logger.warning(f"Failed to save state during stop: {e}")
        try:
            if self.browser:
                await self.browser.close()
        except Exception:
            pass
        try:
            if self.playwright:
                await self.playwright.stop()
        except Exception:
            pass
        self.page = None
        self.context = None
        self.browser = None
        self.playwright = None
        logger.info("Browser stopped")


browser_manager = BrowserManager()