use anyhow::Result;
use std::future::Future;
use std::pin::Pin;
use super::Tester;
#[derive(Clone)]
pub struct StatusChecker {
proxy: Option<String>,
proxy_auth: Option<String>,
timeout: u64,
retries: u32,
random_agent: bool,
insecure: bool,
include_status: Option<Vec<String>>,
exclude_status: Option<Vec<String>>,
}
impl StatusChecker {
pub fn new() -> Self {
StatusChecker {
proxy: None,
proxy_auth: None,
timeout: 30,
retries: 3,
random_agent: false,
insecure: false,
include_status: None,
exclude_status: None,
}
}
pub fn with_include_status(&mut self, status_codes: Option<Vec<String>>) {
self.include_status = status_codes;
}
pub fn with_exclude_status(&mut self, status_codes: Option<Vec<String>>) {
self.exclude_status = status_codes;
}
fn status_matches_pattern(&self, status_code: u16, pattern: &str) -> bool {
if pattern.contains('x') || pattern.contains('X') {
let status_str = status_code.to_string();
let pattern = pattern.to_lowercase();
if status_str.len() != pattern.len() {
return false;
}
for (s, p) in status_str.chars().zip(pattern.chars()) {
if p != 'x' && p != s {
return false;
}
}
true
} else {
if let Ok(pattern_code) = pattern.parse::<u16>() {
status_code == pattern_code
} else {
false
}
}
}
fn matches_any_pattern(&self, status_code: u16, patterns: &[String]) -> bool {
if patterns.is_empty() {
return false;
}
patterns.iter().any(|pattern| {
pattern
.split(',')
.any(|subpattern| self.status_matches_pattern(status_code, subpattern.trim()))
})
}
fn should_include_status(&self, status_code: u16) -> bool {
if let Some(include_patterns) = &self.include_status {
return self.matches_any_pattern(status_code, include_patterns);
}
if let Some(exclude_patterns) = &self.exclude_status {
return !self.matches_any_pattern(status_code, exclude_patterns);
}
true
}
}
impl Tester for StatusChecker {
fn clone_box(&self) -> Box<dyn Tester> {
Box::new(self.clone())
}
fn test_url<'a>(
&'a self,
url: &'a str,
) -> Pin<Box<dyn Future<Output = Result<Vec<String>>> + Send + 'a>> {
Box::pin(async move {
let mut client_builder =
reqwest::Client::builder().timeout(std::time::Duration::from_secs(self.timeout));
if self.insecure {
client_builder = client_builder.danger_accept_invalid_certs(true);
}
if self.random_agent {
let ua = crate::network::random_user_agent();
client_builder = client_builder.user_agent(ua);
}
if let Some(proxy_url) = &self.proxy {
let mut proxy = reqwest::Proxy::all(proxy_url)?;
if let Some(auth) = &self.proxy_auth {
let parts: Vec<&str> = auth.splitn(2, ':').collect();
if parts.len() == 2 {
proxy = proxy.basic_auth(parts[0], parts[1]);
}
}
client_builder = client_builder.proxy(proxy);
}
let client = client_builder.build()?;
let mut last_error = None;
for _ in 0..=self.retries {
match client.get(url).send().await {
Ok(response) => {
let status = response.status();
let status_code = status.as_u16();
if !self.should_include_status(status_code) {
return Ok(vec![]); }
let status_text = format!(
"{} {}",
status_code,
status.canonical_reason().unwrap_or("")
);
return Ok(vec![format!("{} - {}", url, status_text)]);
}
Err(e) => {
last_error = Some(e);
tokio::time::sleep(std::time::Duration::from_millis(500)).await;
continue;
}
}
}
Err(anyhow::anyhow!(
"Failed to check status for {}: {:?}",
url,
last_error
))
})
}
fn with_timeout(&mut self, seconds: u64) {
self.timeout = seconds;
}
fn with_retries(&mut self, count: u32) {
self.retries = count;
}
fn with_random_agent(&mut self, enabled: bool) {
self.random_agent = enabled;
}
fn with_insecure(&mut self, enabled: bool) {
self.insecure = enabled;
}
fn with_proxy(&mut self, proxy: Option<String>) {
self.proxy = proxy;
}
fn with_proxy_auth(&mut self, auth: Option<String>) {
self.proxy_auth = auth;
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_status_matches_pattern() {
let checker = StatusChecker::new();
assert!(checker.status_matches_pattern(200, "200"));
assert!(!checker.status_matches_pattern(200, "404"));
assert!(checker.status_matches_pattern(200, "2xx"));
assert!(checker.status_matches_pattern(200, "20x"));
assert!(checker.status_matches_pattern(201, "20x"));
assert!(checker.status_matches_pattern(404, "4xx"));
assert!(!checker.status_matches_pattern(200, "3xx"));
assert!(!checker.status_matches_pattern(200, "4xx"));
assert!(checker.status_matches_pattern(200, "2XX"));
assert!(checker.status_matches_pattern(404, "4XX"));
}
#[test]
fn test_matches_any_pattern() {
let checker = StatusChecker::new();
assert!(checker.matches_any_pattern(200, &["200".to_string()]));
assert!(!checker.matches_any_pattern(404, &["200".to_string()]));
assert!(checker.matches_any_pattern(200, &["200".to_string(), "404".to_string()]));
assert!(checker.matches_any_pattern(404, &["200".to_string(), "404".to_string()]));
assert!(!checker.matches_any_pattern(301, &["200".to_string(), "404".to_string()]));
assert!(checker.matches_any_pattern(200, &["2xx".to_string()]));
assert!(checker.matches_any_pattern(404, &["2xx".to_string(), "4xx".to_string()]));
assert!(checker.matches_any_pattern(200, &["200,404".to_string()]));
assert!(checker.matches_any_pattern(404, &["200,404".to_string()]));
assert!(checker.matches_any_pattern(200, &["2xx,404".to_string()]));
assert!(!checker.matches_any_pattern(301, &["200,404".to_string()]));
}
#[test]
fn test_should_include_status() {
let mut checker = StatusChecker::new();
assert!(checker.should_include_status(200));
assert!(checker.should_include_status(404));
assert!(checker.should_include_status(500));
checker.with_include_status(Some(vec!["200".to_string(), "3xx".to_string()]));
assert!(checker.should_include_status(200));
assert!(checker.should_include_status(301));
assert!(!checker.should_include_status(404));
assert!(!checker.should_include_status(500));
checker.with_include_status(None);
checker.with_exclude_status(Some(vec!["4xx".to_string(), "500".to_string()]));
assert!(checker.should_include_status(200));
assert!(checker.should_include_status(301));
assert!(!checker.should_include_status(404));
assert!(!checker.should_include_status(500));
checker.with_include_status(Some(vec!["200".to_string()]));
checker.with_exclude_status(Some(vec!["2xx".to_string()]));
assert!(checker.should_include_status(200));
assert!(!checker.should_include_status(201));
}
}