// Copyright 2018 Zach Miller
//
// Licensed under the Apache License, Version 2.0 (http://www.apache.org/licenses/LICENSE-2.0)
// or the MIT License (https://opensource.org/licenses/MIT), at your option.
//
// This file may not be copied, modified, or distributed except according to those terms.

use bincode;
use hmac::{Hmac, Mac};
use sha2::Sha256;

use serde::de::DeserializeOwned;
use serde::ser::Serialize;

use error::Error;
use storage::{CommittedCheckpoint, Storage, UncommittedCheckpoint};

type HmacSha256 = Hmac<Sha256>;

/// Adds HMAC verification to checkpoint data.
///
/// Currently the Sha256 checksum algorithm is used.
#[derive(Debug)]
pub struct HmacWrapper<S: Storage> {
    inner: S,
    secret_key: Vec<u8>,
}

impl<S: Storage> HmacWrapper<S> {
    /// Adds an `HmacWrapper` around a `Storage` object.
    pub fn wrap(inner: S, secret_key: &[u8]) -> Self {
        Self {
            inner,
            secret_key: secret_key.to_vec(),
        }
    }
}

impl<S: Storage> Storage for HmacWrapper<S> {
    type Committed = Committed<S::Committed>;
    type Uncommitted = Uncommitted<S::Uncommitted>;

    fn create_checkpoint(&mut self, identifier: &str) -> Result<Self::Uncommitted, Error> {
        let inner = self.inner.create_checkpoint(identifier)?;
        let hmac = HmacSha256::new_varkey(&self.secret_key).map_err(|_| {
            Error::initialization(format!(
                "Invalid HMAC key length when initializing checkpoint: {} - ",
                identifier
            ))
        })?;

        Ok(Self::Uncommitted { inner, hmac })
    }

    fn commit_checkpoint(
        &mut self,
        uncommitted: Self::Uncommitted,
    ) -> Result<Self::Committed, Error> {
        let inner = self.inner.commit_checkpoint(uncommitted.inner)?;
        Ok(Self::Committed {
            inner,
            hmac: uncommitted.hmac,
        })
    }

    fn load_checkpoint(&mut self, identifier: &str) -> Result<Self::Committed, Error> {
        let inner = self.inner.load_checkpoint(identifier)?;
        let hmac = HmacSha256::new_varkey(&self.secret_key).map_err(|_| {
            Error::initialization(format!(
                "Failed to initialize storage for {}: Invalid HMAC key length",
                identifier
            ))
        })?;
        Ok(Self::Committed { inner, hmac })
    }

    fn remove_checkpoint(&mut self, identifier: &str) -> Result<(), Error> {
        self.inner.remove_checkpoint(identifier)
    }

    fn checkpoint_identifiers(&mut self) -> Result<Vec<String>, Error> {
        self.inner.checkpoint_identifiers()
    }
}

/// Wraps the `CommittedCheckpoint` objects from the underlying storage to verify checkpoint data
/// when it is retrieved.
#[derive(Debug)]
pub struct Committed<C: CommittedCheckpoint> {
    inner: C,
    hmac: HmacSha256,
}

impl<C: CommittedCheckpoint> CommittedCheckpoint for Committed<C> {
    fn get<T: DeserializeOwned + Serialize>(&mut self, key: &str) -> Result<T, Error> {
        let helper: Helper = self.inner.get(key)?;
        let identifier = self.identifier()?.to_string();
        self.hmac.input(identifier.as_bytes());
        self.hmac.input(&helper.bytes);

        if self.hmac.verify(&helper.hmac).is_err() {
            Err(Error::integrity(format!(
                "HMAC integrity check failed for checkpoint: {}, key: {}",
                identifier, key
            )))?;
        }

        let value = bincode::deserialize(&helper.bytes).map_err(|_| {
            Error::deserialize(format!(
                "Failed to deserialize value from checkpoint: {}, key: {}",
                identifier, key
            ))
        })?;

        Ok(value)
    }

    fn keys(&mut self) -> Result<Vec<String>, Error> {
        self.inner.keys()
    }

    fn identifier(&mut self) -> Result<&str, Error> {
        self.inner.identifier()
    }
}

/// Wraps the `UncommittedCheckpoint` objects from the underlying storage to add an HMAC when
/// checkpoint data is stored and to verify checkpoint data when it is retrieved.
#[derive(Debug)]
pub struct Uncommitted<U: UncommittedCheckpoint> {
    inner: U,
    hmac: HmacSha256,
}

impl<U: UncommittedCheckpoint> UncommittedCheckpoint for Uncommitted<U> {
    fn get<T: DeserializeOwned + Serialize>(&mut self, key: &str) -> Result<T, Error> {
        let helper: Helper = self.inner.get(key)?;
        let identifier = self.identifier()?.to_string();
        self.hmac.input(identifier.as_bytes());
        self.hmac.input(&helper.bytes);

        if self.hmac.verify(&helper.hmac).is_err() {
            Err(Error::integrity(format!(
                "HMAC integrity check failed for checkpoint: {}, key: {}",
                identifier, key
            )))?;
        }

        let value = bincode::deserialize(&helper.bytes).map_err(|_| {
            Error::deserialize(format!(
                "Failed to deserialize value from checkpoint: {}, key: {}",
                identifier, key
            ))
        })?;

        Ok(value)
    }

    fn keys(&mut self) -> Result<Vec<String>, Error> {
        self.inner.keys()
    }

    fn identifier(&mut self) -> Result<&str, Error> {
        self.inner.identifier()
    }

    fn put<T: DeserializeOwned + Serialize>(&mut self, key: &str, value: &T) -> Result<(), Error> {
        let identifier = self.identifier()?.to_string();
        let bytes = bincode::serialize(value).map_err(|_| {
            Error::serialize(format!(
                "Failed to serialize value into checkpoint: {}, key: {}",
                identifier, key
            ))
        })?;

        self.hmac.input(identifier.as_bytes());
        self.hmac.input(&bytes);
        let helper = Helper {
            bytes,
            hmac: self.hmac.result().code().to_vec(),
        };

        self.inner.put(key, &helper)?;
        Ok(())
    }

    fn remove(&mut self, key: &str) -> Result<(), Error> {
        self.inner.remove(key)
    }
}

// Used to add an HMAC to the values being inserted into the checkpoint.
#[derive(Debug, Deserialize, Serialize)]
struct Helper {
    bytes: Vec<u8>,
    hmac: Vec<u8>,
}

#[cfg(all(test, feature = "filestorage-tests"))]
mod test {
    use std::fs::File;
    use std::io::{Read, Write};

    use error::ErrorKind;
    use storage::{CommittedCheckpoint, FileStorage, Storage, UncommittedCheckpoint};
    use test_utils::create_test_directory_with_checkpoint_folders;
    use wrappers::HmacWrapper;

    // Verify that checkpoint data can be read properly when data corruption does not occur.
    #[test]
    fn no_corruption() {
        let dir = create_test_directory_with_checkpoint_folders();
        let path = dir.path();

        let mut storage = HmacWrapper::wrap(FileStorage::open(&path).unwrap(), b"Secret Key");
        let mut uncommitted = storage.create_checkpoint("CP").unwrap();
        uncommitted.put("Key", &"Test String".to_string()).unwrap();
        storage.commit_checkpoint(uncommitted).unwrap();

        let mut file_path = path.join("committed");
        file_path.push("CP");
        file_path.push("Key");
        let mut loaded = storage.load_checkpoint("CP").unwrap();
        assert_eq!(loaded.get::<String>("Key").unwrap(), "Test String");
    }

    // Verify that corruption of checkpoint data is detected properly.
    #[test]
    fn corruption_in_stored_data() {
        let dir = create_test_directory_with_checkpoint_folders();
        let path = dir.path();

        let mut storage = HmacWrapper::wrap(FileStorage::open(&path).unwrap(), b"Secret Key");
        let mut uncommitted = storage.create_checkpoint("CP").unwrap();
        uncommitted.put("Key", &"Test String".to_string()).unwrap();
        storage.commit_checkpoint(uncommitted).unwrap();

        let mut file_path = path.join("committed");
        file_path.push("CP");
        file_path.push("Key");

        let mut bytes = Vec::new();
        let mut file = File::open(&file_path).unwrap();
        file.read_to_end(&mut bytes).unwrap();
        bytes[16] = b'?';

        let mut file = File::create(&file_path).unwrap();
        file.write(&bytes).unwrap();

        let mut loaded = storage.load_checkpoint("CP").unwrap();
        assert_eq!(
            loaded.get::<String>("Key").unwrap_err().kind(),
            ErrorKind::Integrity
        );
    }

    // Verify that corruption of the HMAC is detected properly.
    #[test]
    fn corruption_in_hmac() {
        let dir = create_test_directory_with_checkpoint_folders();
        let path = dir.path();

        let mut storage = HmacWrapper::wrap(FileStorage::open(&path).unwrap(), b"Secret Key");
        let mut uncommitted = storage.create_checkpoint("CP").unwrap();
        uncommitted.put("Key", &"Test String".to_string()).unwrap();
        storage.commit_checkpoint(uncommitted).unwrap();

        let mut file_path = path.join("committed");
        file_path.push("CP");
        file_path.push("Key");

        let mut bytes = Vec::new();
        let mut file = File::open(&file_path).unwrap();
        file.read_to_end(&mut bytes).unwrap();
        bytes[35] = b'?';

        let mut file = File::create(&file_path).unwrap();
        file.write(&bytes).unwrap();

        let mut loaded = storage.load_checkpoint("CP").unwrap();
        assert_eq!(
            loaded.get::<String>("Key").unwrap_err().kind(),
            ErrorKind::Integrity
        );
    }

    // Verify that data signed with an unexpected key is detected properly.
    #[test]
    fn wrong_hmac_key() {
        let dir = create_test_directory_with_checkpoint_folders();
        let path = dir.path();

        let mut storage = HmacWrapper::wrap(FileStorage::open(&path).unwrap(), b"Secret Key1");
        let mut uncommitted = storage.create_checkpoint("CP").unwrap();
        uncommitted.put("Key", &"Test String".to_string()).unwrap();
        storage.commit_checkpoint(uncommitted).unwrap();

        let mut storage = HmacWrapper::wrap(FileStorage::open(&path).unwrap(), b"Secret Key2");

        let mut loaded = storage.load_checkpoint("CP").unwrap();
        assert_eq!(
            loaded.get::<String>("Key").unwrap_err().kind(),
            ErrorKind::Integrity
        );
    }
}