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
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
use std::error::Error;

use crate::config::KdfParams;

pub const MAGIC_HEADER: [u8; 4] = [0xde, 0xad, 0xb0, 0x17];

const ENCRYPTION_ALGORITHMS: [u8; 2] = [0x01, 0x02];
const HASH_ALGORITHMS: [u8; 1] = [0x01];
const KDF_ALGORITHMS: [u8; 2] = [0x01, 0x02];
const COMPRESSION_ALGORITHMS: [u8; 2] = [0x00, 0x01];

#[derive(PartialEq, Debug)]
pub struct Header {
    pub version: u16,
    pub data_length: u16,
    pub encryption_alg: u8,
    pub hash_alg: u8,
    pub kdf_alg: u8,
    pub kdf_params: KdfParams,
    pub compression_alg: u8,
    pub master_salt: Vec<u8>,
    pub nonce: Option<Vec<u8>>,
    pub root_key_nonce: Option<Vec<u8>>,
    pub root_key_ciphertext: Option<Vec<u8>>,
}

impl Header {
    pub fn new(
        encryption_alg: Option<u8>,
        hash_alg: Option<u8>,
        kdf_alg: Option<u8>,
        compression_alg: Option<u8>,
        master_salt: Option<Vec<u8>>,
    ) -> Header {
        Header {
            version: 0x01,
            data_length: 0,
            encryption_alg: encryption_alg.unwrap_or(0x01),
            hash_alg: hash_alg.unwrap_or(0x01),
            kdf_alg: kdf_alg.unwrap_or(0x01),
            kdf_params: KdfParams::default(),
            compression_alg: compression_alg.unwrap_or(0x00),
            master_salt: master_salt.unwrap_or(vec![0x00]),
            nonce: Default::default(),
            root_key_nonce: Default::default(),
            root_key_ciphertext: Default::default(),
        }
    }

    /// Read fixed part of header which will return version and the header total length.
    /// Return `Err` if the data is not started with `MAGIC_HEADER` bytes
    pub fn quick_lookup(data: Vec<u8>) -> Result<(u16, u16), Box<dyn Error>> {
        if data[0..4] != MAGIC_HEADER {
            Err("Data is not valid deadbolt database")?;
        }

        let version = ((data[4] as u16) << 8) | (data[5] as u16);
        let data_length = ((data[6] as u16) << 8) | (data[7] as u16);

        Ok((version, data_length))
    }

    /// Initialize new `Header` from `Vec<u8>`
    pub fn new_from_vector(data: Vec<u8>) -> Result<Header, Box<dyn Error>> {
        let (version, data_length) = Header::quick_lookup(data[..8].to_vec())?;
        let mut header = Header::new(None, None, None, None, None);
        header.version = version;
        header.data_length = data_length;
        header.parse_tlv_header(data[8..].to_vec())?;

        Ok(header)
    }

    /// Parse the dynamic part (Type-Length-Value list) of header
    pub fn parse_tlv_header(&mut self, data: Vec<u8>) -> Result<(), Box<dyn Error>> {
        let mut index = 0;

        loop {
            let data_type = data[index];
            index += 1;
            let start = index + 1;
            let end = start + usize::from(data[index]);
            index = end;

            match data_type {
                1 => {
                    self.encryption_alg = u8::from_be_bytes(data[start..end].try_into()?);
                }
                2 => {
                    self.hash_alg = u8::from_be_bytes(data[start..end].try_into()?);
                }
                3 => {
                    self.kdf_alg = u8::from_be_bytes(data[start..end].try_into()?);
                }
                4 => {
                    self.compression_alg = u8::from_be_bytes(data[start..end].try_into()?);
                }
                6 => {
                    self.master_salt = data[start..end].to_vec();
                }
                7 => {
                    self.nonce = Some(data[start..end].to_vec());
                }
                8 => {
                    self.root_key_nonce = Some(data[start..end].to_vec());
                }
                9 => {
                    self.root_key_ciphertext = Some(data[start..end].to_vec());
                }
                10 => {
                    self.kdf_params.iterations = u8::from_be_bytes(data[start..end].try_into()?);
                }
                11 => {
                    self.kdf_params.threads = u8::from_be_bytes(data[start..end].try_into()?);
                }
                12 => {
                    self.kdf_params.memory = u32::from_be_bytes(data[start..end].try_into()?);
                }
                0 => {
                    break;
                }
                _ => {}
            }
        }

        Ok(())
    }

    /// Construct header section and return the buffer.
    /// It will also automatically update `data_length`
    pub fn construct_header(&mut self) -> Result<Vec<u8>, Box<dyn Error>> {
        let mut buffer: Vec<u8> = MAGIC_HEADER.to_vec();
        buffer.extend(self.version.to_be_bytes());

        let mut tmp_buffer: Vec<u8> = vec![];
        // Encryption algorithm
        tmp_buffer.extend([0x01, 0x01]);
        tmp_buffer.push(self.encryption_alg);

        // Hash algorithm
        tmp_buffer.extend([0x02, 0x01]);
        tmp_buffer.push(self.hash_alg);

        // KDF algorithm
        tmp_buffer.extend([0x03, 0x01]);
        tmp_buffer.push(self.kdf_alg);

        // Compression algorithm
        tmp_buffer.extend([0x04, 0x01]);
        tmp_buffer.push(self.compression_alg);

        // Master salt
        tmp_buffer.push(0x06);
        tmp_buffer.push(self.master_salt.len().try_into()?);
        tmp_buffer.extend(self.master_salt.clone());

        // Data encryption nonce
        let nonce = self.nonce.as_ref().unwrap();
        let nonce_size: u8 = nonce.len().try_into()?;
        tmp_buffer.push(0x07);
        tmp_buffer.push(nonce_size);
        tmp_buffer.extend(nonce);

        // Root key encryption nonce
        let root_key_nonce = self.root_key_nonce.as_ref().unwrap();
        let nonce_size: u8 = root_key_nonce.len().try_into()?;
        tmp_buffer.push(0x08);
        tmp_buffer.push(nonce_size);
        tmp_buffer.extend(root_key_nonce);

        // Encrypted root key
        let root_key_ciphertext = self.root_key_ciphertext.as_ref().unwrap();
        let key_size: u8 = root_key_ciphertext.len().try_into()?;
        tmp_buffer.push(0x09);
        tmp_buffer.push(key_size);
        tmp_buffer.extend(root_key_ciphertext);

        // Key Derivation params
        // Iterations
        tmp_buffer.extend([0x0a, 0x01]);
        tmp_buffer.push(self.kdf_params.iterations);
        // Threads
        tmp_buffer.extend([0x0b, 0x01]);
        tmp_buffer.push(self.kdf_params.threads);
        // Memory (in kibibytes)
        tmp_buffer.extend([0x0c, 0x04]);
        tmp_buffer.extend(self.kdf_params.memory.to_be_bytes());

        // End of header section
        tmp_buffer.extend([0x0, 0x0]);

        // Count header section length and append to buffer
        let header_length: u16 = (tmp_buffer.len() + 8).try_into()?;
        buffer.extend(header_length.to_be_bytes());
        buffer.extend(tmp_buffer);

        self.data_length = header_length;

        Ok(buffer)
    }

    /// Check whether the header is valid in the current version
    pub fn validate_version(&mut self) -> bool {
        match self.version {
            0x01 => {
                if !ENCRYPTION_ALGORITHMS.contains(&self.encryption_alg) {
                    return false;
                }
                if !HASH_ALGORITHMS.contains(&self.hash_alg) {
                    return false;
                }
                if !KDF_ALGORITHMS.contains(&self.kdf_alg) {
                    return false;
                }
                if !COMPRESSION_ALGORITHMS.contains(&self.compression_alg) {
                    return false;
                }
                true
            }
            _ => false,
        }
    }
}