captchaforge 0.2.28

//! AWS WAF Captcha dedicated solver.
//!
//! AWS WAF's CAPTCHA action surfaces a visual puzzle (sliding-tile,
//! visual-prompt, or audio fallback) hosted on
//! `captcha.awswaf.com` / `captcha-prod.awswaf.com`. On a successful
//! solve, the AWS WAF JS Integration SDK populates two values that
//! the host application uses to verify the session:
//!
//! - **`aws-waf-token` cookie** — the canonical session token.
//!   Format: `<base64>.<base64>.<base64>.<base64>:<expiry>:<base64>`
//!   (5 colon-separated segments). Long enough that any non-trivial
//!   value is highly likely real.
//! - **`AWSALB` / `AWSALBCORS`** — load-balancer affinity cookies set
//!   alongside; corroborating but not the gate.
//! - **`window.awsWafCaptcha?.iv`** + **`window.awsWafCaptcha?.context`**
//!   — the challenge primitives used by the SDK to construct the
//!   sensor POST. Surfaced here for callers that need to drive a
//!   custom verification flow rather than relying on the cookie.
//!
//! This solver covers the cookie-pass case: with proper stealth
//! ([`crate::stealth`] + a coherent [`crate::stealth_profiles`]),
//! AWS WAF often issues the token without forcing the visual
//! challenge. When the visual challenge IS surfaced (slider variant),
//! the chain falls through to [`super::super::SliderCaptchaSolver`];
//! visual-prompt variant falls through to
//! [`super::super::VlmCaptchaSolver`].

use super::super::*;
use crate::captcha_detect::DetectedCaptcha;
use std::time::{Duration, Instant};

/// Default total wait budget. AWS WAF cookie-pass typically resolves
/// in 2–5s; 12s upper bound matches `WaitForTokenSolver`.
const DEFAULT_MAX_WAIT_MS: u64 = 12_000;

/// AWS WAF Captcha cookie-watch + iv/context capture solver.
pub struct AwsWafSolver {
    max_wait_ms: u64,
}

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

impl AwsWafSolver {
    pub fn new() -> Self {
        Self {
            max_wait_ms: DEFAULT_MAX_WAIT_MS,
        }
    }

    pub fn with_max_wait_ms(mut self, ms: u64) -> Self {
        self.max_wait_ms = ms;
        self
    }
}

/// Pure validator for the `aws-waf-token` cookie shape.
///
/// AWS WAF tokens follow the documented `seg1.seg2.seg3.seg4:expiry:seg5`
/// pattern: at least 4 dot-separated segments before the first colon,
/// then a numeric expiry, then a final segment. We match the
/// structural shape rather than byte-exact format because AWS rotates
/// segment encodings periodically (base64 vs base64url) but the
/// `<dotted>.<colon>:<num>:<base64>` skeleton has held since 2021.
///
/// # Examples
///
/// ```
/// use captchaforge::solver::vendors::aws_waf::is_valid_aws_waf_token;
///
/// assert!(!is_valid_aws_waf_token(""));
/// assert!(!is_valid_aws_waf_token("short"));
/// // Real-shape token: 4 dot-segments, colon-numeric-colon-tail
/// let real = "AQIDB.aGVsbG8.d29ybGQ.dGVzdA:1736899200:c2lnbg";
/// assert!(is_valid_aws_waf_token(real));
/// // Missing the dotted prefix segments
/// assert!(!is_valid_aws_waf_token("token:1736899200:sig"));
/// ```
pub fn is_valid_aws_waf_token(value: &str) -> bool {
    if value.len() < 40 {
        return false;
    }
    let trimmed = value.trim_matches('"');
    let parts: Vec<&str> = trimmed.splitn(3, ':').collect();
    if parts.len() != 3 {
        return false;
    }
    // Part 0: dotted segments. Must have at least 4 segments,
    // each non-empty.
    let dotted: Vec<&str> = parts[0].split('.').collect();
    if dotted.len() < 4 || dotted.iter().any(|s| s.is_empty()) {
        return false;
    }
    // Part 1: numeric expiry (unix epoch seconds).
    let Ok(t) = parts[1].parse::<u64>() else {
        return false;
    };
    if !(1_577_836_800..4_102_444_800).contains(&t) {
        return false;
    }
    // Part 2: signature/tail, non-empty.
    !parts[2].is_empty()
}

/// JS that probes the page for AWS WAF state. Returns:
///   - `{ phase: "passed", cookie: "..." }` — `aws-waf-token` cookie set
///   - `{ phase: "puzzle" }` — captcha iframe rendered, awaiting interaction
///   - `{ phase: "interstitial" }` — sensor SDK loaded, no cookie yet
///   - `{ phase: "passed_with_iv", iv: "...", context: "..." }` — when
///     the SDK has staged iv+context but cookie hasn't landed yet
///     (caller can drive a custom verification flow)
///   - `{ phase: "unknown" }` — no AWS WAF signals on the page
const PHASE_PROBE_JS: &str = r#"
(() => {
    /* Cookie pass — the canonical aws-waf-token. */
    let token = '';
    for (const c of (document.cookie || '').split(';')) {
        const t = c.trim();
        if (t.startsWith('aws-waf-token=')) {
            token = t.slice('aws-waf-token='.length);
            break;
        }
    }
    if (token) return { phase: 'passed', cookie: token };
    /* iv + context staged on the SDK global. */
    if (window.awsWafCaptcha && window.awsWafCaptcha.iv && window.awsWafCaptcha.context) {
        return {
            phase: 'passed_with_iv',
            iv: window.awsWafCaptcha.iv,
            context: window.awsWafCaptcha.context,
        };
    }
    /* Puzzle iframe rendered (visible challenge). */
    if (document.querySelector('iframe[src*="captcha.awswaf.com"], iframe[src*="captcha-prod.awswaf.com"], #captcha-container, [data-aws-waf-captcha]')) {
        return { phase: 'puzzle' };
    }
    /* SDK loaded but nothing surfaced yet. */
    const sdkPresent =
        !!document.querySelector('script[src*="captcha.awswaf.com"], script[src*="captcha-prod.awswaf.com"]') ||
        typeof window.AwsWafIntegration !== 'undefined' ||
        typeof window.AwsWafCaptcha !== 'undefined';
    if (sdkPresent) return { phase: 'interstitial' };
    return { phase: 'unknown' };
})()
"#;

#[derive(Debug, serde::Deserialize)]
struct PhaseProbe {
    phase: String,
    #[serde(default)]
    cookie: String,
    #[serde(default)]
    iv: String,
    #[serde(default)]
    context: String,
}

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

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

    fn supports(&self, kind: &DetectedCaptcha) -> bool {
        // AWS WAF is detected as Custom("aws_waf_captcha") by the
        // bundled community rule. Also accept SliderCaptcha because
        // AWS WAF's puzzle variant sometimes surfaces via the slider
        // detector first.
        matches!(
            kind,
            DetectedCaptcha::Custom(_) | DetectedCaptcha::SliderCaptcha
        )
    }

    async fn solve(&self, page: &Page, _info: &CaptchaInfo) -> Result<CaptchaSolveResult> {
        let t0 = Instant::now();
        let deadline = t0 + Duration::from_millis(self.max_wait_ms);

        loop {
            let raw = page
                .evaluate(PHASE_PROBE_JS)
                .await
                .map_err(|e| anyhow!("AWS WAF probe failed: {e}"))?;
            let probe: PhaseProbe = match raw.into_value() {
                Ok(p) => p,
                Err(_) => PhaseProbe {
                    phase: "interstitial".into(),
                    cookie: String::new(),
                    iv: String::new(),
                    context: String::new(),
                },
            };

            match probe.phase.as_str() {
                "passed" if is_valid_aws_waf_token(&probe.cookie) => {
                    let cookies = crate::cookies::capture_from_page(page)
                        .await
                        .unwrap_or_default();
                    return Ok(CaptchaSolveResult {
                        solution: "aws_waf:cookie".into(),
                        confidence: 1.0,
                        method: SolveMethod::AutoPass,
                        time_ms: t0.elapsed().as_millis() as u64,
                        success: true,
                        screenshot: None,
                        cookies,
                        verified_outcome: None,
                    });
                }
                "passed" => {
                    // Cookie present but malformed — keep waiting in
                    // case it's an in-progress write.
                }
                "passed_with_iv" => {
                    // SDK has staged iv+context — surface them in the
                    // solution payload as JSON so callers driving a
                    // custom verification flow can pick them up.
                    let payload = serde_json::json!({
                        "type": "aws_waf_iv_context",
                        "iv": probe.iv,
                        "context": probe.context,
                    });
                    let cookies = crate::cookies::capture_from_page(page)
                        .await
                        .unwrap_or_default();
                    return Ok(CaptchaSolveResult {
                        solution: payload.to_string(),
                        confidence: 0.85,
                        method: SolveMethod::AutoPass,
                        time_ms: t0.elapsed().as_millis() as u64,
                        success: true,
                        screenshot: None,
                        cookies,
                        verified_outcome: None,
                    });
                }
                "puzzle" => {
                    // Visible challenge — yield to slider/VLM via
                    // screenshot+failure so the chain advances.
                    let shot = super::super::screenshot_b64(page).await.ok();
                    return Ok(CaptchaSolveResult {
                        solution: String::new(),
                        confidence: 0.0,
                        method: SolveMethod::AutoPass,
                        time_ms: t0.elapsed().as_millis() as u64,
                        success: false,
                        screenshot: shot,
                        cookies: Vec::new(),
                        verified_outcome: None,
                    });
                }
                _ => {} // interstitial / unknown — keep waiting
            }

            if Instant::now() >= deadline {
                return Ok(CaptchaSolveResult::failure(
                    SolveMethod::AutoPass,
                    t0.elapsed().as_millis() as u64,
                ));
            }
            tokio::time::sleep(Duration::from_millis(400)).await;
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn defaults_are_sane() {
        assert_eq!(AwsWafSolver::new().max_wait_ms, 12_000);
    }

    #[test]
    fn builders_override() {
        assert_eq!(
            AwsWafSolver::new().with_max_wait_ms(5_000).max_wait_ms,
            5_000
        );
    }

    #[test]
    fn name_method_stable() {
        let s = AwsWafSolver::new();
        assert_eq!(s.name(), "AwsWafSolver");
        assert_eq!(s.method(), SolveMethod::AutoPass);
    }

    #[test]
    fn supports_custom_and_slider() {
        let s = AwsWafSolver::new();
        assert!(s.supports(&DetectedCaptcha::Custom("aws_waf_captcha".into())));
        assert!(s.supports(&DetectedCaptcha::SliderCaptcha));
        assert!(!s.supports(&DetectedCaptcha::Turnstile));
        assert!(!s.supports(&DetectedCaptcha::HCaptcha));
        assert!(!s.supports(&DetectedCaptcha::RecaptchaV2));
    }

    #[test]
    fn token_rejects_short_or_empty() {
        assert!(!is_valid_aws_waf_token(""));
        assert!(!is_valid_aws_waf_token("short"));
        assert!(!is_valid_aws_waf_token(&"a".repeat(39))); // below 40-char floor
    }

    #[test]
    fn token_rejects_missing_segments() {
        // Length passes but no dotted prefix
        assert!(!is_valid_aws_waf_token(
            "longstring1234567890longstring:1736899200:sig"
        ));
        // Only 3 dotted segments (<4)
        assert!(!is_valid_aws_waf_token(
            "a.b.c:1736899200:sig0123456789abcd"
        ));
    }

    #[test]
    fn token_rejects_non_numeric_expiry() {
        assert!(!is_valid_aws_waf_token(
            "AQIDB.aGVsbG8.d29ybGQ.dGVzdA:not-a-time:c2lnbg"
        ));
    }

    #[test]
    fn token_rejects_out_of_range_expiry() {
        // pre-2020 (incidentally also pre-AWS-WAF-Captcha-launch in 2022)
        assert!(!is_valid_aws_waf_token(
            "AQIDB.aGVsbG8.d29ybGQ.dGVzdA:1000000000:c2lnbg"
        ));
        // Far-future garbage
        assert!(!is_valid_aws_waf_token(
            "AQIDB.aGVsbG8.d29ybGQ.dGVzdA:5000000000:c2lnbg"
        ));
    }

    #[test]
    fn token_rejects_empty_tail() {
        assert!(!is_valid_aws_waf_token(
            "AQIDB.aGVsbG8.d29ybGQ.dGVzdA:1736899200:"
        ));
    }

    #[test]
    fn token_accepts_real_shape() {
        let real = "AQIDB.aGVsbG8.d29ybGQ.dGVzdA:1736899200:c2lnbg";
        assert!(is_valid_aws_waf_token(real));
    }

    #[test]
    fn token_accepts_quoted_value() {
        let real = "\"AQIDB.aGVsbG8.d29ybGQ.dGVzdA:1736899200:c2lnbg\"";
        assert!(is_valid_aws_waf_token(real));
    }

    #[test]
    fn token_rejects_empty_dotted_segment() {
        // Empty middle segment violates the structural shape.
        assert!(!is_valid_aws_waf_token(
            "AQIDB..d29ybGQ.dGVzdA:1736899200:c2lnbg"
        ));
    }

    #[test]
    fn phase_probe_js_covers_documented_signals() {
        for needle in [
            "aws-waf-token=",
            "captcha.awswaf.com",
            "captcha-prod.awswaf.com",
            "awswafintegration",
            "awswafcaptcha",
            "iv",
            "context",
            "captcha-container",
        ] {
            assert!(
                PHASE_PROBE_JS.to_lowercase().contains(needle),
                "PHASE_PROBE_JS must reference: {needle}"
            );
        }
    }

    #[test]
    fn phase_probe_js_emits_canonical_phases() {
        for phase in ["passed", "passed_with_iv", "puzzle", "interstitial", "unknown"] {
            assert!(
                PHASE_PROBE_JS.contains(&format!("'{phase}'")),
                "PHASE_PROBE_JS must emit phase: {phase}"
            );
        }
    }
}