cfun 0.2.11

Tidy up common functions
Documentation
#[cfg(feature = "crypto")]
use md5::{Digest, Md5};
#[cfg(feature = "crypto")]
use sha2::Sha256;

/// get md5 of data
#[cfg(feature = "crypto")]
pub fn md5(data: &Vec<u8>, upper: bool) -> String {
    let mut hasher = Md5::new();
    hasher.update(data);
    let result = hasher.finalize();
    if upper {
        format!("{:X}", result)
    } else {
        format!("{:x}", result)
    }
}

/// get sha256 of data
#[cfg(feature = "crypto")]
pub fn sha256(data: &Vec<u8>, upper: bool) -> String {
    let mut hasher = Sha256::new();
    hasher.update(data);
    let result = hasher.finalize();
    if upper {
        format!("{:X}", result)
    } else {
        format!("{:x}", result)
    }
}

/// get sha256 of data
#[cfg(feature = "crypto")]
pub fn sha512(data: &Vec<u8>, upper: bool) -> String {
    use sha2::Sha512;
    let mut hasher = Sha512::new();
    hasher.update(data);
    let result = hasher.finalize();
    if upper {
        format!("{:X}", result)
    } else {
        format!("{:x}", result)
    }
}

/// encode data with base64
#[cfg(feature = "crypto")]
pub fn base64_encode(data: &Vec<u8>) -> String {
    use base64::{prelude::BASE64_STANDARD, Engine};
    BASE64_STANDARD.encode(data)
}

/// decode data with base64
#[cfg(feature = "crypto")]
pub fn base64_decode(data: String) -> Result<Vec<u8>, base64::DecodeError> {
    use base64::{prelude::BASE64_STANDARD, Engine};
    BASE64_STANDARD.decode(data)
}

#[cfg(feature = "crypto")]
pub fn rsa_generate(bits: Option<usize>) -> (String, String) {
    use rsa::{
        pkcs1::{EncodeRsaPrivateKey, EncodeRsaPublicKey},
        rand_core::OsRng,
        RsaPrivateKey, RsaPublicKey,
    };

    let bits = bits.unwrap_or(2048);
    let priv_key = RsaPrivateKey::new(&mut OsRng, bits).expect("failed to generate a key");
    let pub_key = RsaPublicKey::from(&priv_key);
    let priv_key = priv_key
        .to_pkcs1_pem(rsa::pkcs8::LineEnding::CR)
        .unwrap()
        .to_string();
    let pub_key = pub_key.to_pkcs1_pem(rsa::pkcs8::LineEnding::CR).unwrap();
    (pub_key, priv_key)
}

#[cfg(feature = "crypto")]
pub fn rsa_pkcs1_pem_encrypt(pub_key: &str, msg: &Vec<u8>) -> Vec<u8> {
    use rsa::{pkcs1::DecodeRsaPublicKey, rand_core::OsRng, Pkcs1v15Encrypt, RsaPublicKey};
    let pubkey = RsaPublicKey::from_pkcs1_pem(&pub_key).unwrap();
    let mut rng = OsRng;
    let ciphertext = pubkey.encrypt(&mut rng, Pkcs1v15Encrypt, msg).unwrap();
    ciphertext
}

#[cfg(feature = "crypto")]
pub fn rsa_pkcs1_pem_decrypt(priv_key: &str, ciphertext: Vec<u8>) -> Vec<u8> {
    use rsa::{pkcs1::DecodeRsaPrivateKey, Pkcs1v15Encrypt, RsaPrivateKey};
    let privkey = RsaPrivateKey::from_pkcs1_pem(&priv_key).unwrap();
    let data = privkey.decrypt(Pkcs1v15Encrypt, &ciphertext).unwrap();
    data
}

#[cfg(feature = "crypto")]
pub fn rsa_encrypt_pem(pub_pem: &str, cnt: Vec<u8>) -> Result<Vec<u8>, anyhow::Error> {
    use rsa::{pkcs8::DecodePublicKey, rand_core::OsRng, Pkcs1v15Encrypt, RsaPublicKey};
    let pubkey = RsaPublicKey::from_public_key_pem(&pub_pem)?;
    let mut rng = OsRng;
    let ciphertext = pubkey.encrypt(&mut rng, Pkcs1v15Encrypt, &cnt)?;
    Ok(ciphertext)
}

#[cfg(feature = "crypto")]
pub fn rsa_encrypt_key(pub_key: &str, cnt: Vec<u8>) -> Result<Vec<u8>, anyhow::Error> {
    use rsa::{pkcs8::DecodePublicKey, rand_core::OsRng, Pkcs1v15Encrypt, RsaPublicKey};
    let pem_str = format!(
        "-----BEGIN PUBLIC KEY-----\n{}\n-----END PUBLIC KEY-----",
        wrap64(pub_key)
    );
    let pubkey = RsaPublicKey::from_public_key_pem(&pem_str)?;
    let mut rng = OsRng;
    let ciphertext = pubkey.encrypt(&mut rng, Pkcs1v15Encrypt, &cnt)?;
    Ok(ciphertext)
}

#[cfg(feature = "crypto")]
fn wrap64(s: &str) -> String {
    s.as_bytes()
        .chunks(64)
        .map(|chunk| std::str::from_utf8(chunk).unwrap())
        .collect::<Vec<_>>()
        .join("\n")
}

#[cfg(feature = "crypto")]
pub fn encrypt_ecb(plaintext: &[u8]) -> (Vec<u8>, Vec<u8>) {
    use aes::{
        cipher::{block_padding::Pkcs7, BlockEncrypt, KeyInit},
        Aes256Enc,
    };
    use rand::{rngs::OsRng, TryRngCore};
    let mut key = [0u8; 32];
    OsRng.try_fill_bytes(&mut key).unwrap();
    let aes = Aes256Enc::new_from_slice(&key).unwrap();
    let mut block = Vec::new();
    block.resize(plaintext.len() + 16, 0);
    let ret = aes
        .encrypt_padded_b2b::<Pkcs7>(plaintext, &mut block)
        .unwrap();
    (key.to_vec(), ret.to_vec())
}

#[cfg(feature = "crypto")]
pub fn decrypt_ecb(key: &[u8], ciphertext: &[u8]) -> Vec<u8> {
    use aes::{
        cipher::{block_padding::Pkcs7, BlockDecrypt, KeyInit},
        Aes256Dec,
    };

    let cipher = Aes256Dec::new_from_slice(key).unwrap();
    let mut block = Vec::new();
    block.resize(ciphertext.len(), 0);
    let data = cipher
        .decrypt_padded_b2b::<Pkcs7>(ciphertext, &mut block)
        .unwrap();
    data.to_vec()
}

#[cfg(feature = "crypto")]
pub fn encrypt_gcm(plaintext: &[u8], aad_opt: Option<&[u8]>) -> (Vec<u8>, Vec<u8>, Vec<u8>) {
    use aes_gcm::{
        aead::{Aead, OsRng, Payload},
        AeadCore, Aes256Gcm, KeyInit,
    };

    let key = Aes256Gcm::generate_key(OsRng);
    let aes = Aes256Gcm::new(&key);
    let nonce = Aes256Gcm::generate_nonce(&mut OsRng); // 96-bits; unique per message
    let payload = match aad_opt {
        Some(aad) => Payload {
            msg: plaintext,
            aad,
        },
        None => Payload {
            msg: plaintext,
            aad: &[],
        },
    };
    let cipher = aes.encrypt(&nonce, payload).unwrap(); // 4. 如果存在 AAD 就传入,若无则传空切片
    (key.to_vec(), cipher, nonce.to_vec())
}

#[cfg(feature = "crypto")]
pub fn decrypt_gcm(
    key: &Vec<u8>,
    cipher: &Vec<u8>,
    nonce: &Vec<u8>,
    aad_opt: Option<&[u8]>,
) -> Vec<u8> {
    use aes_gcm::{
        aead::{Aead, Payload},
        Aes256Gcm, Key, KeyInit, Nonce,
    };

    let key = Key::<Aes256Gcm>::from_slice(key);
    let nonce = Nonce::from_slice(nonce);
    let aes = Aes256Gcm::new(key);
    let payload = match aad_opt {
        Some(aad) => Payload { msg: &cipher, aad },
        None => Payload {
            msg: &cipher,
            aad: &[],
        },
    };
    let plaintext = aes.decrypt(nonce, payload).unwrap();
    plaintext
}

#[test]
#[cfg(feature = "crypto")]
fn test_md5() {
    let data = b"hello world";
    assert_eq!(
        md5(&data.to_vec(), false),
        "5eb63bbbe01eeed093cb22bb8f5acdc3"
    );
}

#[test]
#[cfg(feature = "crypto")]
fn c_crypto_test_sha256() {
    let data = b"hello world";
    assert_eq!(
        sha256(&data.to_vec(), false),
        "b94d27b9934d3e08a52e52d7da7dabfac484efe37a5380ee9088f7ace2efcde9"
    );
}

#[test]
#[cfg(feature = "crypto")]
fn c_crypto_test_sha512() {
    let data = b"hello world";
    assert_eq!(
        sha512(&data.to_vec(), false),
        "309ecc489c12d6eb4cc40f50c902f2b4d0ed77ee511a7c7a9bcd3ca86d4cd86f989dd35bc5ff499670da34255b45b0cfd830e81f605dcf7dc5542e93ae9cd76f"
    );
}

#[test]
#[cfg(feature = "crypto")]
fn test_base64_encode() {
    let data = b"hello world";
    assert_eq!(base64_encode(&data.to_vec()), "aGVsbG8gd29ybGQ=");
}

#[test]
#[cfg(feature = "crypto")]
fn test_base64_decode() {
    let data = "aGVsbG8gd29ybGQ=";
    assert_eq!(base64_decode(data.to_string()).unwrap(), b"hello world");
}