use aes::cipher::{KeyIvInit, StreamCipher};
use std::vec;
type Aes128Ctr64LE = ctr::Ctr64LE<aes::Aes128>;
pub struct Prg {
seed: Vec<u8>,
counter: u64,
}
impl Prg {
const PRG_NONCE: u64 = 0x0123456789ABCDEF;
const PRG_INITIAL_COUNTER: u64 = 0;
const KEY_LEN: usize = 16;
const IV_LEN: usize = 16;
const BLOCK_LEN: usize = 16;
pub fn new(seed: Option<Vec<u8>>) -> Prg {
let cropped_seed = if let Some(mut value_seed) = seed {
if value_seed.len() > Self::KEY_LEN + Self::IV_LEN {
value_seed[0..Self::KEY_LEN + Self::IV_LEN].to_vec()
} else {
let original_seed_length = value_seed.len();
let mut appended_seed = Vec::new();
appended_seed.append(&mut value_seed);
appended_seed.append(&mut vec![
0;
Self::KEY_LEN + Self::IV_LEN - original_seed_length
]);
appended_seed
}
} else {
vec![0; Self::KEY_LEN + Self::IV_LEN]
};
let counter = Self::PRG_INITIAL_COUNTER;
let mut prg = Prg {
seed: cropped_seed,
counter,
};
prg.init();
prg
}
pub fn init(&mut self) {
self.counter = Self::PRG_INITIAL_COUNTER;
}
pub fn reset(&mut self) {
self.init()
}
pub fn counter(&self) -> u64 {
self.counter
}
pub fn next(&mut self, n_bytes: usize) -> Vec<u8> {
if n_bytes == 0 {
return Vec::new();
}
let mut n_blocks = n_bytes / Self::BLOCK_LEN;
if n_bytes % Self::BLOCK_LEN != 0 {
n_blocks += 1;
}
let key = &self.seed[0..Self::KEY_LEN];
let iv = &self.seed[Self::KEY_LEN..];
let mut cipher = Aes128Ctr64LE::new(key.into(), iv.into());
let mut out = Vec::new();
for _ in 0..n_blocks {
let mut buffer = [Self::PRG_NONCE.to_ne_bytes(), self.counter.to_ne_bytes()].concat();
cipher.apply_keystream(&mut buffer);
out.append(&mut buffer);
self.counter += 1;
}
out[..n_bytes].to_vec()
}
}