#![allow(dead_code)]

pub mod algorithms;
mod compression;
pub mod config;
pub mod header;
pub mod password;

use std::error::Error;
use std::fs::{self, File};
use std::io::Read;
use std::string::String;

use config::Config;
use header::Header;
use secrecy::{ExposeSecret, Secret};

use deadbolt_crypto::encryption::{Encryption, StreamEncryption};
use deadbolt_crypto::hash::{Hkdf, Hmac, PasswordHash};
use deadbolt_crypto::rand::Random;

type MasterKeySignature = (Secret<Vec<u8>>, Secret<Vec<u8>>);

#[derive(Debug, Clone, PartialEq)]
pub struct Data {
    pub key: String,
    pub metadata: String,
    pub value: Vec<u8>,
}

impl Data {
    pub fn new(key: String, metadata: String, value: Vec<u8>) -> Data {
        Data {
            key,
            metadata,
            value,
        }
    }
}

pub struct Parser {
    pub header: Header,
    pub data: Vec<Data>,
    pub signature_key: Secret<Vec<u8>>,
    pub root_key: Secret<Vec<u8>>,
    pub filepath: String,
    raw_content: Vec<u8>,
    pub config: Config,
}

impl Parser {
    pub fn new(
        signature_key: Secret<Vec<u8>>,
        root_key: Secret<Vec<u8>>,
        master_salt: Vec<u8>,
        encrypted_root_key: Option<Vec<u8>>,
        nonce: Option<Vec<u8>>,
    ) -> Parser {
        let mut header = Header::new(None, None, None, None, Some(master_salt));
        header.root_key_ciphertext = encrypted_root_key;
        header.root_key_nonce = nonce;

        Parser {
            header,
            data: vec![],
            signature_key,
            root_key,
            filepath: "".to_string(),
            raw_content: vec![],
            config: config::default_config(),
        }
    }

    pub fn new_default() -> Parser {
        Parser {
            header: Header::new(None, None, None, None, None),
            data: vec![],
            signature_key: Secret::new(vec![]),
            root_key: Secret::new(vec![]),
            filepath: "".to_string(),
            raw_content: vec![],
            config: config::default_config(),
        }
    }

    /// Read db from file and store the content.
    /// Only header section will be parsed
    pub fn read_from_file(&mut self, filepath: &str) -> Result<(), Box<dyn Error>> {
        let raw_data = fs::read(filepath)?;
        let header = Header::new_from_vector(raw_data.clone())?;

        self.header = header;
        self.raw_content = raw_data;
        self.filepath = filepath.to_string();

        Ok(())
    }

    /// Verify password without having to read the whole file content
    pub fn verify_password_from_file(
        filepath: &str,
        master_password: Secret<String>,
    ) -> Result<MasterKeySignature, Box<dyn Error>> {
        // Read only the first 10 bytes of file
        let mut file = File::open(filepath)?;
        let mut fixed_header_data: Vec<u8> = vec![0; 8];
        file.read_exact(&mut fixed_header_data)?;

        // Read only header section of the file
        let (version, header_length) = Header::quick_lookup(fixed_header_data)?;
        let mut header_data: Vec<u8> = vec![0; header_length.try_into().unwrap()];
        file.read_exact(&mut header_data)?;

        let mut header = Header::new(None, None, None, None, None);
        header.version = version;
        header.data_length = header_length;
        header.parse_tlv_header(header_data)?;

        let (master_key, signature_key) = Parser::derive_master_key(
            master_password,
            &header.master_salt,
            algorithms::KdfAlg::from_u8(header.kdf_alg)
                .to_str()
                .to_owned(),
            header.kdf_params.iterations,
            header.kdf_params.threads,
            header.kdf_params.memory,
        );

        let mut cipher = Encryption::new(&master_key);
        let nonce = header.root_key_nonce.as_ref().unwrap()[..]
            .try_into()
            .unwrap();
        let decryption_result = cipher.decrypt(
            nonce,
            header.root_key_ciphertext.as_ref().unwrap().to_vec(),
            header.encryption_alg,
        );

        match decryption_result {
            Ok(_) => Ok((master_key, signature_key)),
            Err(e) => Err(e),
        }
    }

    /// Derive `master_key` and `signature_key` from `master_password`
    pub fn derive_master_key(
        master_password: Secret<String>,
        salt: &[u8],
        algorithm: String,
        iterations: u8,
        threads: u8,
        memory: u32,
    ) -> MasterKeySignature {
        let mut password_hash: Secret<Vec<u8>> = Secret::new(vec![]);
        if algorithm == "argon2d" {
            password_hash =
                PasswordHash::argon2d(master_password, salt, iterations, threads, memory);
        } else if algorithm == "argon2id" {
            password_hash =
                PasswordHash::argon2id(master_password, salt, iterations, threads, memory);
        }
        let key_material: Secret<Vec<u8>> =
            Hkdf::expand(password_hash.expose_secret().to_vec().into());
        let master_key = Secret::new(key_material.expose_secret()[..32].to_vec());
        let signature_key = Secret::new(key_material.expose_secret()[32..].to_vec());

        (master_key, signature_key)
    }

    /// Decrypt `encrypted_root_key` and set it as current `root_key`
    pub fn get_root_key(&mut self, master_key: &Secret<Vec<u8>>) {
        let mut cipher = Encryption::new(master_key);
        let root_key = cipher
            .decrypt(
                self.header
                    .root_key_nonce
                    .as_ref()
                    .unwrap()
                    .to_vec()
                    .try_into()
                    .unwrap(),
                self.header.root_key_ciphertext.as_ref().unwrap().to_vec(),
                self.header.encryption_alg,
            )
            .expect("Invalid master key");

        self.root_key = Secret::new(root_key.data);
    }

    /// Generate new random `root_key` and its encrypted form
    pub fn generate_root_key(
        master_key: Secret<Vec<u8>>,
        encryption_alg: u8,
    ) -> (Secret<Vec<u8>>, Vec<u8>, [u8; 12]) {
        let root_key = Secret::new(Random::get_rand_bytes());
        let mut cipher = Encryption::new(&master_key);
        let encrypted_root_key = cipher
            .encrypt(root_key.expose_secret().to_vec(), encryption_alg)
            .unwrap();

        (root_key, encrypted_root_key.data, encrypted_root_key.nonce)
    }

    /// Get the `child_key` from the current `root_key`
    pub fn generate_child_key(&mut self) -> Secret<Vec<u8>> {
        let mut hasher = Hmac::new(&self.signature_key);
        hasher.update(self.root_key.expose_secret());

        Secret::new(hasher.finalize())
    }

    /// Add new entry to the `Parser.data`
    pub fn new_password_entry(&mut self, key: String, metadata: String, password: Secret<String>) {
        let ciphertext = StreamEncryption::encrypt(
            &self.generate_child_key(),
            password.expose_secret().as_bytes(),
        );

        self.data.push(Data {
            key,
            metadata,
            value: ciphertext,
        })
    }

    /// Parse the TLV data
    pub fn read_tlv(data: Vec<u8>) -> Result<Vec<Data>, Box<dyn Error>> {
        let mut index = 0;
        let mut current_data_index = 0;

        let mut buffer = vec![];

        loop {
            let data_type = data[index];

            match data_type {
                1 => {
                    index += 1;
                    let start = index + 1;
                    let end = start + usize::from(data[index]);
                    let current_key = String::from_utf8_lossy(&data[start..end]).to_string();
                    index = end;

                    buffer.push(Data {
                        key: current_key,
                        metadata: "".to_string(),
                        value: vec![],
                    });
                    current_data_index += 1;
                }
                2 => {
                    index += 1;
                    let start = index + 4;
                    let length = u32::from_be_bytes(data[index..index + 4].try_into()?);
                    let end = start + usize::try_from(length)?;
                    let current_metadata = String::from_utf8_lossy(&data[start..end]).to_string();
                    index = end;

                    buffer[current_data_index - 1].metadata = current_metadata;
                }
                3 => {
                    index += 1;
                    let start = index + 1;
                    let end = start + usize::from(data[index]);
                    let current_value = data[start..end].to_vec();
                    index = end;

                    buffer[current_data_index - 1].value = current_value;
                }
                _ => {
                    break;
                }
            }
        }

        Ok(buffer)
    }

    /// Read the whole data section and store it in `Parser.data`
    fn parse_data(&mut self) -> Result<(), Box<dyn Error>> {
        let mut cipher = Encryption::new(&self.root_key);
        let nonce = self.header.nonce.as_ref().unwrap()[..].try_into()?;
        let header_length: usize = (self.header.data_length + 32).try_into()?;
        let decryption_result = cipher.decrypt(
            nonce,
            self.raw_content[header_length..].to_vec(),
            self.header.encryption_alg,
        );

        match decryption_result {
            Ok(decrypted_data) => {
                let mut plain_data = decrypted_data.data;
                if self.header.compression_alg != 0 {
                    plain_data = compression::Gzip::decompress(plain_data)?;
                }
                self.data = Parser::read_tlv(plain_data)?;

                Ok(())
            }

            Err(e) => Err(e),
        }
    }

    /// Construct the db content
    fn construct_content(&mut self) -> Result<(), Box<dyn Error>> {
        // Construct plain data section
        let mut plain_data: Vec<u8> = vec![];
        for data in &self.data {
            // Key type
            plain_data.push(0x01);
            plain_data.push(data.key.len().try_into()?);
            plain_data.extend(data.key.as_bytes());

            // Metadata type
            let metadata_len: u32 = data.metadata.len().try_into()?;
            plain_data.push(0x02);
            plain_data.extend(metadata_len.to_be_bytes().to_vec());
            plain_data.extend(data.metadata.as_bytes());

            // Password type
            plain_data.push(0x03);
            plain_data.push(data.value.len().try_into()?);
            plain_data.extend(data.value.clone());
        }

        // Terminator (end of data section)
        plain_data.push(0x00);

        // Compress plain data (optional)
        if self.config.algorithms.compression {
            plain_data = compression::Gzip::compress(plain_data)?;
        }

        // Encrypt data section
        let mut cipher = Encryption::new(&self.root_key);
        let encrypted_data = cipher.encrypt(
            plain_data,
            algorithms::EncryptionAlg::from_str(&self.config.algorithms.encryption).value(),
        )?;

        // Update existing variables of parser
        self.header.encryption_alg =
            algorithms::EncryptionAlg::from_str(&self.config.algorithms.encryption).value();
        self.header.hash_alg = algorithms::HashAlg::from_str(&self.config.algorithms.hash).value();
        self.header.kdf_alg =
            algorithms::KdfAlg::from_str(&self.config.algorithms.kdf.algorithm).value();
        self.header.compression_alg =
            algorithms::CompressionAlg::from_bool(self.config.algorithms.compression).value();
        self.header.nonce = Some(encrypted_data.nonce.to_vec());

        match &self.config.algorithms.kdf.parameters {
            Some(params) => {
                self.header.kdf_params.iterations = params.iterations;
                self.header.kdf_params.threads = params.threads;
                self.header.kdf_params.memory = params.memory;
            }
            None => {}
        }

        // Construct header section
        let mut buffer: Vec<u8> = self.header.construct_header()?;

        // Construct signature section
        let mut hmac = Hmac::new(&self.signature_key);
        hmac.update(&buffer);
        let tag = hmac.finalize();
        buffer.extend(tag);

        // And encrypted data section to constructed buffer
        buffer.extend(encrypted_data.data);

        self.raw_content = buffer;

        Ok(())
    }

    /// Verify whether the signature is correct with the header data
    pub fn verify_signature(&mut self) -> Result<(), Box<dyn Error>> {
        let mut hmac = Hmac::new(&self.signature_key);
        let header_len: usize = self.header.data_length.try_into()?;
        hmac.update(&self.raw_content[..header_len]);
        let expected_mac = &self.raw_content[header_len..header_len + 32];

        hmac.verify(expected_mac)?;
        Ok(())
    }

    /// Parse `Parser.raw_content` into `Parser.data`
    pub fn read_raw_content(&mut self) -> Result<(), Box<dyn Error>> {
        if self.verify_signature().is_ok() {
            if !self.header.validate_version() {
                return Err("Invalid header data".to_string().into());
            }
            return self.parse_data();
        }
        Err("File might be corrupted".to_string().into())
    }

    /// Write content into file
    pub fn write_to_file(&mut self) -> Result<(), Box<dyn Error>> {
        self.construct_content()?;
        fs::write(&self.filepath, &self.raw_content)?;
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use std::env::temp_dir;

    use crate::{Data, Parser};
    use deadbolt_crypto::encryption::StreamEncryption;
    use secrecy::{ExposeSecret, Secret};

    fn generate_temp_test_file(filename: &str) -> String {
        let mut test_path = temp_dir();
        test_path.push(filename);

        test_path.to_string_lossy().to_string()
    }

    #[test]
    fn test_read_tlv() {
        let test_data: [Data; 3] = [
            Data {
                key: "test1".to_string(),
                metadata: "test1".to_string(),
                value: vec![0x01, 0x02, 0x03],
            },
            Data {
                key: "test2".to_string(),
                metadata: "test2".to_string(),
                value: vec![],
            },
            Data {
                key: "test3".to_string(),
                metadata: "".to_string(),
                value: vec![],
            },
        ];
        let length: u32 = test_data[0].metadata.len().try_into().unwrap();
        let mut buffer: Vec<u8> = vec![0x01];
        buffer.push(test_data[0].key.len().try_into().unwrap());
        buffer.extend(test_data[0].key.as_bytes());
        buffer.push(0x02);
        buffer.extend(length.to_be_bytes());
        buffer.extend(test_data[0].metadata.as_bytes());
        buffer.push(0x03);
        buffer.push(test_data[0].value.len().try_into().unwrap());
        buffer.extend(test_data[0].value.clone());

        buffer.push(0x01);
        buffer.push(test_data[1].key.len().try_into().unwrap());
        buffer.extend(test_data[1].key.as_bytes());
        buffer.push(0x02);
        buffer.extend(length.to_be_bytes());
        buffer.extend(test_data[1].metadata.as_bytes());

        buffer.push(0x01);
        buffer.push(test_data[2].key.len().try_into().unwrap());
        buffer.extend(test_data[2].key.as_bytes());

        buffer.push(0x00);

        let result = Parser::read_tlv(buffer).unwrap();

        assert_eq!(result[0], test_data[0]);
        assert_eq!(result[1], test_data[1]);
        assert_eq!(result[2], test_data[2]);
    }

    #[test]
    fn test_read_write_file() {
        let path = &generate_temp_test_file("test-1.dblt");
        let master_password = Secret::new("password123".to_string());

        let (master_key, signature_key) = Parser::derive_master_key(
            master_password.clone(),
            "yellow_submarine".as_bytes(),
            "argon2d".to_string(),
            2,
            1,
            19 * 1024,
        );
        let (root_key, encrypted_root_key, root_key_nonce) =
            Parser::generate_root_key(Secret::new(master_key.expose_secret().clone()), 1);
        let mut write_parser = Parser::new(
            Secret::new(signature_key.expose_secret().clone()),
            Secret::new(root_key.expose_secret().clone()),
            "yellow_submarine".as_bytes().to_vec(),
            Some(encrypted_root_key.clone()),
            Some(root_key_nonce.to_vec()),
        );
        let test_str = "test";
        write_parser.new_password_entry(
            test_str.to_string(),
            test_str.to_string(),
            Secret::new(test_str.to_string()),
        );
        write_parser.filepath = path.to_string();
        write_parser.write_to_file().unwrap();

        let mut read_parser = Parser::new(
            signature_key,
            root_key,
            "yellow_submarine".as_bytes().to_vec(),
            Some(encrypted_root_key),
            Some(root_key_nonce.to_vec()),
        );
        read_parser.read_from_file(path).unwrap();
        read_parser.read_raw_content().unwrap();

        assert!(Parser::verify_password_from_file(path, master_password).is_ok());
        assert_eq!(write_parser.header, read_parser.header);
        assert_eq!(write_parser.data.len(), read_parser.data.len());

        assert_eq!(write_parser.data[0].key, read_parser.data[0].key);
        assert_eq!(write_parser.data[0].metadata, read_parser.data[0].metadata);
        assert_eq!(write_parser.data[0].value, read_parser.data[0].value);

        let child_key = write_parser.generate_child_key();
        let writer_plaintext = StreamEncryption::decrypt(&child_key, &write_parser.data[0].value);
        let reader_plaintext = StreamEncryption::decrypt(&child_key, &read_parser.data[0].value);

        assert_eq!(writer_plaintext, reader_plaintext);
    }

    #[test]
    fn test_invalid_password() {
        let path = &generate_temp_test_file("test-2.dblt");
        let (master_key, signature_key) = Parser::derive_master_key(
            Secret::new("password123".to_string()),
            "yellow_submarine".as_bytes(),
            "argon2d".to_string(),
            2,
            1,
            19 * 1024,
        );
        let (root_key, encrypted_root_key, root_key_nonce) =
            Parser::generate_root_key(master_key, 1);
        let mut write_parser = Parser::new(
            Secret::new(signature_key.expose_secret().clone()),
            Secret::new(root_key.expose_secret().clone()),
            "yellow_submarine".as_bytes().to_vec(),
            Some(encrypted_root_key),
            Some(root_key_nonce.to_vec()),
        );
        write_parser.filepath = path.to_string();
        write_parser.write_to_file().unwrap();

        let invalid_password = Secret::new("password1234".to_string());

        let expected_from_sample = Parser::verify_password_from_file(path, invalid_password);

        assert!(expected_from_sample.is_err());
    }
}