use clap::Parser;
use futures::{stream, StreamExt};
use indicatif::ProgressBar;
use indicatif::ProgressStyle;
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::Arc;
use std::{collections::HashMap, fs};
use tokio::sync::mpsc::{channel, Sender};
use trust_dns_resolver::{config::*, TokioAsyncResolver};

enum OutputMessage {
    UrlResult(String),
    SubDomainResult(String),
    Shutdown,
}

/// Subdomain Bruteforcer, and directory bruteforcer written in rust, and multi threaded!!!
#[derive(Parser, Debug)]
#[command(version, about, long_about = None)]
struct Args {
    /// target to attack, can be a URL or a dns name, you can specify more than one by comma separating the targets
    #[arg(short, long)]
    target: Option<String>,

    /// wordlist to use for subdomains
    #[arg(short, long, default_value_t = String::from("none"))]
    subwordlist: String,

    ///wordlist to use for directories
    #[arg(short, long, default_value_t = String::from("none"))]
    dirwordlist: String,

    ///output files to write, will create two files yourinput_urls.txt and yourinput_subs.txt
    #[arg(short, long)]
    outfile: Option<String>,

    ///input file full of urls formatted to one url per line.
    #[arg(short, long)]
    urlfile: Option<PathBuf>,

    ///input file full of domains formatted to one domain per line.
    #[arg(short, long)]
    indomainfile: Option<PathBuf>,
}

async fn try_sub(
    domain: String,
    wildcard_reses: &Arc<HashMap<&String, Vec<String>>>,
    tx: Sender<OutputMessage>,
    output: bool,
    resolver: Arc<TokioAsyncResolver>,
    pb: Arc<ProgressBar>,
) {
    let mut ips = Vec::new();
    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).await {
        ipv6s.iter().for_each(|ip| ips.push(ip.0.to_string()));
    }
    let mut wild_card = None;
    for target in wildcard_reses.keys() {
        if domain.contains(&target.to_string()) {
            wild_card = wildcard_reses.get(target);
        }
    }
    if !ips.is_empty() {
        if let Some(wild_ips) = wild_card {
            if ips.len() == wild_ips.len() {
                let mut wild_only = true;
                for ip in &ips {
                    if !wild_ips.contains(ip) {
                        wild_only = false;
                    }
                }
                if wild_only {
                    return;
                }
            }
        }
        if output {
            if let Err(e) = tx
                .send(OutputMessage::SubDomainResult(format!(
                    "{}: {}",
                    &domain,
                    ips.join(", ")
                )))
                .await
            {
                eprintln!("error sending output! {}", e);
            }
        }
        pb.println(&format!("{} {}", domain, ips.join(", ")));
    }
    pb.inc(1);
}

async fn try_dir(url: String, tx: Sender<OutputMessage>, output: bool, pb: Arc<ProgressBar>) {
    let resp_stat = reqwest::get(&url).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)))
                        .await
                    {
                        eprintln!("error sending output! {}", e);
                    }
                }
                pb.println(&format!("{} {}", resp, &url));
            }
            StatusCode::ACCEPTED => {
                if output {
                    if let Err(e) = tx
                        .send(OutputMessage::UrlResult(format!("{} {}", resp, url)))
                        .await
                    {
                        eprintln!("error sending output! {}", e);
                    }
                }
                pb.println(&format!("{} {}", resp, url));
            }
            StatusCode::CONTINUE => {
                if output {
                    if let Err(e) = tx
                        .send(OutputMessage::UrlResult(format!("{} {}", resp, url)))
                        .await
                    {
                        eprintln!("error sending output! {}", e);
                    }
                }
                pb.println(&format!("{} {}", resp, url));
            }
            StatusCode::CREATED => {
                if output {
                    if let Err(e) = tx
                        .send(OutputMessage::UrlResult(format!("{} {}", resp, url)))
                        .await
                    {
                        eprintln!("error sending output! {}", e);
                    }
                }
                pb.println(&format!("{} {}", resp, url));
            }
            StatusCode::FOUND => {
                if output {
                    if let Err(e) = tx
                        .send(OutputMessage::UrlResult(format!("{} {}", resp, url)))
                        .await
                    {
                        eprintln!("error sending output! {}", e);
                    }
                }
                pb.println(&format!("{} {}", resp, url));
            }
            StatusCode::IM_A_TEAPOT => {
                if output {
                    if let Err(e) = tx
                        .send(OutputMessage::UrlResult(format!("{} {}", resp, url)))
                        .await
                    {
                        eprintln!("error sending output! {}", e);
                    }
                }
                pb.println(&format!(
                    "{} {} - and what a beautiful teapot you are!",
                    resp, url
                ));
            }
            StatusCode::MOVED_PERMANENTLY => {
                if output {
                    if let Err(e) = tx
                        .send(OutputMessage::UrlResult(format!("{} {}", resp, url)))
                        .await
                    {
                        eprintln!("error sending output! {}", e);
                    }
                }
                pb.println(&format!("{} {}", resp, url));
            }
            StatusCode::PERMANENT_REDIRECT => {
                if output {
                    if let Err(e) = tx
                        .send(OutputMessage::UrlResult(format!("{} {}", resp, url)))
                        .await
                    {
                        eprintln!("error sending output! {}", e);
                    }
                }
                pb.println(&format!("{} {}", resp, url));
            }
            StatusCode::TEMPORARY_REDIRECT => {
                if output {
                    if let Err(e) = tx
                        .send(OutputMessage::UrlResult(format!("{} {}", resp, url)))
                        .await
                    {
                        eprintln!("error sending output! {}", e);
                    }
                }
                pb.println(&format!("{} {}", resp, url));
            }
            _ => (),
        }
    }
    pb.inc(1);
}

#[tokio::main]
async fn main() {
    print!(
        "




                                                           ▄

                                                        ╓╦▀

                                                      ╓█▀

                                                     ▀▀



                                                              ╓▓

                                            ╓▄              ▄█▀

                                          ▄█▀▓▌           ▄█▀            ▄▀

                                        ▄█╙   █          ▀╙           ╓█▀

                                      ▄▀      █▄   ▄▄▄              ╓█▀

                                     ▐▌       ▐█ ▄▀└ ▐█             └

                                     ▐▌       ╓█▀   ▄█┘

                                     ╟▌     ▄█▀   ▄█▌

                                  ┌▄▄██╖  ▄█▀  ╓▄▀  └╙▀▀▀▀▓▄▄

                              ▄▄▓▀╙    ▀██╙  ╓▓▀           ▄▀

                          ╓▄▓▀└          ╙█▄█▀          ╓█▀

                       ▄██▌                 ▀▄        ▄█▀

                    ▄█▀   ▀▌                 ▐█▀▀▀▀▀▀▀╘

                 ▄█▀ ▀▌    ╙█▄              ╓█

               ▄▀╙    ╙█┐    ▀▓╖           ▄█

            ╓▓▀         ▀▌     ▀█▄        ▄█

          ▄█▀             ▀▄     ╙▀▓▄    ▄▀

        ╓█▀                 ▀█╖     ╙▀▓▄█▀

       ██                     ╙▀▄╓    ▄▌

      █╛▀█                       ╙▀▓▄█▀

     ▓▌   █▄                       ▄█

     █░    ▀█┐                   ▄█╘

     █▌      ▀█╖               ▄█╙

      █─       ▀█▄           ▄█╙

      ▄█▄         ▀█▄     ┌▄▀└

     █▌ ▀█▄          ▀▀▄▄█▀

     ╙█▄▄▄▀▀▀▓▄▄▄▄▄▄▄▓▀▀


"
    );
    let mut dirs_to_try = Vec::new();
    let mut subs_to_try = Vec::new();
    let args = Args::parse();
    let mut targets = Vec::new();
    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());
        }
    }
    if let Some(url_file) = args.urlfile {
        if let Ok(url_string) = fs::read_to_string(url_file) {
            url_string.lines().into_iter().for_each(|url| {
                targets.push(url.to_string());
            });
        }
    }
    if let Some(domain_file) = args.indomainfile {
        if let Ok(domain_string) = fs::read_to_string(domain_file) {
            domain_string.lines().into_iter().for_each(|domain| {
                targets.push(domain.to_string());
            });
        }
    }
    if targets.is_empty() {
        println!("no targets provided!");
        exit(1);
    }
    println!("detecting if a wild card record exists...");
    let mut opts = ResolverOpts::default();
    opts.edns0 = true;
    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))
                .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))
                .await
            {
                println!("wildcard found!");
                ipv6s.iter().for_each(|ip| {
                    ips.push(ip.0.to_string());
                });
            }
            if !ips.is_empty() {
                wild_card_results.insert(target, ips);
            }
        } else {
            dirs_to_try.push(target.clone());
        }
    }
    if args.dirwordlist != String::from("none") {
        let dirwordlist =
            fs::read_to_string(args.dirwordlist).expect("error reading directory wordlist");
        for dir in dirwordlist.split("\n").collect::<Vec<&str>>() {
            for target in &targets {
                if target.contains("http") {
                    if dir.starts_with("/") || target.ends_with("/") {
                        dirs_to_try.push(format!("{}{}", target, dir.trim()));
                    } else {
                        dirs_to_try.push(format!("{}/{}", target, dir.trim()));
                    }
                }
            }
        }
    }
    if args.subwordlist != String::from("none") {
        let subwordlist =
            fs::read_to_string(args.subwordlist).expect("error reading subdomain word list");
        for sub in subwordlist.split("\n").collect::<Vec<&str>>() {
            for target in &targets {
                if !target.contains("http") {
                    subs_to_try.push(format!("{}.{}", sub.trim(), target));
                }
            }
        }
    }
    println!("INFO:::loading targets, and wordlists...:::");
    let urls_len = dirs_to_try.len() as u64;
    let domains_len = subs_to_try.len() as u64;
    println!(
        "INFO:::DONE!, {} URLS to try, and {} subdomains to try:::",
        &urls_len, &domains_len
    );
    let mut output = false;
    let mut urlpath = PathBuf::new();
    let mut subpath = PathBuf::new();
    if args.outfile.is_some() {
        output = true;
        let path = args.outfile.clone().unwrap();
        urlpath = PathBuf::from(format!("{}_urls.txt", &path));
        subpath = PathBuf::from(format!("{}_subs.txt", &path));
        println!(
            "INFO:::results will be saved to {}_urls.txt and {}_subs.txt:::",
            path, path
        );
    }
    let write_handler = tokio::spawn(async move {
        if output {
            let urlfile = File::create(urlpath).unwrap();
            let subfile = File::create(subpath).unwrap();
            let mut urlwrite = BufWriter::new(urlfile);
            let mut subwrite = BufWriter::new(subfile);
            const FLUSHEVERY: usize = 10;
            let mut url_counter = 0;
            let mut sub_counter = 0;
            loop {
                let msg = rx.recv().await.unwrap();
                match msg {
                    OutputMessage::UrlResult(res) => {
                        if let Err(e) = writeln!(urlwrite, "{}", res) {
                            println!("error writing url file! {}", e);
                        } else {
                            url_counter += 1;
                        }
                    }
                    OutputMessage::SubDomainResult(res) => {
                        if let Err(e) = writeln!(subwrite, "{}", res) {
                            println!("error writing sub file! {}", e);
                        } else {
                            sub_counter += 1;
                        }
                    }
                    OutputMessage::Shutdown => {
                        urlwrite.flush().unwrap();
                        subwrite.flush().unwrap();
                        break;
                    }
                }
                if url_counter == FLUSHEVERY {
                    if let Err(e) = urlwrite.flush() {
                        println!("error writing url file! {e}");
                    }
                }
                if sub_counter == FLUSHEVERY {
                    if let Err(e) = subwrite.flush() {
                        println!("error writing sub file! {e}");
                    }
                }
            }
        } else {
            println!("INFO:::results will be printed to console and not saved...:::");
        }
    });
    let total = urls_len + domains_len;
    let pb = Arc::new(ProgressBar::new(total));
    pb.set_style(
        ProgressStyle::with_template("[{elapsed_precise}] {bar:40.orange/red} {pos}/{len} ({eta})")
            .unwrap(),
    );
    let resolver = TokioAsyncResolver::tokio(ResolverConfig::default(), opts);
    let resolver = Arc::new(resolver);
    pb.println("starting...");
    pb.force_draw();
    let dir_stream = stream::iter(dirs_to_try.into_iter().map(|url| {
        let pb = pb.clone();
        let tx = tx.clone();
        let output = output.clone();
        async move {
            try_dir(url, tx, output, pb).await;
        }
    }));
    let wc_arc = Arc::new(wild_card_results);
    let sub_stream = stream::iter(subs_to_try.into_iter().map(|sub| {
        let tx = tx.clone();
        let output = output.clone();
        let wc = &wc_arc;
        let resolver = resolver.clone();
        let pb = pb.clone();
        async move {
            try_sub(sub, wc, tx, output, resolver.clone(), pb).await;
        }
    }));

    let concurrency = 1000;
    dir_stream
        .buffer_unordered(concurrency)
        .collect::<Vec<_>>()
        .await;

    sub_stream
        .buffer_unordered(concurrency)
        .collect::<Vec<_>>()
        .await;
    pb.println("INFO:::enumeration finished!:::");
    if output {
        pb.println("INFO:::waiting for output writer to finish...:::");
        tx.send(OutputMessage::Shutdown).await.unwrap();
    }
    write_handler.await.unwrap();
    pb.println("INFO:::done bruteforcing, happy hunting!:::");
    pb.finish_with_message("Done!");
}
