use anyhow::Result;
use crossterm::{cursor, terminal, ExecutableCommand};
use reqwest;
use reqwest::StatusCode;
use std::collections::HashMap;
use tokio::sync::mpsc;
use url::Url;
use std::{
fs::File,
io::{prelude::*, BufReader},
};
const DEFAULT: &str = "default";
fn load_urls() -> (HashMap<String, Vec<String>>, usize) {
let mut url_groups: HashMap<String, Vec<String>> = HashMap::new();
let mut total_urls = 0;
let mut args = std::env::args().skip(1);
let source: String = if args.len() == 0 {
std::env::var("HASWEBS_SOURCE")
.expect("if not passing file or urls as arg, need to pass as env var HASWEBS_SOURCE")
} else {
args.next().expect("should have had")
};
if source.contains(",") {
let mut current_group = String::from(DEFAULT);
for line in source.split(",") {
if Url::parse(&line).is_ok() {
url_groups
.entry(current_group.clone())
.and_modify(|urls| urls.push(line.to_string()))
.or_insert(vec![line.to_string()]);
total_urls += 1;
} else if line != "" {
current_group = line.to_string();
}
}
return (url_groups, total_urls);
} else if Url::parse(&source).is_ok() {
url_groups.insert(String::from("default"), vec![source]);
total_urls += 1;
return (url_groups, total_urls);
}
let mut current_group = String::from(DEFAULT);
for line in BufReader::new(
File::open(source)
.expect("expected arg or env var to be a file since it was not a url or group of urls"),
)
.lines()
{
let line = line.unwrap();
if Url::parse(&line).is_ok() {
url_groups
.entry(current_group.clone())
.and_modify(|urls| urls.push(line.to_string()))
.or_insert(vec![line.to_string()]);
total_urls += 1;
} else if line != "" {
current_group = line.to_string();
}
}
(url_groups, total_urls)
}
#[derive(Debug)]
enum Outcome {
Ok(String),
NotOk(String, String, String),
}
fn print_progress(groups: &Vec<(String, usize)>, outcomes: &HashMap<String, (usize, Vec<String>)>) {
let mut msg = String::from("");
for (group, group_url_count) in groups {
if let Some((successes, errors)) = outcomes.get(group) {
let num_errors = errors.len();
msg = msg
+ group
+ ": "
+ &(successes + num_errors).to_string()
+ " of "
+ &group_url_count.to_string();
if num_errors > 0 {
msg = msg + " (" + &num_errors.to_string() + " failed)";
for err in errors {
msg = msg + "\n " + err;
}
}
msg += "\n";
};
}
let mut stdo = std::io::stdout();
stdo.execute(terminal::Clear(terminal::ClearType::FromCursorDown))
.unwrap();
stdo.execute(cursor::SavePosition).unwrap();
print!("\r{}", msg);
stdo.execute(cursor::RestorePosition).unwrap();
}
#[tokio::main]
async fn main() -> Result<()> {
let (url_groups, total_urls) = load_urls();
println!("checking {} urls", total_urls);
let ordered_groups: Vec<_> = url_groups
.iter()
.map(|(group, urls)| (group.clone(), urls.len()))
.collect();
let mut outcomes: HashMap<String, (usize, Vec<String>)> = HashMap::new();
print_progress(&ordered_groups, &outcomes);
let (tx, mut rx) = mpsc::channel(std::cmp::min(100, total_urls));
for (group, urls) in url_groups {
for url in urls {
let tx = tx.clone();
let group = group.clone();
tokio::spawn(async move {
match reqwest::get(&url).await {
Ok(resp) => {
if resp.status() == StatusCode::OK {
let send_value: Outcome = Outcome::Ok(group);
tx.send(send_value).await.unwrap();
} else {
tx.send(Outcome::NotOk(group, url, resp.status().to_string()))
.await
.unwrap();
}
}
Err(x) => {
tx.send(Outcome::NotOk(group, url, x.to_string()))
.await
.unwrap();
}
};
});
}
}
let mut total_errors = 0;
for _ in 0..total_urls {
match rx.recv().await.unwrap() {
Outcome::Ok(group) => {
outcomes
.entry(group)
.and_modify(|(successes, _)| *successes += 1)
.or_insert((1, vec![]));
}
Outcome::NotOk(group, url, error_message) => {
let descriptive_error = format!("{} ({})", url, error_message);
total_errors += 1;
outcomes
.entry(group)
.and_modify(|(_, url_errors)| url_errors.push(descriptive_error.clone()))
.or_insert((0, vec![descriptive_error]));
}
};
print_progress(&ordered_groups, &outcomes);
}
let mut stdo = std::io::stdout();
stdo.execute(cursor::MoveDown(
(ordered_groups.len() + total_errors) as u16,
))
.unwrap();
Ok(())
}