deadbolt-parser 0.1.1

use std::error::Error;

use deadbolt_crypto::rand::Random;

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_else(Random::get_rand_bytes),
            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,
        }
    }
}