robotex/src/main.py

import json
import logging
import traceback
import uuid
from pathlib import Path
from typing import Annotated
from urllib.parse import urlsplit

from fastapi import Cookie
from fastapi import FastAPI
from fastapi import Request
from fastapi import Response
from patchright.async_api import Page, ProxySettings

from api import actions
from api.geoip import resolve_geo
from api.schemas import SolveRequest
from engine import AsyncStealthySession

logger = logging.getLogger(__name__)
logger.setLevel(logging.DEBUG)
logger.addHandler(logging.StreamHandler())
BROWSER_SLOTS = 10


def parse_proxy(proxy_url: str) -> ProxySettings:
	parts = urlsplit(proxy_url)
	server = f"{parts.scheme}://{parts.hostname}:{parts.port}"
	settings: ProxySettings = {"server": server}
	if parts.username:
		settings["username"] = parts.username
	if parts.password:
		settings["password"] = parts.password
	return settings


def make_page_action(payload: SolveRequest):
	"""Замыкание для AsyncStealthySession.fetch(page_action=...): выполняет очередь действий
	(включая прохождение капчи/антибота через api.actions.execute) и достаёт из уже
	навигированной страницы всё, что нужно для ответа API. Возврат этой функции — то, что
	вернёт fetch() (см. engine/stealthy.py)."""

	async def page_action(page: Page) -> dict:
		actions_result = await actions.execute(
			page, payload.actions, payload.captcha,
			antibot_name=payload.captcha_type,
			success_locator=payload.captcha_success_locator,
			detect_locator=payload.captcha_detect_locator,
		)
		return {
			"actions": actions_result,
			"cookies": await page.context.cookies(),
			"html": await page.inner_html("html"),
			"user_agent": await page.evaluate("() => navigator.userAgent"),
		}

	return page_action


app = FastAPI()


@app.get("/")
async def root():
	return {"message": "Hello World"}


@app.post("/v1/solve")
async def solve(
	request: Request,
	payload: SolveRequest,
	response: Response,
	session_id: Annotated[str | None, Cookie()] = None
):
	# TODO: добавить поддержку api ключей
	api_key = request.headers.get('X-API-Key')

	logger.info("Incoming payload: [%s]\n%s", session_id, payload.model_dump_json(indent=2, ensure_ascii=True))
	global BROWSER_SLOTS

	if BROWSER_SLOTS <= 0:
		return {
			"ok": False,
			"error": "No slots available",
		}
	BROWSER_SLOTS -= 1

	try:
		if not session_id:
			session_id = str(uuid.uuid4())
			response.set_cookie("session_id", session_id)
			logger.info("Setup new session with session_id: %s", session_id)

		proxy = parse_proxy(payload.proxy)

		# профиль (cookies/local storage/fingerprint.json) хранится на диске по
		# session_id — если он уже приходил в куках, ниже подгружаем существующий
		# fingerprint из этой папки; если нет (новый session_id) — генерируем новый
		# и сохраняем туда же (см. api/fingerprint.py:load_or_create_fingerprint)
		user_data_dir = Path(__file__).parent / f"../extra/user_data_dir/{session_id}"
		if not user_data_dir.exists():
			user_data_dir.mkdir(parents=True)

		# geo — до создания fingerprint'а, чтобы производная от прокси locale
		# (см. api/geoip.py) попала в сам fingerprint, а не подменяла его
		# задним числом
		geo = await resolve_geo(proxy)
		locale = geo.get("locale") if geo else None

		fp_kwargs = {}
		if payload.user_agent:
			fp_kwargs["user_agent"] = payload.user_agent
		# TODO: если есть session_id, то user_agent нужно взять из настроек пользователя
		if locale:
			fp_kwargs["locale"] = locale

		nav_timeout_ms = payload.timeout * 1000 if payload.timeout else 60_000
		blocked_domains = set(payload.blocked_domains) if payload.blocked_domains else None

		kwargs = dict(
			headless=False,
			proxy=proxy,
			user_data_dir=str(user_data_dir),
			real_chrome=True,
			locale=locale,
			timezone_id=geo.get("timezone_id") if geo else None,
			cookies=payload.cookies,
			disable_resources=payload.disable_resources,
			blocked_domains=blocked_domains,
			timeout=nav_timeout_ms,
			retries=max(1, payload.retries),
			retry_delay=payload.retry_delay,
			wait=payload.wait,
			google_search=payload.google_search
		)

		async with AsyncStealthySession(**kwargs) as engine:
			result = await engine.fetch(payload.url, page_action=make_page_action(payload))
	except Exception as e:
		traceback.print_exc()
		result_payload = {
			"error": str(e),
			"status": False
		}
	else:
		status = not "error" in result["actions"][-1]
		error = ["Something error", None][status]

		result_payload = {
			"status": ["err", "ok"][status],
			"user_agent": result["user_agent"],
			"cookies": result["cookies"],
			"actions": result["actions"],
		}
		if payload.html:
			result_payload["html"] = result["html"]
		if error:
			result_payload["error"] = error
	finally:
		BROWSER_SLOTS += 1

	logger.info("Outgoing payload:\n%s", json.dumps(result_payload, indent=2, sort_keys=True, ensure_ascii=True))
	return result_payload