use gumdrop::Options;
use httpmock::MockServer;
use std::io::{self, BufRead};
use swanling::metrics::SwanlingMetrics;
use swanling::swanling::{SwanlingTask, SwanlingTaskSet};
use swanling::{SwanlingAttack, SwanlingConfiguration};
type WorkerHandles = Vec<std::thread::JoinHandle<()>>;
#[allow(dead_code)]
pub fn build_configuration(server: &MockServer, custom: Vec<&str>) -> SwanlingConfiguration {
let mut configuration: Vec<&str> = vec![];
let server_url = server.base_url();
configuration.extend_from_slice(&custom);
if !configuration.contains(&"--worker") {
if !configuration.contains(&"--host") {
configuration.extend_from_slice(&["--host", &server_url]);
}
if !configuration.contains(&"--users") {
configuration.extend_from_slice(&["--users", "1"]);
}
if !configuration.contains(&"--hatch-rate") {
configuration.extend_from_slice(&["--hatch-rate", "1"]);
}
if !configuration.contains(&"--run-time") {
configuration.extend_from_slice(&["--run-time", "1"]);
}
if !configuration.contains(&"--co-mitigation") {
configuration.extend_from_slice(&["--co-mitigation", "disabled"]);
}
}
SwanlingConfiguration::parse_args_default(&configuration)
.expect("failed to parse options and generate a configuration")
}
#[allow(dead_code)]
pub fn launch_gaggle_workers(
swanling_attack: SwanlingAttack,
expect_workers: usize,
) -> WorkerHandles {
let mut worker_handles = Vec::new();
for _ in 0..expect_workers {
let worker_swanling_attack = swanling_attack.clone();
worker_handles.push(std::thread::spawn(move || {
run_load_test(worker_swanling_attack, None);
}));
}
worker_handles
}
#[allow(dead_code)]
pub fn build_load_test(
configuration: SwanlingConfiguration,
taskset: &SwanlingTaskSet,
start_task: Option<&SwanlingTask>,
stop_task: Option<&SwanlingTask>,
) -> SwanlingAttack {
let mut swanling = crate::SwanlingAttack::initialize_with_config(configuration)
.unwrap()
.register_taskset(taskset.clone());
if let Some(task) = start_task {
swanling = swanling.test_start(task.clone());
}
if let Some(task) = stop_task {
swanling = swanling.test_stop(task.clone());
}
swanling
}
pub fn run_load_test(
swanling_attack: SwanlingAttack,
worker_handles: Option<WorkerHandles>,
) -> SwanlingMetrics {
let swanling_metrics = swanling_attack.execute().unwrap();
if let Some(handles) = worker_handles {
for handle in handles {
let _ = handle.join();
}
}
swanling_metrics
}
#[allow(dead_code)]
pub fn file_length(file_name: &str) -> usize {
if let Ok(file) = std::fs::File::open(std::path::Path::new(file_name)) {
io::BufReader::new(file).lines().count()
} else {
0
}
}
#[allow(dead_code)]
pub fn cleanup_files(files: Vec<&str>) {
for file in files {
if std::path::Path::new(file).exists() {
std::fs::remove_file(file).expect("failed to remove file");
}
}
}