rustbust/
request.rs

1use reqwest::Client;
2use std::{
3    error::Error,
4    io::{stdout, BufRead, BufWriter, Lines, Write},
5    sync::{mpsc::channel, Arc},
6    thread::{sleep, spawn},
7    time::Duration,
8};
9
10pub struct Response {
11    pub name: String,
12    pub status: u16,
13}
14
15pub fn request_n<T, U>(
16    lines: &mut Lines<T>,
17    n: usize,
18    url: &str,
19    outfile: &mut BufWriter<U>,
20) -> Result<bool, Box<Error>>
21where
22    T: BufRead,
23    U: Write,
24{
25    let cli = Arc::new(Client::new());
26    let (tx, rx) = channel();
27    let lines = lines.take(n).collect::<Result<Vec<String>, _>>()?;
28    let len = lines.len();
29
30    print!("{}[2K", 27 as char);
31    print!("fuzzing {}..{}\r", &lines[0], &lines[len - 1]);
32    stdout().flush()?;
33
34    for line in lines {
35        let fuzz_url = format!("{}/{}", url, &line);
36        let cli = Arc::clone(&cli);
37        let tx = tx.clone();
38
39        spawn(move || {
40            let mut counter = 0;
41
42            loop {
43                match cli.get(&fuzz_url).send() {
44                    Ok(res) => {
45                        let s = res.status().as_u16();
46
47                        if s != 404 {
48                            println!("\n{} returned {}", &line, s);
49
50                            tx.send(Some(Response {
51                                status: s,
52                                name: line,
53                            }))
54                            .unwrap();
55                        } else {
56                            tx.send(None).unwrap();
57                        }
58
59                        break;
60                    }
61
62                    Err(e) => {
63                        counter += 1;
64
65                        println!(
66                            "\nerror encountered while trying to fuzz {}: {}\nretrying in 1s",
67                            &line, e
68                        );
69
70                        if counter > 5 {
71                            println!(
72                                "\ntoo many errors encountered trying to fuzz {}, skipped",
73                                line
74                            );
75
76                            tx.send(Some(Response {
77                                name: line,
78                                status: 500,
79                            }))
80                            .unwrap();
81
82                            break;
83                        }
84
85                        sleep(Duration::from_secs(1));
86                    }
87                }
88            }
89        });
90    }
91
92    for _ in 0..len {
93        if let Some(res) = rx.recv()? {
94            outfile.write(&format!("{}\t{}\n", res.name, res.status).into_bytes())?;
95        }
96    }
97
98    outfile.flush()?;
99
100    Ok(n == len)
101}