compliance-scanner-agent/compliance-dast/src/tools/rate_limit_tester.rs

use std::time::Instant;

use compliance_core::error::CoreError;
use compliance_core::models::dast::{DastEvidence, DastFinding, DastVulnType};
use compliance_core::models::Severity;
use compliance_core::traits::pentest_tool::{PentestTool, PentestToolContext, PentestToolResult};
use serde_json::json;
use tracing::{info, warn};

/// Tool that tests whether a target enforces rate limiting.
pub struct RateLimitTesterTool {
    http: reqwest::Client,
}

impl RateLimitTesterTool {
    pub fn new(http: reqwest::Client) -> Self {
        Self { http }
    }
}

impl PentestTool for RateLimitTesterTool {
    fn name(&self) -> &str {
        "rate_limit_tester"
    }

    fn description(&self) -> &str {
        "Tests whether an endpoint enforces rate limiting by sending rapid sequential requests. \
         Checks for 429 responses and measures response time degradation."
    }

    fn input_schema(&self) -> serde_json::Value {
        json!({
            "type": "object",
            "properties": {
                "url": {
                    "type": "string",
                    "description": "URL of the endpoint to test for rate limiting"
                },
                "method": {
                    "type": "string",
                    "description": "HTTP method to use",
                    "enum": ["GET", "POST", "PUT", "PATCH", "DELETE"],
                    "default": "GET"
                },
                "request_count": {
                    "type": "integer",
                    "description": "Number of rapid requests to send (default: 50)",
                    "default": 50,
                    "minimum": 10,
                    "maximum": 200
                },
                "body": {
                    "type": "string",
                    "description": "Optional request body (for POST/PUT/PATCH)"
                }
            },
            "required": ["url"]
        })
    }

    fn execute<'a>(
        &'a self,
        input: serde_json::Value,
        context: &'a PentestToolContext,
    ) -> std::pin::Pin<Box<dyn std::future::Future<Output = Result<PentestToolResult, CoreError>> + Send + 'a>> {
        Box::pin(async move {
        let url = input
            .get("url")
            .and_then(|v| v.as_str())
            .ok_or_else(|| CoreError::Dast("Missing required 'url' parameter".to_string()))?;

        let method = input
            .get("method")
            .and_then(|v| v.as_str())
            .unwrap_or("GET");

        let request_count = input
            .get("request_count")
            .and_then(|v| v.as_u64())
            .unwrap_or(50)
            .min(200) as usize;

        let body = input.get("body").and_then(|v| v.as_str());

        let target_id = context
            .target
            .id
            .map(|oid| oid.to_hex())
            .unwrap_or_else(|| "unknown".to_string());

        // Respect the context rate limit if set
        let max_requests = if context.rate_limit > 0 {
            request_count.min(context.rate_limit as usize * 5)
        } else {
            request_count
        };

        let mut status_codes: Vec<u16> = Vec::with_capacity(max_requests);
        let mut response_times: Vec<u128> = Vec::with_capacity(max_requests);
        let mut got_429 = false;
        let mut rate_limit_at_request: Option<usize> = None;

        for i in 0..max_requests {
            let start = Instant::now();

            let request = match method {
                "POST" => {
                    let mut req = self.http.post(url);
                    if let Some(b) = body {
                        req = req.body(b.to_string());
                    }
                    req
                }
                "PUT" => {
                    let mut req = self.http.put(url);
                    if let Some(b) = body {
                        req = req.body(b.to_string());
                    }
                    req
                }
                "PATCH" => {
                    let mut req = self.http.patch(url);
                    if let Some(b) = body {
                        req = req.body(b.to_string());
                    }
                    req
                }
                "DELETE" => self.http.delete(url),
                _ => self.http.get(url),
            };

            match request.send().await {
                Ok(resp) => {
                    let elapsed = start.elapsed().as_millis();
                    let status = resp.status().as_u16();
                    status_codes.push(status);
                    response_times.push(elapsed);

                    if status == 429 && !got_429 {
                        got_429 = true;
                        rate_limit_at_request = Some(i + 1);
                        info!(url, request_num = i + 1, "Rate limit triggered (429)");
                    }

                    // Check for rate limit headers even on 200
                    if !got_429 {
                        let headers = resp.headers();
                        let has_rate_headers = headers.contains_key("x-ratelimit-limit")
                            || headers.contains_key("x-ratelimit-remaining")
                            || headers.contains_key("ratelimit-limit")
                            || headers.contains_key("ratelimit-remaining")
                            || headers.contains_key("retry-after");

                        if has_rate_headers && rate_limit_at_request.is_none() {
                            // Server has rate limit headers but hasn't blocked yet
                        }
                    }
                }
                Err(e) => {
                    let elapsed = start.elapsed().as_millis();
                    status_codes.push(0);
                    response_times.push(elapsed);
                }
            }
        }

        let mut findings = Vec::new();
        let total_sent = status_codes.len();
        let count_429 = status_codes.iter().filter(|&&s| s == 429).count();
        let count_success = status_codes.iter().filter(|&&s| (200..300).contains(&s)).count();

        // Calculate response time statistics
        let avg_time = if !response_times.is_empty() {
            response_times.iter().sum::<u128>() / response_times.len() as u128
        } else {
            0
        };

        let first_half_avg = if response_times.len() >= 4 {
            let half = response_times.len() / 2;
            response_times[..half].iter().sum::<u128>() / half as u128
        } else {
            avg_time
        };

        let second_half_avg = if response_times.len() >= 4 {
            let half = response_times.len() / 2;
            response_times[half..].iter().sum::<u128>() / (response_times.len() - half) as u128
        } else {
            avg_time
        };

        // Significant time degradation suggests possible (weak) rate limiting
        let time_degradation = if first_half_avg > 0 {
            (second_half_avg as f64 / first_half_avg as f64) - 1.0
        } else {
            0.0
        };

        let rate_data = json!({
            "total_requests_sent": total_sent,
            "status_429_count": count_429,
            "success_count": count_success,
            "rate_limit_at_request": rate_limit_at_request,
            "avg_response_time_ms": avg_time,
            "first_half_avg_ms": first_half_avg,
            "second_half_avg_ms": second_half_avg,
            "time_degradation_pct": (time_degradation * 100.0).round(),
        });

        if !got_429 && count_success == total_sent {
            // No rate limiting detected at all
            let evidence = DastEvidence {
                request_method: method.to_string(),
                request_url: url.to_string(),
                request_headers: None,
                request_body: body.map(String::from),
                response_status: 200,
                response_headers: None,
                response_snippet: Some(format!(
                    "Sent {total_sent} rapid requests. All returned success (2xx). \
                     No 429 responses received. Avg response time: {avg_time}ms."
                )),
                screenshot_path: None,
                payload: None,
                response_time_ms: Some(avg_time as u64),
            };

            let mut finding = DastFinding::new(
                String::new(),
                target_id.clone(),
                DastVulnType::RateLimitAbsent,
                format!("No rate limiting on {} {}", method, url),
                format!(
                    "The endpoint {} {} does not enforce rate limiting. \
                     {total_sent} rapid requests were all accepted with no 429 responses \
                     or noticeable degradation. This makes the endpoint vulnerable to \
                     brute force attacks and abuse.",
                    method, url
                ),
                Severity::Medium,
                url.to_string(),
                method.to_string(),
            );
            finding.cwe = Some("CWE-770".to_string());
            finding.evidence = vec![evidence];
            finding.remediation = Some(
                "Implement rate limiting on this endpoint. Use token bucket or sliding window \
                 algorithms. Return 429 Too Many Requests with a Retry-After header when the \
                 limit is exceeded."
                    .to_string(),
            );
            findings.push(finding);
            warn!(url, method, total_sent, "No rate limiting detected");
        } else if got_429 {
            info!(
                url,
                method,
                rate_limit_at = ?rate_limit_at_request,
                "Rate limiting is enforced"
            );
        }

        let count = findings.len();

        Ok(PentestToolResult {
            summary: if got_429 {
                format!(
                    "Rate limiting is enforced on {method} {url}. \
                     429 response received after {} requests.",
                    rate_limit_at_request.unwrap_or(0)
                )
            } else if count > 0 {
                format!(
                    "No rate limiting detected on {method} {url} after {total_sent} requests."
                )
            } else {
                format!("Rate limit testing complete for {method} {url}.")
            },
            findings,
            data: rate_data,
        })
        })
    }
}