pwtool 0.4.0

use base64::prelude::*;
use hmac::{Hmac, Mac};
use pbkdf2::pbkdf2_hmac;
use rand::distributions::Alphanumeric;
use rand::thread_rng;
use rand::Rng;
use sha1::Digest;
use sha2::Sha256;
use std::fs::File;
use std::io::BufRead;
use std::io::BufReader;

pub struct Config {
    pub number: u32,
    pub len: u32,
    pub pw_type: Option<u32>,
    pub word_list: Option<String>,
    pub words: Option<Vec<String>>,
    pub username: Option<String>,
    pub database: Option<String>,
}

impl Config {
    pub fn new() -> Config {
        Config {
            len: 15,
            pw_type: None,
            number: 20,
            word_list: None,
            words: None,
            username: None,
            database: None,
        }
    }
}

impl Default for Config {
    fn default() -> Self {
        Self::new()
    }
}

pub enum PwClass {
    Num = 1 << 0,
    Alpha = 1 << 1,
    Ext = 1 << 2,
    Lower = 1 << 3,
    Upper = 1 << 4,
}

pub fn valid_word(s: &str, word_set: Option<u32>) -> bool {
    if word_set.is_some() {
        let mut valid = false;

        if word_set.as_ref().unwrap() & PwClass::Num as u32 != 0 {
            for j in s.chars() {
                if j.is_ascii_digit() {
                    valid = true;
                }
            }
        }

        if word_set.as_ref().unwrap() & PwClass::Lower as u32 != 0 {
            for j in s.chars() {
                if j.is_ascii_lowercase() {
                    valid = true;
                }
            }
        }

        if word_set.as_ref().unwrap() & PwClass::Upper as u32 != 0 {
            for j in s.chars() {
                if j.is_ascii_uppercase() {
                    valid = true;
                }
            }
        }

        if word_set.as_ref().unwrap() & PwClass::Ext as u32 != 0 {
            for j in s.chars() {
                if !j.is_ascii_digit() && !j.is_ascii_lowercase() && !j.is_ascii_uppercase() {
                    valid = true;
                }
            }
        }
        return valid;
    };

    true
}

pub fn prng_string(c: &mut Config) -> String {
    let mut set = c.pw_type;
    if set.is_none() {
        set = Some(
            PwClass::Num as u32
                | PwClass::Alpha as u32
                | PwClass::Lower as u32
                | PwClass::Upper as u32,
        );
    }

    let mut rng = rand::thread_rng();

    if c.word_list.is_some() {
        let word_list = c.word_list.as_ref().unwrap();
        if c.words.is_none() {
            let f = File::open(word_list);
            if f.is_err() {
                eprintln!("Cannot open {}: {}", word_list, f.err().unwrap());
                std::process::exit(1);
            }

            let f = f.unwrap();

            let mut line = String::new();
            let mut br = BufReader::new(f);
            let mut wv: Vec<String> = vec![];
            loop {
                line.clear();
                let l = br.read_line(&mut line);
                match l {
                    Ok(i) => {
                        if i == 0 {
                            break;
                        }
                    }
                    Err(_) => {
                        break;
                    }
                }

                // skip words with apostrophes etc

                if line.find('\'').is_some() {
                    continue;
                }

                let s = line.clone().trim().to_string();
                if s.is_empty() {
                    continue;
                }

                if !valid_word(&s, c.pw_type) {
                    continue;
                }

                wv.push(s);
            }
            c.words = Some(wv);
        }

        let mut phrase = vec![];

        let words_len = c.words.as_ref().unwrap().len();
        if c.words.as_ref().unwrap().is_empty() {
            eprintln!("no words to process");
            std::process::exit(1);
        }

        for _j in 0..c.len {
            phrase.push(
                c.words.as_ref().unwrap()[rng.gen_range(0..words_len)]
                    .clone()
                    .to_string(),
            );
        }

        return phrase.join(" ");
    }

    let set = set.unwrap();
    let mut chars = "".to_string();
    if set & PwClass::Num as u32 != 0 {
        chars += "0123456789";
    };

    if set & PwClass::Alpha as u32 != 0 {
        chars += "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ";
    };

    if set & PwClass::Lower as u32 != 0 {
        chars += "abcdefghijklmnopqrstuvwxyz";
    };

    if set & PwClass::Upper as u32 != 0 {
        chars += "ABCDEFGHIJKLMNOPQRSTUVWXYZ";
    };

    if set & PwClass::Ext as u32 != 0 {
        chars += r#"!"$%^&*()-={}[]:;@'<>,./\|"#;
    };

    let one_char = || chars.chars().nth(rng.gen_range(0..chars.len())).unwrap();
    std::iter::repeat_with(one_char)
        .take(c.len as usize)
        .collect()
}

pub fn gen_salt_str(chars: usize) -> String {
    let rng = thread_rng();
    rng.sample_iter(&Alphanumeric)
        .take(chars)
        .map(|x| x as char)
        .collect()
}

pub fn postgres_pass(pw: &str) -> String {
    let salt_size = 16;
    let iterations = 4096;

    let salt = gen_salt_str(salt_size);
    let mut key = [0u8; 32];
    pbkdf2_hmac::<Sha256>(pw.as_bytes(), salt.as_bytes(), iterations, &mut key);

    let mut client_key = Hmac::<Sha256>::new_from_slice(&key).unwrap();
    client_key.update(b"Client Key");
    let client_result = client_key.finalize();

    let mut server_key = Hmac::<Sha256>::new_from_slice(&key).unwrap();
    server_key.update(b"Server Key");
    let server_result = server_key.finalize();

    let mut stored_key = Sha256::new();
    stored_key.update(client_result.clone().into_bytes());
    let stored_key = stored_key.finalize();

    let bstored_key = BASE64_STANDARD.encode(stored_key);
    let bsalt = BASE64_STANDARD.encode(salt);
    let bserver_key = BASE64_STANDARD.encode(server_result.clone().into_bytes());

    format!("SCRAM-SHA-256${iterations}:{bsalt}${bstored_key}:{bserver_key}")
}