enigma-storage 0.0.1

use std::fs;
use std::path::{Path, PathBuf};

use blake3::Hasher;
use chacha20poly1305::aead::{Aead, Payload};
use chacha20poly1305::{Key, KeyInit, XChaCha20Poly1305, XNonce};
use rand::RngCore;
use zeroize::Zeroize;

use crate::codec::decode_envelope;
use crate::error::{EnigmaStorageError, Result};
#[cfg(any(feature = "provider-password", test))]
use crate::kdf::{derive_wrapping_key, KdfParams};
#[cfg(any(feature = "provider-password", test))]
use zeroize::Zeroizing;

pub trait KeyProvider: Send + Sync {
    fn get_or_create_master_key(&self) -> Result<MasterKey>;
    fn get_master_key(&self) -> Result<MasterKey>;
}

pub struct MasterKey(pub [u8; 32]);

impl MasterKey {
    pub fn new(bytes: [u8; 32]) -> Self {
        MasterKey(bytes)
    }

    pub fn as_bytes(&self) -> &[u8; 32] {
        &self.0
    }
}

impl Clone for MasterKey {
    fn clone(&self) -> Self {
        MasterKey(self.0)
    }
}

impl Drop for MasterKey {
    fn drop(&mut self) {
        self.0.zeroize();
    }
}

#[cfg(any(feature = "provider-file-sealed", test))]
pub struct FileSealedKeyProvider {
    root: PathBuf,
}

#[cfg(any(feature = "provider-file-sealed", test))]
impl FileSealedKeyProvider {
    pub fn new<P: AsRef<Path>>(root: P) -> Self {
        FileSealedKeyProvider {
            root: root.as_ref().to_path_buf(),
        }
    }

    fn salt_path(&self) -> PathBuf {
        self.root.join(".enigma_storage_salt")
    }

    fn sealed_path(&self) -> PathBuf {
        self.root.join(".enigma_storage_sealed_key")
    }

    fn ensure_root(&self) -> Result<()> {
        fs::create_dir_all(&self.root).map_err(|e| EnigmaStorageError::KeyProviderError(e.to_string()))
    }

    fn load_or_create_salt(&self) -> Result<[u8; 32]> {
        let path = self.salt_path();
        if path.exists() {
            let data = fs::read(&path).map_err(|e| EnigmaStorageError::KeyProviderError(e.to_string()))?;
            if data.len() != 32 {
                return Err(EnigmaStorageError::KeyProviderError("invalid salt length".into()));
            }
            let mut salt = [0u8; 32];
            salt.copy_from_slice(&data);
            return Ok(salt);
        }
        let mut salt = [0u8; 32];
        rand::thread_rng().fill_bytes(&mut salt);
        fs::write(&path, &salt).map_err(|e| EnigmaStorageError::KeyProviderError(e.to_string()))?;
        Ok(salt)
    }

    fn device_key(&self, salt: &[u8; 32]) -> [u8; 32] {
        let mut hasher = Hasher::new();
        hasher.update(b"enigma:devicekey:v1");
        hasher.update(salt);
        let hash = hasher.finalize();
        *hash.as_bytes()
    }

    fn seal(&self, master: &MasterKey, device_key: &[u8; 32]) -> Result<Vec<u8>> {
        let mut nonce = [0u8; 24];
        rand::thread_rng().fill_bytes(&mut nonce);
        let cipher = XChaCha20Poly1305::new(Key::from_slice(device_key));
        let ciphertext = cipher
            .encrypt(XNonce::from_slice(&nonce), Payload { msg: master.as_bytes(), aad: &[] })
            .map_err(|_| EnigmaStorageError::AeadError)?;
        let mut out = Vec::with_capacity(1 + nonce.len() + ciphertext.len());
        out.push(1);
        out.extend_from_slice(&nonce);
        out.extend_from_slice(&ciphertext);
        Ok(out)
    }

    fn unseal(&self, device_key: &[u8; 32], data: &[u8]) -> Result<MasterKey> {
        let (version, nonce, ciphertext) = decode_envelope(data)?;
        if version != 1 {
            return Err(EnigmaStorageError::UnsupportedVersion);
        }
        let cipher = XChaCha20Poly1305::new(Key::from_slice(device_key));
        let plaintext = cipher
            .decrypt(XNonce::from_slice(&nonce), Payload { msg: ciphertext.as_slice(), aad: &[] })
            .map_err(|_| EnigmaStorageError::AeadError)?;
        if plaintext.len() != 32 {
            return Err(EnigmaStorageError::InvalidKey);
        }
        let mut key = [0u8; 32];
        key.copy_from_slice(&plaintext);
        Ok(MasterKey(key))
    }
}

#[cfg(any(feature = "provider-file-sealed", test))]
impl KeyProvider for FileSealedKeyProvider {
    fn get_or_create_master_key(&self) -> Result<MasterKey> {
        self.ensure_root()?;
        let salt = self.load_or_create_salt()?;
        let device_key = self.device_key(&salt);
        let sealed_path = self.sealed_path();
        if sealed_path.exists() {
            return self.get_master_key();
        }
        let mut master_bytes = [0u8; 32];
        rand::thread_rng().fill_bytes(&mut master_bytes);
        let master = MasterKey(master_bytes);
        let wrapped = self.seal(&master, &device_key)?;
        fs::write(&sealed_path, wrapped).map_err(|e| EnigmaStorageError::KeyProviderError(e.to_string()))?;
        Ok(master)
    }

    fn get_master_key(&self) -> Result<MasterKey> {
        let salt = self.load_or_create_salt()?;
        let device_key = self.device_key(&salt);
        let data = fs::read(self.sealed_path()).map_err(|e| EnigmaStorageError::KeyProviderError(e.to_string()))?;
        self.unseal(&device_key, &data).map_err(|e| match e {
            EnigmaStorageError::AeadError => EnigmaStorageError::KeyProviderError("failed to unseal key".into()),
            other => other,
        })
    }
}

#[cfg(any(feature = "provider-password", test))]
pub struct PasswordKeyProvider {
    root: PathBuf,
    password: Zeroizing<Vec<u8>>,
    params: KdfParams,
}

#[cfg(any(feature = "provider-password", test))]
impl PasswordKeyProvider {
    pub fn new<P: AsRef<Path>>(root: P, password: impl Into<Vec<u8>>) -> Self {
        PasswordKeyProvider {
            root: root.as_ref().to_path_buf(),
            password: Zeroizing::new(password.into()),
            params: KdfParams::default(),
        }
    }

    fn salt_path(&self) -> PathBuf {
        self.root.join(".enigma_storage_pwd_salt")
    }

    fn wrapped_path(&self) -> PathBuf {
        self.root.join(".enigma_storage_pwd_wrapped_key")
    }

    fn ensure_root(&self) -> Result<()> {
        fs::create_dir_all(&self.root).map_err(|e| EnigmaStorageError::KeyProviderError(e.to_string()))
    }

    fn load_or_create_salt(&self) -> Result<Vec<u8>> {
        let path = self.salt_path();
        if path.exists() {
            let data = fs::read(&path).map_err(|e| EnigmaStorageError::KeyProviderError(e.to_string()))?;
            if data.len() < 8 {
                return Err(EnigmaStorageError::KeyProviderError("invalid salt".into()));
            }
            return Ok(data);
        }
        let mut salt = vec![0u8; 16];
        rand::thread_rng().fill_bytes(&mut salt);
        fs::write(&path, &salt).map_err(|e| EnigmaStorageError::KeyProviderError(e.to_string()))?;
        Ok(salt)
    }

    fn derive_key(&self, salt: &[u8]) -> Result<[u8; 32]> {
        if self.password.is_empty() {
            return Err(EnigmaStorageError::InvalidKey);
        }
        derive_wrapping_key(self.password.as_slice(), salt, &self.params)
    }
}

#[cfg(any(feature = "provider-password", test))]
impl KeyProvider for PasswordKeyProvider {
    fn get_or_create_master_key(&self) -> Result<MasterKey> {
        self.ensure_root()?;
        let salt = self.load_or_create_salt()?;
        let wrapping_key = self.derive_key(&salt)?;
        let wrapped_path = self.wrapped_path();
        if wrapped_path.exists() {
            return self.get_master_key();
        }
        let mut master = [0u8; 32];
        rand::thread_rng().fill_bytes(&mut master);
        let mut nonce = [0u8; 24];
        rand::thread_rng().fill_bytes(&mut nonce);
        let cipher = XChaCha20Poly1305::new(Key::from_slice(&wrapping_key));
        let ciphertext = cipher
            .encrypt(XNonce::from_slice(&nonce), Payload { msg: &master, aad: &[] })
            .map_err(|_| EnigmaStorageError::AeadError)?;
        let mut out = Vec::with_capacity(1 + nonce.len() + ciphertext.len());
        out.push(1);
        out.extend_from_slice(&nonce);
        out.extend_from_slice(&ciphertext);
        fs::write(&wrapped_path, out).map_err(|e| EnigmaStorageError::KeyProviderError(e.to_string()))?;
        Ok(MasterKey(master))
    }

    fn get_master_key(&self) -> Result<MasterKey> {
        let salt = self.load_or_create_salt()?;
        let wrapping_key = self.derive_key(&salt)?;
        let data = fs::read(self.wrapped_path()).map_err(|e| EnigmaStorageError::KeyProviderError(e.to_string()))?;
        let (version, nonce, ciphertext) = decode_envelope(&data)?;
        if version != 1 {
            return Err(EnigmaStorageError::UnsupportedVersion);
        }
        let cipher = XChaCha20Poly1305::new(Key::from_slice(&wrapping_key));
        let plaintext = cipher
            .decrypt(XNonce::from_slice(&nonce), Payload { msg: ciphertext.as_slice(), aad: &[] })
            .map_err(|_| EnigmaStorageError::KeyProviderError("invalid password".into()))?;
        if plaintext.len() != 32 {
            return Err(EnigmaStorageError::InvalidKey);
        }
        let mut key = [0u8; 32];
        key.copy_from_slice(&plaintext);
        Ok(MasterKey(key))
    }
}

pub use crate::platform::ForeignKeyProvider;
#[cfg(all(windows, feature = "provider-windows-dpapi"))]
pub use crate::platform::WindowsDpapiKeyProvider;
#[cfg(all(target_os = "macos", feature = "provider-macos-keychain"))]
pub use crate::platform::MacosKeychainKeyProvider;
#[cfg(all(target_os = "linux", feature = "provider-linux-secret-service"))]
pub use crate::platform::LinuxSecretServiceKeyProvider;