dekun
268 lines
8.9 KiB
Python
268 lines
8.9 KiB
Python
import asyncio
|
|
import base64
|
|
import time
|
|
import uuid
|
|
from dataclasses import dataclass, field
|
|
from typing import Any, Optional
|
|
|
|
from playwright.async_api import Browser, BrowserContext, Page, Playwright, async_playwright
|
|
|
|
from app.security import get_idle_timeout, get_screencast_quality, get_viewport_size
|
|
|
|
|
|
@dataclass
|
|
class BrowserSession:
|
|
session_id: str
|
|
url: str
|
|
playwright: Playwright
|
|
browser: Browser
|
|
context: BrowserContext
|
|
page: Page
|
|
cdp: Any
|
|
created_at: float = field(default_factory=time.time)
|
|
last_activity: float = field(default_factory=time.time)
|
|
subscribers: set[asyncio.Queue] = field(default_factory=set)
|
|
screencast_task: Optional[asyncio.Task] = None
|
|
idle_task: Optional[asyncio.Task] = None
|
|
closed: bool = False
|
|
viewport_width: int = 1280
|
|
viewport_height: int = 720
|
|
|
|
|
|
class BrowserManager:
|
|
def __init__(self, max_sessions: int = 1) -> None:
|
|
self.max_sessions = max_sessions
|
|
self._sessions: dict[str, BrowserSession] = {}
|
|
self._lock = asyncio.Lock()
|
|
|
|
async def create_session(self, url: str) -> BrowserSession:
|
|
async with self._lock:
|
|
if len(self._sessions) >= self.max_sessions:
|
|
raise RuntimeError(
|
|
f"已达最大会话数 ({self.max_sessions})，请先关闭现有会话"
|
|
)
|
|
|
|
session_id = str(uuid.uuid4())
|
|
width, height = get_viewport_size()
|
|
quality = get_screencast_quality()
|
|
|
|
playwright = await async_playwright().start()
|
|
browser = await playwright.chromium.launch(
|
|
headless=True,
|
|
args=[
|
|
"--no-sandbox",
|
|
"--disable-setuid-sandbox",
|
|
"--disable-dev-shm-usage",
|
|
"--disable-gpu",
|
|
],
|
|
)
|
|
context = await browser.new_context(
|
|
viewport={"width": width, "height": 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"
|
|
),
|
|
)
|
|
page = await context.new_page()
|
|
cdp = await context.new_cdp_session(page)
|
|
|
|
session = BrowserSession(
|
|
session_id=session_id,
|
|
url=url,
|
|
playwright=playwright,
|
|
browser=browser,
|
|
context=context,
|
|
page=page,
|
|
cdp=cdp,
|
|
viewport_width=width,
|
|
viewport_height=height,
|
|
)
|
|
|
|
await page.goto(url, wait_until="domcontentloaded", timeout=60000)
|
|
session.url = page.url
|
|
|
|
page.on(
|
|
"framenavigated",
|
|
lambda frame: asyncio.create_task(
|
|
self._on_frame_navigated(session, frame)
|
|
),
|
|
)
|
|
|
|
session.screencast_task = asyncio.create_task(
|
|
self._run_screencast(session, quality)
|
|
)
|
|
session.idle_task = asyncio.create_task(self._watch_idle(session))
|
|
|
|
self._sessions[session_id] = session
|
|
return session
|
|
|
|
def session_count(self) -> int:
|
|
return len(self._sessions)
|
|
|
|
async def get_session(self, session_id: str) -> Optional[BrowserSession]:
|
|
return self._sessions.get(session_id)
|
|
|
|
async def close_session(self, session_id: str) -> None:
|
|
async with self._lock:
|
|
session = self._sessions.pop(session_id, None)
|
|
if session:
|
|
await self._cleanup_session(session)
|
|
|
|
async def close_all(self) -> None:
|
|
async with self._lock:
|
|
session_ids = list(self._sessions.keys())
|
|
for session_id in session_ids:
|
|
await self.close_session(session_id)
|
|
|
|
def touch(self, session: BrowserSession) -> None:
|
|
session.last_activity = time.time()
|
|
|
|
async def navigate(self, session: BrowserSession, url: str) -> str:
|
|
self.touch(session)
|
|
await session.page.goto(url, wait_until="domcontentloaded", timeout=60000)
|
|
session.url = session.page.url
|
|
return session.url
|
|
|
|
async def go_back(self, session: BrowserSession) -> str:
|
|
self.touch(session)
|
|
await session.page.go_back(wait_until="domcontentloaded", timeout=30000)
|
|
session.url = session.page.url
|
|
return session.url
|
|
|
|
async def go_forward(self, session: BrowserSession) -> str:
|
|
self.touch(session)
|
|
await session.page.go_forward(wait_until="domcontentloaded", timeout=30000)
|
|
session.url = session.page.url
|
|
return session.url
|
|
|
|
async def reload(self, session: BrowserSession) -> str:
|
|
self.touch(session)
|
|
await session.page.reload(wait_until="domcontentloaded", timeout=60000)
|
|
session.url = session.page.url
|
|
return session.url
|
|
|
|
def subscribe(self, session: BrowserSession) -> asyncio.Queue:
|
|
queue: asyncio.Queue = asyncio.Queue(maxsize=8)
|
|
session.subscribers.add(queue)
|
|
return queue
|
|
|
|
def unsubscribe(self, session: BrowserSession, queue: asyncio.Queue) -> None:
|
|
session.subscribers.discard(queue)
|
|
|
|
async def _broadcast(self, session: BrowserSession, message: dict) -> None:
|
|
dead: list[asyncio.Queue] = []
|
|
for queue in session.subscribers:
|
|
try:
|
|
queue.put_nowait(message)
|
|
except asyncio.QueueFull:
|
|
try:
|
|
queue.get_nowait()
|
|
queue.put_nowait(message)
|
|
except asyncio.QueueEmpty:
|
|
pass
|
|
except Exception:
|
|
dead.append(queue)
|
|
for queue in dead:
|
|
session.subscribers.discard(queue)
|
|
|
|
async def _on_frame_navigated(self, session: BrowserSession, frame) -> None:
|
|
if session.closed or frame != session.page.main_frame:
|
|
return
|
|
session.url = session.page.url
|
|
await self._broadcast(session, {"type": "url_update", "url": session.url})
|
|
|
|
async def _run_screencast(self, session: BrowserSession, quality: int) -> None:
|
|
async def on_screencast_frame(params: dict) -> None:
|
|
if session.closed:
|
|
return
|
|
data = params.get("data", "")
|
|
session_id = params.get("sessionId")
|
|
try:
|
|
await session.cdp.send(
|
|
"Page.screencastFrameAck", {"sessionId": session_id}
|
|
)
|
|
except Exception:
|
|
return
|
|
try:
|
|
frame_bytes = base64.b64decode(data)
|
|
except Exception:
|
|
return
|
|
await self._broadcast(
|
|
session,
|
|
{
|
|
"type": "frame",
|
|
"data": frame_bytes,
|
|
"url": session.url,
|
|
"width": session.viewport_width,
|
|
"height": session.viewport_height,
|
|
},
|
|
)
|
|
|
|
def schedule_frame(params: dict) -> None:
|
|
asyncio.create_task(on_screencast_frame(params))
|
|
|
|
session.cdp.on("Page.screencastFrame", schedule_frame)
|
|
|
|
await session.cdp.send(
|
|
"Page.startScreencast",
|
|
{
|
|
"format": "jpeg",
|
|
"quality": quality,
|
|
"maxWidth": session.viewport_width,
|
|
"maxHeight": session.viewport_height,
|
|
"everyNthFrame": 1,
|
|
},
|
|
)
|
|
|
|
try:
|
|
while not session.closed:
|
|
await asyncio.sleep(1)
|
|
finally:
|
|
try:
|
|
await session.cdp.send("Page.stopScreencast")
|
|
except Exception:
|
|
pass
|
|
|
|
async def _watch_idle(self, session: BrowserSession) -> None:
|
|
timeout = get_idle_timeout()
|
|
try:
|
|
while not session.closed:
|
|
await asyncio.sleep(30)
|
|
if time.time() - session.last_activity > timeout:
|
|
await self.close_session(session.session_id)
|
|
break
|
|
except asyncio.CancelledError:
|
|
pass
|
|
|
|
async def _cleanup_session(self, session: BrowserSession) -> None:
|
|
if session.closed:
|
|
return
|
|
session.closed = True
|
|
|
|
for task in (session.screencast_task, session.idle_task):
|
|
if task and not task.done():
|
|
task.cancel()
|
|
try:
|
|
await task
|
|
except asyncio.CancelledError:
|
|
pass
|
|
|
|
await self._broadcast(session, {"type": "closed", "reason": "session_ended"})
|
|
|
|
try:
|
|
await session.context.close()
|
|
except Exception:
|
|
pass
|
|
try:
|
|
await session.browser.close()
|
|
except Exception:
|
|
pass
|
|
try:
|
|
await session.playwright.stop()
|
|
except Exception:
|
|
pass
|
|
|
|
|
|
browser_manager = BrowserManager()
|