passgen-rs 0.1.0

use clap::{Args, Parser};
use passgen::{random::BoxedRandomSource, *};
use std::{path::PathBuf, str::FromStr};

#[derive(Clone)]
pub struct WordlistSpec {
    name: String,
    path: PathBuf,
}

#[derive(thiserror::Error, Debug)]
pub enum WordlistSpecParseError {
    #[error("missing colon")]
    MissingColon,
}

impl FromStr for WordlistSpec {
    type Err = WordlistSpecParseError;

    fn from_str(input: &str) -> Result<Self, Self::Err> {
        let (name, path) = input
            .split_once(':')
            .ok_or(WordlistSpecParseError::MissingColon)?;
        Ok(Self {
            name: name.into(),
            path: path.into(),
        })
    }
}

#[derive(Parser)]
#[command(author, version)]
pub struct Options {
    #[clap(short, long)]
    pub preset: Option<String>,

    #[clap(short, long, default_value = "1")]
    pub amount: usize,

    #[clap(short, long)]
    pub random: Option<Random>,

    #[clap(flatten)]
    pub master: MasterOptions,

    #[clap(short, long, default_value = "text")]
    pub format: String,

    #[clap(short = 'j', long)]
    pub threads: Option<usize>,

    #[clap(short, long)]
    pub verbose: Option<usize>,

    #[clap(short, long)]
    pub wordlist: Vec<WordlistSpec>,

    #[clap(required_unless_present("preset"))]
    pub pattern: Option<String>,
}

/// Options for master-pass mode.
#[derive(Args)]
#[group()]
pub struct MasterOptions {
    #[clap(short, long)]
    pub master: Option<String>,

    #[clap(short, long)]
    pub domain: Option<String>,

    #[clap(short, long)]
    pub token: Option<String>,
}

impl Options {
    fn random(&self) -> Random {
        self.master
            .random()
            .or(self.random.clone())
            .unwrap_or_default()
    }
}

impl MasterOptions {
    fn random(&self) -> Option<Random> {
        self.master
            .as_deref()
            .map(|master| Random::from_master_pass(master, self.domain.as_deref().unwrap_or("")))
    }
}

fn main() {
    let options = Options::parse();
    let mut config = Config::default();

    for wordlist in &options.wordlist {
        config.wordlist_add(&wordlist.name, &wordlist.path);
    }

    // get randomness source
    let random = options.random();

    // parse pattern
    let pattern = Pattern::parse(&options.pattern.unwrap()).unwrap();

    // optimize pattern and get dependencies
    let (pattern, dependencies) = pattern.optimize();

    // resolve dependencies and load in data store
    let mut data = DataStore::default();
    for dependency in &dependencies {
        data.resolve(&config, dependency).unwrap();
    }

    let pattern = pattern.prepare(&data).unwrap();

    if let Some(threads) = options.threads {
        use std::sync::{atomic::*, *};
        let pattern = Arc::new(pattern);
        let counter = Arc::new(AtomicUsize::new(0));
        let amount = options.amount;

        let (sender, receiver) = mpsc::sync_channel(1024);
        let threads = (0..threads)
            .map(move |_index| {
                let random = random.clone();
                let counter = counter.clone();
                let pattern = pattern.clone();
                let sender = sender.clone();
                std::thread::spawn(move || {
                    loop {
                        let index = counter.fetch_add(1, Ordering::SeqCst);
                        if index >= options.amount {
                            break;
                        }

                        let mut rng = random.get_rng_boxed(index);
                        let output = pattern.generate(&mut rng).unwrap();
                        sender.send(output);
                    }
                    drop(sender);
                })
            })
            .collect::<Vec<_>>();

        while let Ok(message) = receiver.recv() {
            println!("{message}");
        }

        for thread in threads {
            thread.join().unwrap();
        }
    } else {
        for index in 0..options.amount {
            let mut rng = random.get_rng_boxed(index);
            let output = pattern.generate(&mut rng).unwrap();
            println!("{output}");
        }
    }
}