duralumin/src/key_gen.rs

use russh_keys::PublicKeyBase64;
use sha2::Digest;
use zeroize::Zeroizing;

use crate::ed25519::randomart;

pub mod cli {
    use zeroize::Zeroizing;

    use crate::key_gen::{HashDesc, KeygenDesc};

    fn fix_newline_ref(line: &mut String) {
        if line.ends_with('\n') {
            line.pop();

            if line.ends_with('\r') {
                line.pop();
            }
        }
    }

    #[allow(dead_code)]
    fn fix_newline(mut line: String) -> String {
        fix_newline_ref(&mut line);

        line
    }
    pub fn read_line() -> std::io::Result<String> {
        let mut line = String::new();
        std::io::stdin().read_line(&mut line)?;
        fix_newline_ref(&mut line);

        Ok(line)
    }

    pub fn read_non_empty_line() -> std::io::Result<Option<String>> {
        let line = read_line()?;

        Ok(if line.is_empty() { None } else { Some(line) })
    }

    #[derive(thiserror::Error, Debug)]
    pub enum Error {
        #[error(transparent)]
        Io(#[from] std::io::Error),
        #[error(transparent)]
        ParseInt(#[from] std::num::ParseIntError),
        #[error("Mismatching secret.")]
        MismatchingSecret,
        #[error("Invalid hash method.")]
        InvalidHashMethod,
    }

    pub type Result<T> = std::result::Result<T, Error>;

    fn read_passphrase() -> Result<Zeroizing<String>> {
        let passphrase = Zeroizing::new(rpassword::prompt_password_stdout("Enter a passphrase: ")?);
        let passphrase2 =
            Zeroizing::new(rpassword::prompt_password_stdout("Enter a passphrase: ")?);

        if passphrase == passphrase2 {
            println!(
                "Passphrase entropy (rough estimate): {:.2}",
                crate::entropy(&passphrase)
            );
            Ok(passphrase)
        } else {
            Err(Error::MismatchingSecret)
        }
    }

    pub fn read_argon_desc() -> Result<HashDesc> {
        print!("Use argon2 variant (argon2i, argon2d, argon2id) [argon2id]: ");
        let variant = match read_non_empty_line()?.as_ref().map(|s| s.as_str()) {
            Some("argon2i") => argon2::Algorithm::Argon2i,
            Some("argon2d") => argon2::Algorithm::Argon2d,
            Some("argon2id") | None => argon2::Algorithm::Argon2id,
            _ => {
                return Err(Error::InvalidHashMethod);
            }
        };

        print!("Use argon2 version (16,19) [argon2id]: ");
        let version = match read_non_empty_line()?.as_ref().map(|s| s.as_str()) {
            Some("16") | Some("10") => argon2::Version::V0x10,
            Some("19") | Some("13") | None => argon2::Version::V0x13,
            _ => {
                return Err(Error::InvalidHashMethod);
            }
        };

        print!("memory size (memory cost) (KiB if no unit specified) (min 19 MiB) [64 MiB]: ");
        let memory = 'mem: {
            let Some(line) = read_non_empty_line()? else {
                break 'mem 64 * 1024;
            };
            let line = line.trim();
            let chars = line.char_indices();
            let digits_end = chars
                .take_while(|(_, c)| ('0'..='9').contains(c))
                .map(|(i, _)| i)
                .last()
                .unwrap_or(line.len());
            let digits = &line[..digits_end];
            let unit = &line[digits_end..];

            let number = digits.parse::<u32>()?;
            let factor = match unit {
                "KiB" | "kib" | "" => 1,
                "MiB" | "mib" => 1024,
                _ => {
                    return Err(Error::InvalidHashMethod);
                }
            };

            number * factor
        };

        print!("Iterations (time cost) [4]: ");
        let iterations = read_non_empty_line()?
            .map(|s| s.parse::<u16>())
            .unwrap_or(Ok(4))?;

        print!("Lanes (parallelism cost) [1]: ");
        let lanes = read_non_empty_line()?
            .map(|s| s.parse::<u16>())
            .unwrap_or(Ok(1))?;

        Ok(HashDesc::Argon {
            variant,
            version,
            memory,
            time: iterations as u32,
            lanes: lanes as u32,
            hash_length: 32,
        })
    }

    pub fn keygen_desc_from_stdin() -> Result<super::KeygenDesc> {
        let passphrase = read_passphrase()?;

        print!("Enter an optional ID []: ");
        let tag = read_non_empty_line()?.map(|s| Zeroizing::new(s));

        print!("Encrypt keypair with passphrase? [Y/n]: ");
        let encrypt = read_line()? == "Y";

        print!("Use hash algorithm (sha256, argon2) [argon2]: ");
        let hash = match read_non_empty_line()?.as_ref().map(|s| s.as_str()) {
            Some("argon2") | None => read_argon_desc()?,
            Some("sha256") => {
                print!("Iterations [4]: ");
                let iterations = read_non_empty_line()?
                    .map(|s| s.parse::<u16>())
                    .unwrap_or(Ok(4))?;
                HashDesc::Sha256 {
                    iterations: iterations as u32,
                }
            }
            _ => return Err(Error::InvalidHashMethod),
        };

        Ok(KeygenDesc {
            hash,
            salt: KeygenDesc::SALT.to_vec(),
            tag,
            passphrase,
            encrypt,
        })
    }
}

#[derive(Debug, thiserror::Error)]
pub enum Error {
    #[error("Argon2 Error: {0}")]
    Argon2(argon2::Error),
    #[error(transparent)]
    SshKeys(#[from] russh_keys::Error),
}

impl From<argon2::Error> for Error {
    fn from(value: argon2::Error) -> Self {
        Self::Argon2(value)
    }
}

pub type Result<T> = core::result::Result<T, Error>;

pub struct KeyPair {
    pub passphrase: Option<Zeroizing<String>>,
    pub private_key: ed25519_dalek::SigningKey,
    pub public_key: ed25519_dalek::VerifyingKey,
}

impl KeyPair {
    pub fn fingerprint(&self) -> Vec<u8> {
        sha2::Sha256::digest(self.public_key.as_bytes()).to_vec()
    }

    pub fn fingerprint_base64(&self) -> String {
        base64::encode(&self.fingerprint())
    }

    pub fn randomart(&self) -> randomart::RandomArt {
        randomart::RandomArt::from_digest(&self.fingerprint())
    }

    pub fn encode_keys(&self) -> Result<(Zeroizing<String>, String)> {
        let keypair = russh_keys::key::KeyPair::Ed25519(self.private_key.clone());
        let public_key = keypair.clone_public_key()?;

        let mut private_key = Vec::new();
        match &self.passphrase {
            Some(passphrase) => {
                russh_keys::encode_pkcs8_pem_encrypted(
                    &keypair,
                    passphrase.as_bytes(),
                    1,
                    &mut private_key,
                )?;
            }
            None => {
                russh_keys::encode_pkcs8_pem(&keypair, &mut private_key)?;
            }
        }

        let private_key = Zeroizing::new(core::str::from_utf8(&private_key).unwrap().to_string());
        let public_key = public_key.public_key_base64();
        Ok((private_key, public_key))
    }
}

pub fn generate_key(desc: KeygenDesc) -> Result<KeyPair> {
    let seed = Zeroizing::new(
        desc.passphrase
            .chars()
            .chain(desc.tag.unwrap_or_default().chars())
            .collect::<String>(),
    );

    let mut hash = [0u8; 32];
    match desc.hash {
        HashDesc::Sha256 { iterations } => {
            let mut hasher = sha2::Sha256::new();
            hasher.update(seed.as_bytes());
            hasher.update(desc.salt.as_slice());

            let mut digest = hasher.finalize();

            for _ in 0..(iterations - 1) {
                digest = sha2::Sha256::digest(digest.as_slice());
            }

            hash.copy_from_slice(&digest[..32]);
        }
        HashDesc::Argon {
            variant,
            version,
            memory,
            time,
            lanes,
            hash_length,
        } => {
            let argon = argon2::Argon2::new(
                variant,
                version,
                argon2::Params::new(memory, time, lanes, Some(hash_length as usize))?,
            );

            argon.hash_password_into(seed.as_bytes(), &desc.salt, &mut hash)?;
        }
    };

    // let hash = rand_chacha::ChaChaRng::from_seed(hash).gen::<[u8; 32]>();
    let private_key = ed25519_dalek::SigningKey::from_bytes(&hash);
    let public_key = private_key.verifying_key();

    return Ok(KeyPair {
        passphrase: desc.encrypt.then_some(desc.passphrase),
        private_key,
        public_key,
    });
    // let keypair = russh_keys::key::KeyPair::Ed25519(private_key);
    // let public_key = keypair.clone_public_key().expect("pubkey");
    // let fingerprint = public_key.fingerprint();

    // let pubbytes = public_key.public_key_bytes();

    // let randomart = randomart::RandomArt::from_digest(sha2::Sha256::digest(&pubbytes).as_slice())
    //     .render("ED25519 256", "SHA256")
    //     .expect("randomart");

    // let mut private_key = Vec::new();
    // if desc.encrypt {
    //     russh_keys::encode_pkcs8_pem_encrypted(
    //         &keypair,
    //         desc.passphrase.as_bytes(),
    //         1,
    //         &mut private_key,
    //     )
    //     .expect("writing private key");
    // } else {
    //     russh_keys::encode_pkcs8_pem(&keypair, &mut private_key).expect("writing private key");
    // }
}

#[derive(Debug)]
pub struct KeygenDesc {
    pub hash: HashDesc,
    pub salt: Vec<u8>,
    pub tag: Option<Zeroizing<String>>,
    pub passphrase: Zeroizing<String>,
    pub encrypt: bool,
}

impl Default for KeygenDesc {
    fn default() -> Self {
        Self {
            hash: Default::default(),
            salt: Self::SALT.to_vec(),
            tag: None,
            passphrase: Zeroizing::new("password123".to_string()),
            encrypt: true,
        }
    }
}

impl KeygenDesc {
    /// known random salt.
    pub const SALT: &'static [u8; 32] = b"10ru4YBD5wvdQZc2Mw7Brx45jCxGUM3F";
}

#[derive(Debug)]
pub enum HashDesc {
    Sha256 {
        iterations: u32,
    },
    Argon {
        variant: argon2::Algorithm,
        version: argon2::Version,
        memory: u32,
        time: u32,
        lanes: u32,
        hash_length: u32,
    },
}

impl Default for HashDesc {
    fn default() -> Self {
        Self::Argon {
            variant: argon2::Algorithm::Argon2id,
            version: argon2::Version::V0x13,
            memory: 65536,
            time: 4,
            lanes: 1,
            hash_length: 32,
        }
    }
}

#[cfg(test)]
mod tests {
    use super::{generate_key, KeygenDesc};

    #[test]
    fn test_edkey() {
        generate_key(KeygenDesc::default());
    }
}