use serde::{Deserialize, Serialize};
#[derive(Debug, Clone, Deserialize, Serialize)]
#[serde(untagged)]
pub enum WaitBehavior {
Fixed(u64),
Range {
#[serde(rename = "min")]
min_ms: u64,
#[serde(rename = "max")]
max_ms: u64,
},
Function(String),
}
impl WaitBehavior {
pub fn get_duration_ms(&self) -> u64 {
match self {
WaitBehavior::Fixed(ms) => *ms,
WaitBehavior::Range { min_ms, max_ms } => {
use rand::Rng;
rand::thread_rng().gen_range(*min_ms..=*max_ms)
}
WaitBehavior::Function(js_func) => {
Self::execute_js_wait_function(js_func).unwrap_or(100)
}
}
}
fn execute_js_wait_function(js_func: &str) -> Option<u64> {
let trimmed = js_func.trim();
if !trimmed.starts_with("function") {
return None;
}
if let Some(body) = extract_function_body(trimmed) {
if body.contains("var min") && body.contains("var max") {
return Self::parse_solo_wait_pattern(&body);
}
let body = body
.replace("return ", "")
.trim_end_matches(';')
.to_string();
if body.contains("Math.random()") {
use rand::Rng;
let re = regex::Regex::new(
r"Math\.floor\s*\(\s*Math\.random\s*\(\s*\)\s*\*\s*(\d+)\s*\)\s*\+\s*(\d+)",
)
.ok()?;
if let Some(caps) = re.captures(&body) {
let range = caps.get(1)?.as_str().parse::<u64>().ok()?;
let offset = caps.get(2)?.as_str().parse::<u64>().ok()?;
return Some(rand::thread_rng().gen_range(offset..=offset + range));
}
let re = regex::Regex::new(r"Math\.random\s*\(\s*\)\s*\*\s*(\d+)").ok()?;
if let Some(caps) = re.captures(&body) {
let range = caps.get(1)?.as_str().parse::<u64>().ok()?;
return Some(rand::thread_rng().gen_range(0..=range));
}
}
body.trim().parse::<u64>().ok()
} else {
None
}
}
fn parse_solo_wait_pattern(body: &str) -> Option<u64> {
use rand::Rng;
let min_re = regex::Regex::new(r"var\s+min\s*=\s*Math\.ceil\s*\(\s*(\d+)\s*\)").ok()?;
let min_val = min_re
.captures(body)
.and_then(|c| c.get(1))
.and_then(|m| m.as_str().parse::<u64>().ok())
.unwrap_or(0);
let max_re = regex::Regex::new(r"var\s+max\s*=\s*Math\.floor\s*\(\s*(\d+)\s*\)").ok()?;
let max_val = max_re
.captures(body)
.and_then(|c| c.get(1))
.and_then(|m| m.as_str().parse::<u64>().ok())
.unwrap_or(0);
if max_val >= min_val {
Some(rand::thread_rng().gen_range(min_val..=max_val))
} else {
Some(min_val)
}
}
}
fn extract_function_body(js_func: &str) -> Option<String> {
let start = js_func.find('{')?;
let end = js_func.rfind('}')?;
if start < end {
Some(js_func[start + 1..end].trim().to_string())
} else {
None
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_wait_behavior_fixed() {
let wait = WaitBehavior::Fixed(100);
assert_eq!(wait.get_duration_ms(), 100);
}
#[test]
fn test_wait_behavior_range() {
let wait = WaitBehavior::Range {
min_ms: 100,
max_ms: 200,
};
for _ in 0..10 {
let duration = wait.get_duration_ms();
assert!((100..=200).contains(&duration));
}
}
#[test]
fn test_wait_behavior_serde() {
let yaml = "100";
let wait: WaitBehavior = serde_yaml::from_str(yaml).unwrap();
assert!(matches!(wait, WaitBehavior::Fixed(100)));
let yaml = "min: 100\nmax: 200";
let wait: WaitBehavior = serde_yaml::from_str(yaml).unwrap();
assert!(matches!(
wait,
WaitBehavior::Range {
min_ms: 100,
max_ms: 200
}
));
}
#[test]
fn test_wait_behavior_solo_js_pattern() {
let js = " function() { var min = Math.ceil(0); var max = Math.floor(0); var num = Math.floor(Math.random() * (max - min + 1)); var wait = (num + min); return wait; } ";
let wait = WaitBehavior::Function(js.to_string());
assert_eq!(wait.get_duration_ms(), 0);
let js = "function() { var min = Math.ceil(50); var max = Math.floor(100); var num = Math.floor(Math.random() * (max - min + 1)); var wait = (num + min); return wait; }";
let wait = WaitBehavior::Function(js.to_string());
for _ in 0..10 {
let duration = wait.get_duration_ms();
assert!(
(50..=100).contains(&duration),
"Duration {} not in range 50-100",
duration
);
}
}
#[test]
fn test_wait_behavior_js_function() {
let js = "function() { return Math.floor(Math.random() * 100) + 50; }";
let wait = WaitBehavior::Function(js.to_string());
for _ in 0..10 {
let duration = wait.get_duration_ms();
assert!(
(50..=150).contains(&duration),
"Duration {} not in range 50-150",
duration
);
}
}
}