1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145

use rand::Rng;
use serde::{Deserialize, Serialize};
use serde_big_array::big_array;

pub type Key128Bits = [u8; 16];
pub type Key256Bits = [u8; 32];
pub type Key512Bits = [u8; 64];

big_array! { BigArray; }

#[derive(Serialize, Deserialize, Clone, Copy)]
pub enum Key {
    Key128Bits(Key128Bits),
    Key256Bits(Key256Bits),

    #[serde(with = "BigArray")]
    Key512Bits(Key512Bits),
}

impl Key {
    /// Converts slice of bytes to a key if it is the right size,
    /// otherwise returns nothing
    pub fn from_slice(slice: &[u8]) -> Option<Self> {
        use std::convert::TryInto;
        if slice.len() == KeySize::Key128Bits.size_in_bytes() {
            slice.try_into().map(Self::Key128Bits).ok()
        } else if slice.len() == KeySize::Key256Bits.size_in_bytes() {
            slice.try_into().map(Self::Key256Bits).ok()
        } else if slice.len() == KeySize::Key512Bits.size_in_bytes() {
            // NOTE: 64-byte array requires special handling due to
            //       limitations in rust right now
            let mut key = [0; 64];
            key.copy_from_slice(slice);
            Some(Self::Key512Bits(key))
        } else {
            None
        }
    }

    pub fn as_slice(&self) -> &[u8] {
        match self {
            Self::Key128Bits(key) => key,
            Self::Key256Bits(key) => key,
            Self::Key512Bits(key) => key,
        }
    }

    pub fn key_size(&self) -> KeySize {
        match self {
            Self::Key128Bits(_) => KeySize::Key128Bits,
            Self::Key256Bits(_) => KeySize::Key256Bits,
            Self::Key512Bits(_) => KeySize::Key512Bits,
        }
    }
}

impl std::fmt::Debug for Key {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::Key128Bits(k) => write!(f, "Key {:?}", k),
            Self::Key256Bits(k) => write!(f, "Key {:?}", k),
            Self::Key512Bits(k) => {
                let k_str = k
                    .iter()
                    .map(|n| n.to_string())
                    .collect::<Vec<String>>()
                    .join(",");
                write!(f, "Key [{:?}]", k_str)
            }
        }
    }
}

impl From<Key128Bits> for Key {
    fn from(buffer: Key128Bits) -> Self {
        Self::Key128Bits(buffer)
    }
}

impl From<Key256Bits> for Key {
    fn from(buffer: Key256Bits) -> Self {
        Self::Key256Bits(buffer)
    }
}

impl From<Key512Bits> for Key {
    fn from(buffer: Key512Bits) -> Self {
        Self::Key512Bits(buffer)
    }
}

#[derive(Debug, Copy, Clone, Serialize, Deserialize)]
pub enum KeySize {
    Key128Bits,
    Key256Bits,
    Key512Bits,
}

impl KeySize {
    pub fn size_in_bytes(self) -> usize {
        match self {
            KeySize::Key128Bits => 16,
            KeySize::Key256Bits => 32,
            KeySize::Key512Bits => 64,
        }
    }
}

impl From<KeySize> for usize {
    fn from(key_size: KeySize) -> Self {
        match key_size {
            KeySize::Key128Bits => 16,
            KeySize::Key256Bits => 32,
            KeySize::Key512Bits => 64,
        }
    }
}

impl From<KeySize> for Vec<u8> {
    fn from(key_size: KeySize) -> Self {
        match key_size {
            KeySize::Key128Bits => new_128bit_key().to_vec(),
            KeySize::Key256Bits => new_256bit_key().to_vec(),
            KeySize::Key512Bits => new_512bit_key().to_vec(),
        }
    }
}

pub fn new_128bit_key() -> Key128Bits {
    let mut buffer = [0; 16];
    rand::thread_rng().fill(&mut buffer);
    buffer
}

pub fn new_256bit_key() -> Key256Bits {
    let mut buffer = [0; 32];
    rand::thread_rng().fill(&mut buffer);
    buffer
}

pub fn new_512bit_key() -> Key512Bits {
    let mut buffer = [0; 64];
    rand::thread_rng().fill(&mut buffer);
    buffer
}