use crate::{bail, error::Error};
use ctr::Ctr64LE;
use rand::RngExt;
use sm4::{
Sm4,
cipher::{KeyIvInit, StreamCipher},
};
const KEY_SIZE: usize = 16;
const NONCE_SIZE: usize = 16;
pub fn encrypt(plaintext: &[u8], key: &[u8]) -> Result<Vec<u8>, Error> {
let mut rng = rand::rng();
let mut nonce = [0u8; NONCE_SIZE];
rng.fill(&mut nonce);
let mut buf = plaintext.to_vec();
let key = padded_key(key).into();
let iv = nonce.into();
Ctr64LE::<Sm4>::new(&key, &iv).apply_keystream(&mut buf);
buf.extend_from_slice(&nonce);
Ok(buf)
}
pub fn decrypt(data: &[u8], key: &[u8]) -> Result<Vec<u8>, Error> {
if data.len() <= NONCE_SIZE {
bail!("invalid data length");
}
let (ciphertext, bytes) = data.split_at(data.len() - NONCE_SIZE);
let nonce: [u8; NONCE_SIZE] = bytes.try_into()?;
let mut buf = ciphertext.to_vec();
let key = padded_key(key).into();
let iv = nonce.into();
Ctr64LE::<Sm4>::new(&key, &iv).apply_keystream(&mut buf);
Ok(buf)
}
fn padded_key(key: &[u8]) -> [u8; KEY_SIZE] {
let mut padded_key = [0_u8; KEY_SIZE];
let key_len = key.len();
if key_len > KEY_SIZE {
padded_key.copy_from_slice(&key[0..KEY_SIZE]);
} else {
padded_key[0..key_len].copy_from_slice(key);
}
padded_key
}