robotex/src/api/schemas.py

from typing import Any, Literal
from urllib.parse import urlsplit

from pydantic import BaseModel, Field, field_validator, model_validator

CAPTCHA_TYPES = ("cloudflare", "recaptcha", "hcaptcha")
# максимум browser_idle_timeout в минутах: открытое окно держит слот браузера (#78)
MAX_BROWSER_IDLE_TIMEOUT = 120
PROXY_SCHEMES = ("http", "https", "socks5")

# тип действия -> обязательные поля (кроме type)
_REQUIRED_FIELDS: dict[str, tuple[str, ...]] = {
	"fill": ("selector", "value"),
	"captcha": ("selector", "value"),
	"delay": ("value",),
	"mouse_move": ("selector",),
	"click": ("selector",),
	"submit": ("selector",),
	"inner_html": ("selector",),
	"inner_text": ("selector",),
	"wait_for": ("selector",),
	"pass_challenge": (),
}
WAIT_FOR_DEFAULT_TIMEOUT = 30


class ActionRequest(BaseModel):
	"""
	Действие автоматизации. Характеризуется типом и производится над селектором.
	Набор параметров зависит от типа, этот класс перечисляет все параметры.
	Обязательные поля для каждого типа проверяются здесь же — неверное действие
	отклоняется с 422 до запуска браузера.
	"""
	type: str
	selector: str | None = None
	value: str | None = None
	timeout: int | None = None

	@model_validator(mode="after")
	def check_fields(self):
		if self.type not in _REQUIRED_FIELDS:
			raise ValueError(f"unknown action type {self.type!r}, expected one of: {', '.join(_REQUIRED_FIELDS)}")

		missing = [f for f in _REQUIRED_FIELDS[self.type] if getattr(self, f) in (None, "")]
		if missing:
			raise ValueError(f"action {self.type!r} requires: {', '.join(missing)}")

		if self.type == "wait_for":
			if self.timeout is None:
				self.timeout = WAIT_FOR_DEFAULT_TIMEOUT
			if not 1 <= self.timeout <= 7200:
				raise ValueError("wait_for timeout must be between 1 and 7200 seconds")
		elif self.type == "delay":
			try:
				delay = int(self.value)
			except ValueError:
				raise ValueError("delay value must be an integer number of seconds") from None
			if not 0 <= delay <= 600:
				raise ValueError("delay value must be between 0 and 600 seconds")
		elif self.type == "captcha" and self.value not in CAPTCHA_TYPES:
			raise ValueError(f"captcha value must be one of: {', '.join(CAPTCHA_TYPES)}")
		return self


class SolveRequest(BaseModel):
	"""
	Объект запроса для автоматизации на сайте.
	"""
	url: str
	proxy: str | None = None
	timeout: int | None = None
	cookies: list[dict[str, Any]] | None = None
	# List[{name: str, value: str, domain: str, path: str, expires: float, httpOnly: bool, secure: bool, sameSite: Union["Lax", "None", "Strict"], partitionKey: Union[str, None]}]

	actions: list[ActionRequest]

	html: bool = False
	screenshot: bool = False

	captcha: str | None = None
	captcha_detect_locator: str | None = None
	captcha_success_locator: str | None = None
	captcha_type: str | None = None

	# ниже — опции навигации, портированные из scrapling.engines._browsers (StealthySession.fetch),
	# все выключены/нейтральны по умолчанию, чтобы поведение для текущих вызывающих не менялось
	disable_resources: bool = False
	blocked_domains: list[str] | None = None
	retries: int = 1
	retry_delay: float = 1.0
	wait: int = 0
	google_search: bool = False

	@model_validator(mode="before")
	@classmethod
	def reject_browser_identity(cls, data):
		# UA и размер окна задаёт сам браузер: подмена расходится с userAgentData и движком,
		# и Cloudflare такой браузер не пропускает (см. docs/Usage.md)
		if isinstance(data, dict) and "user_agent" in data:
			raise ValueError(
				"user_agent is not supported: take user_agent from the response instead (see docs/Usage.md)"
			)
		return data

	@field_validator("proxy")
	@classmethod
	def check_proxy(cls, proxy: str | None) -> str | None:
		if not proxy:
			return None
		try:
			parts = urlsplit(proxy)
			port = parts.port
		except ValueError as e:
			raise ValueError(f"invalid proxy url: {e}") from None
		if parts.scheme not in PROXY_SCHEMES:
			raise ValueError(f"proxy scheme must be one of: {', '.join(PROXY_SCHEMES)}")
		if not parts.hostname or not port:
			raise ValueError("proxy must be in form scheme://[user:pass@]host:port")
		return proxy

	@model_validator(mode="after")
	def check_captcha_key(self):
		if not self.captcha and any(a.type == "captcha" for a in self.actions):
			raise ValueError("action 'captcha' requires request field 'captcha' (2captcha key)")
		return self


class ActionResult(BaseModel):
	type: str
	ok: bool
	result: str | None = None
	duration_ms: int


class SolveResponse(BaseModel):
	"""
	Ответ /v1/solve. status — "ok" (антибот пройден, все действия выполнены) или "err".
	code — машиночитаемая причина ошибки, error — текст для человека.
	"""
	status: Literal["ok", "err"]
	code: Literal["no_slots", "request_failed", "challenge_failed", "action_failed"] | None = None
	error: str | None = None
	user_agent: str | None = None
	cookies: list[dict[str, Any]] | None = None
	actions: list[ActionResult] = []
	html: str | None = None
	screenshot: str | None = None
	duration_ms: int | None = None


class TaskRequest(SolveRequest):
	"""
	Задача POST /v2/tasks: тело как у /v1/solve плюс удержание окна браузера.

	url необязателен, если окно сессии открыто: тогда действия выполняются на текущей странице.
	"""
	url: str | None = None
	# минуты простоя, через которые окно сессии закроется; 0 — закрыть сразу после задачи
	browser_idle_timeout: int = Field(0, ge=0, le=MAX_BROWSER_IDLE_TIMEOUT)


class TaskCreated(BaseModel):
	task_id: str
	status: str
	url: str


class TaskResponse(BaseModel):
	"""
	GET /v2/tasks/{id}. status: queued → running → done | failed.
	done — результат в result (в нём свой status ok/err, как в /v1/solve);
	failed — задачу не удалось выполнить (воркер упал), причина в error.
	"""
	task_id: str
	status: Literal["queued", "running", "done", "failed", "cancelled"]
	error: str | None = None
	result: SolveResponse | None = None