rs-random/src/main.rs

use rand::{rngs::OsRng, seq::SliceRandom};
use std::collections::HashSet;
use std::env;
use std::process;

static CHARACTER_SETS: [(&str, &str, &str); 11] = [
    ("lowercase", "abcdefghijklmnopqrstuvwxyz", "English lowercase letters (a-z)"),
    ("uppercase", "ABCDEFGHIJKLMNOPQRSTUVWXYZ", "English uppercase letters (A-Z)"),
    ("numbers", "0123456789", "Numbers (0-9)"),
    ("hex", "0123456789abcdef", "Hexadecimal (0-9, a-f)"),
    ("special", "!@#$%^&*()_+-=[]{}|;:,.<>?", "Special characters"),
    ("special-safe", "!@#$%^&*_+-=<>?", "Safe special chars (no pipes/brackets)"),
    ("cyrillic-lower", "абвгдежзийклмнопрстуфхцчшщъыьэюя", "Cyrillic lowercase"),
    ("cyrillic-upper", "АБВГДЕЖЗИЙКЛМНОПРСТУФХЦЧШЩЪЫЬЭЮЯ", "Cyrillic uppercase"),
    ("greek-lower", "αβγδεζηθικλμνξοπρστυφχψω", "Greek lowercase"),
    ("greek-upper", "ΑΒΓΔΕΖΗΘΙΚΛΜΝΞΟΠΡΣΤΥΦΧΨΩ", "Greek uppercase"),
    ("symbols", "©®™€£¥§¶†‡•…‰′″‹›\"\'–—", "Extended symbols"),
];

#[inline(never)]
fn print_help() {
    println!("rs-random: Secure String Generator\n");
    println!("Usage:");
    println!("  rs-random [LENGTH]                   (uses safe defaults)");
    println!("  rs-random -l <LEN> [-s <SETS>] [-c <COUNT>]");
    println!("  rs-random -h\n");
    println!("Available sets (for -s, comma-separated):");
    for &(name, _, desc) in &CHARACTER_SETS {
        println!("  {:<15} - {}", name, desc);
    }
}

#[inline]
fn get_chars(selected_sets: &[&str]) -> Result<Vec<char>, String> {
    // Pre-allocate with estimated capacity
    let mut all_chars = Vec::with_capacity(selected_sets.len() * 32);
    let mut unknown = Vec::new();

    for &set_name in selected_sets {
        match CHARACTER_SETS.iter().find(|&&(name, _, _)| name == set_name) {
            Some(&(_, chars, _)) => all_chars.extend(chars.chars()),
            None => unknown.push(set_name),
        }
    }

    if !unknown.is_empty() {
        return Err(format!("Unknown sets: {}", unknown.join(",")));
    }
    if all_chars.is_empty() {
        return Err("Empty character pool".to_string());
    }

    // Optimize for common case where no duplicates exist
    if selected_sets.len() == 1 || 
       (selected_sets.len() <= 4 && !selected_sets.iter().any(|&s| 
           s == "numbers" && (selected_sets.contains(&"hex") || 
           selected_sets.contains(&"lowercase") || 
           selected_sets.contains(&"uppercase")))) {
        return Ok(all_chars);
    }

    // Remove duplicates only when necessary
    let mut unique_chars = Vec::with_capacity(all_chars.len());
    let mut seen = HashSet::with_capacity(all_chars.len());
    for ch in all_chars {
        if seen.insert(ch) {
            unique_chars.push(ch);
        }
    }
    unique_chars.shrink_to_fit();

    Ok(unique_chars)
}

#[inline]
fn generate_random_string(chars: &[char], length: usize) -> String {
    let mut rng = OsRng;
    let mut result = String::with_capacity(length);
    
    // Unroll loop for better performance on small strings
    let chunks = length / 4;
    let remainder = length % 4;
    
    for _ in 0..chunks {
        result.push(*chars.choose(&mut rng).unwrap());
        result.push(*chars.choose(&mut rng).unwrap());
        result.push(*chars.choose(&mut rng).unwrap());
        result.push(*chars.choose(&mut rng).unwrap());
    }
    
    for _ in 0..remainder {
        result.push(*chars.choose(&mut rng).unwrap());
    }
    
    result
}

fn main() {
    let args: Vec<String> = env::args().collect();

    if args.len() < 2 || args.contains(&"-h".to_string()) {
        print_help();
        return;
    }

    let mut length = 16;
    let mut count = 1;
    let mut sets = "lowercase,uppercase,numbers,special-safe";

    let mut i = 1;
    while i < args.len() {
        match args[i].as_str() {
            "-l" => {
                if i + 1 < args.len() {
                    match args[i + 1].parse::<usize>() {
                        Ok(n) => length = n,
                        Err(_) => {
                            eprintln!("Error: Invalid length");
                            process::exit(1);
                        }
                    }
                    i += 2;
                } else {
                    eprintln!("Error: Missing length value");
                    process::exit(1);
                }
            }
            "-s" => {
                if i + 1 < args.len() {
                    sets = &args[i + 1];
                    i += 2;
                } else {
                    eprintln!("Error: Missing sets value");
                    process::exit(1);
                }
            }
            "-c" => {
                if i + 1 < args.len() {
                    match args[i + 1].parse::<usize>() {
                        Ok(n) => count = n,
                        Err(_) => {
                            eprintln!("Error: Invalid count");
                            process::exit(1);
                        }
                    }
                    i += 2;
                } else {
                    eprintln!("Error: Missing count value");
                    process::exit(1);
                }
            }
            arg if i == 1 && !arg.starts_with('-') => {
                match arg.parse::<usize>() {
                    Ok(n) => length = n,
                    Err(_) => {
                        eprintln!("Error: Invalid length");
                        process::exit(1);
                    }
                }
                i += 1;
            }
            _ => {
                i += 1;
            }
        }
    }

    if length == 0 {
        eprintln!("Error: Length must be > 0");
        process::exit(1);
    }

    let selected_sets: Vec<&str> = sets.split(',').map(|s| s.trim()).filter(|s| !s.is_empty()).collect();
    
    match get_chars(&selected_sets) {
        Ok(chars) => {
            for i in 0..count {
                let random_string = generate_random_string(&chars, length);
                if count > 1 {
                    print!("{}: ", i + 1);
                }
                println!("{}", random_string);
            }
        }
        Err(e) => {
            eprintln!("Error: {}", e);
            process::exit(1);
        }
    }
}