captchaforge 0.2.38

//! YOLOv8-based image grid CAPTCHA solver.
//!
//! Solves reCAPTCHA v2 / hCaptcha image grids by running local
//! object detection on the challenge screenshot. Only activates when
//! the challenge task maps to a known COCO class (traffic lights,
//! buses, cars, bicycles, etc.). Falls back to VLM for unknown
//! classes (crosswalks, chimneys, stairs, palm trees, etc.).
//!
//! Requires the `vision` feature for actual inference. Without it,
//! the solver is inert (`supports` returns false).

use super::*;
#[cfg(feature = "vision")]
use super::grid_click;
#[cfg(feature = "vision")]
use anyhow::Context;
#[cfg(feature = "vision")]
use base64::Engine as _;



/// Solver that uses YOLOv8 ONNX inference for image grid CAPTCHAs.
///
/// When the `vision` feature is enabled, the model is loaded lazily
/// on the first `solve()` call. If the model is not cached and cannot
/// be downloaded, the solver is inert.
pub struct YoloGridSolver {
    config: SolveConfig,
}

impl Default for YoloGridSolver {
    fn default() -> Self {
        Self::new()
    }
}

impl YoloGridSolver {
    pub fn new() -> Self {
        Self {
            config: SolveConfig::default(),
        }
    }

    pub fn with_config(mut self, config: SolveConfig) -> Self {
        self.config = config;
        self
    }

    /// Extract the challenge task text from the page (e.g. "Select all
    /// images with traffic lights").
    #[cfg(feature = "vision")]
    async fn extract_task(&self, page: &Page) -> Option<String> {
        let probes = [
            ".rc-imageselect-desc-no-canonical",
            ".rc-imageselect-desc",
            ".task-text",
            ".captcha-prompt",
            "#widget > div",
            ".captcha-task",
        ];

        for sel in &probes {
            if let Ok(result) = page.evaluate(format!(r#"
                (() => {{
                    const el = document.querySelector('{}');
                    return el && el.textContent ? el.textContent.trim() : '';
                }})()
            "#, sel)).await {
                if let Ok(text) = result.into_value::<String>() {
                    if !text.is_empty() {
                        return Some(text);
                    }
                }
            }
        }
        None
    }
}

#[async_trait]
impl CaptchaSolver for YoloGridSolver {
    fn name(&self) -> &'static str {
        "YoloGridSolver"
    }

    fn method(&self) -> SolveMethod {
        SolveMethod::VisionLLM
    }

    fn supports(&self, kind: &crate::captcha_detect::DetectedCaptcha) -> bool {
        #[cfg(feature = "vision")]
        {
            // We can't know if the model is available without trying to load it,
            // but we declare support for the types and let solve() handle failure.
            matches!(
                kind,
                crate::captcha_detect::DetectedCaptcha::RecaptchaV2
                    | crate::captcha_detect::DetectedCaptcha::HCaptcha
                    | crate::captcha_detect::DetectedCaptcha::ImageCaptcha
            )
        }
        #[cfg(not(feature = "vision"))]
        {
            let _ = kind;
            false
        }
    }

    async fn solve(
        &self,
        page: &Page,
        _captcha_info: &CaptchaInfo,
    ) -> Result<CaptchaSolveResult> {
        let _ = page;
        let t0 = Instant::now();

        #[cfg(not(feature = "vision"))]
        {
            return Ok(CaptchaSolveResult::failure(self.method(), t0.elapsed().as_millis() as u64));
        }

        #[cfg(feature = "vision")]
        {
            // Lazy-load detector.
            let hub = crate::vision::ModelHub::new();
            let path = hub.resolve(crate::vision::ModelId::YoloV8n).await
                .map_err(|e| anyhow!("YOLO model unavailable: {e}"))?;
            let mut detector = crate::vision::YoloDetector::load(&path)
                .map_err(|e| anyhow!("YOLO load failed: {e}"))?;

            // 1. Extract task text.
            let task = self
                .extract_task(page)
                .await
                .unwrap_or_else(|| "Select all matching images".to_string());

            // 2. Check if task maps to known COCO classes.
            let target_classes = crate::vision::yolo::task_to_coco_classes(&task);
            if target_classes.is_none() {
                return Ok(CaptchaSolveResult::failure(
                    self.method(),
                    t0.elapsed().as_millis() as u64,
                ));
            }
            let target_classes = target_classes.unwrap();

            // 3. Locate the challenge tiles across the frame tree FIRST. The
            //    reCAPTCHA/hCaptcha image grid renders inside a cross-origin
            //    OOPIF, so its tiles are invisible to main-document JS; they are
            //    found via BiDi per-frame evaluation. If no grid is visible
            //    there is nothing to click and detection would be wasted work.
            let Some((tiles, _tile_sel)) = grid_click::locate_grid_tiles(page).await else {
                return Ok(CaptchaSolveResult::failure(
                    self.method(),
                    t0.elapsed().as_millis() as u64,
                ));
            };

            // 4. Screenshot the page, then crop to JUST the grid region so the
            //    detector's normalized [0,1] coordinates map 1:1 onto the tile
            //    grid. Without the crop, detections are normalized over the whole
            //    viewport but mapped as if the screenshot were the grid — correct
            //    only for a full-bleed local fixture, never for a floating real
            //    reCAPTCHA popup. The crop box is scaled from CSS to device
            //    pixels (the BiDi screenshot is captured at the device pixel
            //    ratio).
            let screenshot_b64 = super::util::screenshot_b64(page).await?;
            let screenshot_bytes = base64::engine::general_purpose::STANDARD
                .decode(screenshot_b64.as_bytes())
                .map_err(|e| anyhow::anyhow!("decode screenshot: {e}"))?;
            let full = image::load_from_memory(&screenshot_bytes)
                .context("load screenshot as image")?;
            let (img_w, img_h) = (full.width(), full.height());
            let (vw, vh) = viewport_css_size(page)
                .await
                .unwrap_or((img_w as f64, img_h as f64));
            let crop_box = grid_click::grid_region(&tiles)
                .and_then(|r| grid_click::scale_rect_to_image(r, vw, vh, img_w, img_h));
            let image = match crop_box {
                Some((x, y, w, h)) => full.crop_imm(x, y, w, h),
                None => full,
            };

            // 5. Run YOLO detection (synchronous — done before any await).
            let detections = detector.detect(&image)?;

            if detections.is_empty() {
                return Ok(CaptchaSolveResult::failure(
                    self.method(),
                    t0.elapsed().as_millis() as u64,
                ));
            }

            // 6. Determine grid layout from the tiles' own geometry (cluster
            //    distinct row/column centres) and map detections — now
            //    normalized over the grid crop — to tile indices.
            let (grid_w, grid_h) = grid_click::infer_grid_dims(&tiles).unwrap_or((3, 3));
            let tile_w = 1.0 / grid_w as f32;
            let tile_h = 1.0 / grid_h as f32;

            let mut selected_tiles: std::collections::HashSet<usize> = std::collections::HashSet::new();

            for det in &detections {
                if !target_classes.contains(&det.class.as_str()) {
                    continue;
                }

                // Compute which tile(s) the bounding box center falls into.
                let cx = (det.bbox[0] + det.bbox[2]) / 2.0;
                let cy = (det.bbox[1] + det.bbox[3]) / 2.0;

                let tile_x = (cx / tile_w).floor() as usize;
                let tile_y = (cy / tile_h).floor() as usize;

                if tile_x < grid_w && tile_y < grid_h {
                    let tile_idx = tile_y * grid_w + tile_x;
                    selected_tiles.insert(tile_idx);
                }
            }

            if selected_tiles.is_empty() {
                return Ok(CaptchaSolveResult::failure(
                    self.method(),
                    t0.elapsed().as_millis() as u64,
                ));
            }

            // 7. Click the selected tiles with TRUSTED, cross-origin-correct
            //    input — BiDi `input.performActions` (isTrusted === true) routed
            //    into the owning frame, the only click a real image CAPTCHA
            //    accepts — then click the verify button. Reuses the tiles already
            //    located in step 3, so there is no second frame walk.
            let indices: Vec<usize> = selected_tiles.into_iter().collect();
            let clicked_count =
                grid_click::click_located_tiles(page, &tiles, &indices).await;

            if clicked_count == 0 {
                return Ok(CaptchaSolveResult::failure(
                    self.method(),
                    t0.elapsed().as_millis() as u64,
                ));
            }

            // 8. Poll for token.
            tokio::time::sleep(Duration::from_millis(500)).await;

            let token = poll_for_token(page, self.config).await;
            let elapsed = t0.elapsed().as_millis() as u64;

            Ok(CaptchaSolveResult {
                solution: token.clone(),
                confidence: if token.is_empty() { 0.5 } else { 0.85 },
                method: self.method(),
                time_ms: elapsed,
                success: !token.is_empty(),
                screenshot: None,
                cookies: Vec::new(),
                verified_outcome: None,
            })
        }
    }
}

/// Poll the page for a populated CAPTCHA response token.
#[cfg(feature = "vision")]
async fn poll_for_token(page: &Page, config: SolveConfig) -> String {
    let deadline = Instant::now() + Duration::from_millis(config.token_max_attempts as u64 * config.token_poll_interval_ms);
    loop {
        let raw = page
            .evaluate(crate::solver::wait_for_token::TOKEN_PROBE_JS)
            .await;
        if let Ok(Ok(Some(t))) = raw.map(|r| r.into_value::<Option<String>>()) {
            if !t.is_empty() {
                return t;
            }
        }
        if Instant::now() >= deadline {
            return String::new();
        }
        tokio::time::sleep(Duration::from_millis(config.token_poll_interval_ms)).await;
    }
}

/// Read the CSS viewport size (`window.innerWidth`/`innerHeight`) so a CSS-pixel
/// crop box can be scaled to the device-pixel screenshot. Returns `None` if the
/// page reports a non-positive dimension.
#[cfg(feature = "vision")]
async fn viewport_css_size(page: &Page) -> Option<(f64, f64)> {
    let result = page
        .evaluate("[window.innerWidth, window.innerHeight]")
        .await
        .ok()?;
    let arr: Vec<f64> = result.into_value().ok()?;
    if arr.len() == 2 && arr[0] > 0.0 && arr[1] > 0.0 {
        Some((arr[0], arr[1]))
    } else {
        None
    }
}