pyro/rustbuster
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases 5 Wiki Activity

Switched from rayon for multithreading to tokio for async.

Browse Source
This commit is contained in:
pyro
2026-05-15 15:43:18 -05:00
parent 050238f46a
commit e46f987db2
3 changed files with 184 additions and 228 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/main.rs
+97 -70
View File
@@ -1,17 +1,15 @@
use clap::Parser;
use rayon::iter::{IntoParallelRefIterator, ParallelIterator};
use rayon::ThreadPoolBuilder;
use futures::{stream, StreamExt};
use reqwest::{self, StatusCode};
use std::fs::File;
use std::io::BufWriter;
use std::io::Write;
use std::path::PathBuf;
use std::process::exit;
use std::sync::mpsc::{channel, Sender};
use std::thread;
use std::sync::Arc;
use std::{collections::HashMap, fs};
use trust_dns_resolver::config::*;
use trust_dns_resolver::Resolver;
use tokio::sync::mpsc::{channel, Sender};
use trust_dns_resolver::{config::*, TokioAsyncResolver};
use urlencoding::encode;
enum OutputMessage {
@@ -47,26 +45,20 @@ struct Args {
///input file full of domains formatted to one domain per line.
#[arg(short, long)]
domainfile: Option<PathBuf>,
///number of threads to use by default it will use the rayon global pool default.
#[arg(long)]
threads: Option<usize>,
}
fn try_sub(
async fn try_sub(
domain: String,
wildcard_reses: HashMap<&String, Vec<String>>,
tx: Sender<OutputMessage>,
output: bool,
resolver: Arc<TokioAsyncResolver>,
) {
let mut opts = ResolverOpts::default();
opts.edns0 = true;
let resolver = Resolver::new(ResolverConfig::default(), opts).unwrap();
let mut ips = Vec::new();
if let Ok(ipv4s) = resolver.ipv4_lookup(&domain) {
if let Ok(ipv4s) = resolver.ipv4_lookup(&domain).await {
ipv4s.iter().for_each(|ip| ips.push(ip.0.to_string()));
}
if let Ok(ipv6s) = resolver.ipv6_lookup(&domain) {
if let Ok(ipv6s) = resolver.ipv6_lookup(&domain).await {
ipv6s.iter().for_each(|ip| ips.push(ip.0.to_string()));
}
let mut wild_card = None;
@@ -90,11 +82,14 @@ fn try_sub(
}
}
if output {
if let Err(e) = tx.send(OutputMessage::SubDomainResult(format!(
"{}: {}",
&domain,
ips.join(", ")
))) {
if let Err(e) = tx
.send(OutputMessage::SubDomainResult(format!(
"{}: {}",
&domain,
ips.join(", ")
)))
.await
{
eprintln!("error sending output! {}", e);
}
}
@@ -102,14 +97,17 @@ fn try_sub(
}
}
fn try_dir(url: String, tx: Sender<OutputMessage>, output: bool) {
let resp_stat = reqwest::blocking::get(&encode(&url).to_string());
async fn try_dir(url: String, tx: Sender<OutputMessage>, output: bool) {
let resp_stat = reqwest::get(&encode(&url).to_string()).await;
if resp_stat.is_ok() {
let resp = resp_stat.unwrap().status();
match resp {
StatusCode::OK => {
if output {
if let Err(e) = tx.send(OutputMessage::UrlResult(format!("{} {}", resp, url))) {
if let Err(e) = tx
.send(OutputMessage::UrlResult(format!("{} {}", resp, url)))
.await
{
eprintln!("error sending output! {}", e);
}
}
@@ -117,7 +115,10 @@ fn try_dir(url: String, tx: Sender<OutputMessage>, output: bool) {
}
StatusCode::ACCEPTED => {
if output {
if let Err(e) = tx.send(OutputMessage::UrlResult(format!("{} {}", resp, url))) {
if let Err(e) = tx
.send(OutputMessage::UrlResult(format!("{} {}", resp, url)))
.await
{
eprintln!("error sending output! {}", e);
}
}
@@ -125,7 +126,10 @@ fn try_dir(url: String, tx: Sender<OutputMessage>, output: bool) {
}
StatusCode::CONTINUE => {
if output {
if let Err(e) = tx.send(OutputMessage::UrlResult(format!("{} {}", resp, url))) {
if let Err(e) = tx
.send(OutputMessage::UrlResult(format!("{} {}", resp, url)))
.await
{
eprintln!("error sending output! {}", e);
}
}
@@ -133,7 +137,10 @@ fn try_dir(url: String, tx: Sender<OutputMessage>, output: bool) {
}
StatusCode::CREATED => {
if output {
if let Err(e) = tx.send(OutputMessage::UrlResult(format!("{} {}", resp, url))) {
if let Err(e) = tx
.send(OutputMessage::UrlResult(format!("{} {}", resp, url)))
.await
{
eprintln!("error sending output! {}", e);
}
}
@@ -141,7 +148,10 @@ fn try_dir(url: String, tx: Sender<OutputMessage>, output: bool) {
}
StatusCode::FOUND => {
if output {
if let Err(e) = tx.send(OutputMessage::UrlResult(format!("{} {}", resp, url))) {
if let Err(e) = tx
.send(OutputMessage::UrlResult(format!("{} {}", resp, url)))
.await
{
eprintln!("error sending output! {}", e);
}
}
@@ -149,7 +159,10 @@ fn try_dir(url: String, tx: Sender<OutputMessage>, output: bool) {
}
StatusCode::IM_A_TEAPOT => {
if output {
if let Err(e) = tx.send(OutputMessage::UrlResult(format!("{} {}", resp, url))) {
if let Err(e) = tx
.send(OutputMessage::UrlResult(format!("{} {}", resp, url)))
.await
{
eprintln!("error sending output! {}", e);
}
}
@@ -157,7 +170,10 @@ fn try_dir(url: String, tx: Sender<OutputMessage>, output: bool) {
}
StatusCode::MOVED_PERMANENTLY => {
if output {
if let Err(e) = tx.send(OutputMessage::UrlResult(format!("{} {}", resp, url))) {
if let Err(e) = tx
.send(OutputMessage::UrlResult(format!("{} {}", resp, url)))
.await
{
eprintln!("error sending output! {}", e);
}
}
@@ -165,7 +181,10 @@ fn try_dir(url: String, tx: Sender<OutputMessage>, output: bool) {
}
StatusCode::PERMANENT_REDIRECT => {
if output {
if let Err(e) = tx.send(OutputMessage::UrlResult(format!("{} {}", resp, url))) {
if let Err(e) = tx
.send(OutputMessage::UrlResult(format!("{} {}", resp, url)))
.await
{
eprintln!("error sending output! {}", e);
}
}
@@ -173,7 +192,10 @@ fn try_dir(url: String, tx: Sender<OutputMessage>, output: bool) {
}
StatusCode::TEMPORARY_REDIRECT => {
if output {
if let Err(e) = tx.send(OutputMessage::UrlResult(format!("{} {}", resp, url))) {
if let Err(e) = tx
.send(OutputMessage::UrlResult(format!("{} {}", resp, url)))
.await
{
eprintln!("error sending output! {}", e);
}
}
@@ -184,7 +206,8 @@ fn try_dir(url: String, tx: Sender<OutputMessage>, output: bool) {
}
}
fn main() {
#[tokio::main]
async fn main() {
print!(
"
@@ -262,7 +285,7 @@ fn main() {
let mut subs_to_try = Vec::new();
let args = Args::parse();
let mut targets = Vec::new();
let (tx, rx) = channel();
let (tx, mut rx) = channel(1024);
if let Some(given_targets) = args.target {
for target in given_targets.split(",").collect::<Vec<&str>>() {
targets.push(target.to_owned());
@@ -286,32 +309,28 @@ fn main() {
println!("no targets provided!");
exit(1);
}
if let Some(thread) = args.threads {
ThreadPoolBuilder::new()
.num_threads(thread)
.build_global()
.unwrap();
}
println!(
"INFO:::{} threads will be used:::",
rayon::current_num_threads()
);
println!("detecting if a wild card record exists...");
let mut opts = ResolverOpts::default();
opts.edns0 = true;
let resolver = Resolver::new(ResolverConfig::default(), opts).unwrap();
let resolver = TokioAsyncResolver::tokio(ResolverConfig::default(), opts);
let mut wild_card_results = HashMap::new();
for target in &targets {
if !target.contains("http") {
subs_to_try.push(target.clone());
let mut ips = Vec::new();
if let Ok(ipv4s) = resolver.ipv4_lookup(&format!("burstpyrofoo.{}", target)) {
if let Ok(ipv4s) = resolver
.ipv4_lookup(&format!("burstpyrofoo.{}", target))
.await
{
ipv4s.iter().for_each(|ip| {
ips.push(ip.0.to_string());
});
println!("INFO:::wildcard found: {} {}:::", target, ips.join(", "));
}
if let Ok(ipv6s) = resolver.ipv6_lookup(&format!("burstpyrofoo.{}", target)) {
if let Ok(ipv6s) = resolver
.ipv6_lookup(&format!("burstpyrofoo.{}", target))
.await
{
println!("wildcard found!");
ipv6s.iter().for_each(|ip| {
ips.push(ip.0.to_string());
@@ -366,7 +385,7 @@ fn main() {
path, path
);
}
let write_handler = thread::spawn(move || {
let write_handler = tokio::spawn(async move {
if output {
let urlfile = File::create(urlpath).unwrap();
let subfile = File::create(subpath).unwrap();
@@ -376,7 +395,7 @@ fn main() {
let mut url_counter = 0;
let mut sub_counter = 0;
loop {
let msg = rx.recv().unwrap();
let msg = rx.recv().await.unwrap();
match msg {
OutputMessage::UrlResult(res) => {
if let Err(e) = writeln!(urlwrite, "{}", res) {
@@ -413,32 +432,40 @@ fn main() {
println!("INFO:::results will be printed to console and not saved...:::");
}
});
rayon::scope(|s| {
if !dirs_to_try.is_empty() {
s.spawn(|_| {
dirs_to_try.par_iter().for_each(|url| {
try_dir(url.to_string(), tx.clone(), output);
});
});
let resolver = TokioAsyncResolver::tokio(ResolverConfig::default(), opts);
let resolver = Arc::new(resolver);
let dir_stream = stream::iter(dirs_to_try.into_iter().map(|url| {
let tx = tx.clone();
let output = output.clone();
async move {
try_dir(url, tx, output).await;
}
if !subs_to_try.is_empty() {
s.spawn(|_| {
subs_to_try.par_iter().for_each(|sub| {
try_sub(
sub.to_string(),
wild_card_results.clone(),
tx.clone(),
output,
);
});
})
}));
let sub_stream = stream::iter(subs_to_try.into_iter().map(|sub| {
let tx = tx.clone();
let output = output.clone();
let wc = wild_card_results.clone();
let resolver = resolver.clone();
async move {
try_sub(sub, wc, tx, output, resolver.clone()).await;
}
});
}));
let concurrency = 100;
dir_stream
.buffer_unordered(concurrency)
.collect::<Vec<_>>()
.await;
sub_stream
.buffer_unordered(concurrency)
.collect::<Vec<_>>()
.await;
println!("INFO:::enumeration finished!:::");
if output {
println!("INFO:::waiting for output writer to finish...:::");
tx.send(OutputMessage::Shutdown).unwrap();
tx.send(OutputMessage::Shutdown).await.unwrap();
}
write_handler.join().unwrap();
write_handler.await.unwrap();
println!("INFO:::done bruteforcing, happy hunting!:::");
}