ERitors-Scribe-Desktop/src-tauri/src/crypto/key_manager.rs

use std::collections::HashMap;
use std::fs;
use std::path::PathBuf;

use aes::cipher::{BlockDecryptMut, BlockEncryptMut, KeyIvInit};
use base64::{engine::general_purpose::STANDARD as BASE64, Engine};
use rand::RngCore;
use serde::{Deserialize, Serialize};
use sha2::{Digest, Sha256};

use crate::error::{AppError, AppResult};

const SERVICE_NAME: &str = "com.eritors.scribe.desktop";
const IV_LENGTH: usize = 16;

type Aes256CbcEnc = cbc::Encryptor<aes::Aes256>;
type Aes256CbcDec = cbc::Decryptor<aes::Aes256>;

// ===== Secure vault: encrypted JSON file (like Electron's safeStorage) =====

#[derive(Serialize, Deserialize, Default)]
struct SecureVault {
    token: Option<String>,
    encryption_keys: HashMap<String, String>,
    last_user_id: Option<String>,
    pin_hashes: HashMap<String, String>,
}

fn vault_path() -> PathBuf {
    dirs_next::data_dir()
        .unwrap_or_else(|| PathBuf::from("."))
        .join(SERVICE_NAME)
        .join("secure-config.json")
}

fn derive_machine_key() -> [u8; 32] {
    let mut hasher = Sha256::new();
    hasher.update(SERVICE_NAME.as_bytes());
    if let Ok(name) = hostname::get() {
        hasher.update(name.as_encoded_bytes());
    }
    if let Some(dir) = dirs_next::home_dir() {
        hasher.update(dir.to_string_lossy().as_bytes());
    }
    hasher.finalize().into()
}

fn encrypt_vault(data: &[u8], key: &[u8; 32]) -> AppResult<Vec<u8>> {
    let mut iv = [0u8; IV_LENGTH];
    rand::rng().fill_bytes(&mut iv);
    let block_size = 16;
    let padding_len = block_size - (data.len() % block_size);
    let mut buf = Vec::with_capacity(data.len() + padding_len);
    buf.extend_from_slice(data);
    buf.extend(std::iter::repeat(padding_len as u8).take(padding_len));
    let padded_len = buf.len();
    let cipher = Aes256CbcEnc::new(key.into(), &iv.into());
    cipher.encrypt_padded_mut::<aes::cipher::block_padding::NoPadding>(&mut buf, padded_len)
        .map_err(|e| AppError::Encryption(format!("Vault encrypt failed: {}", e)))?;
    let mut result = Vec::with_capacity(IV_LENGTH + buf.len());
    result.extend_from_slice(&iv);
    result.extend_from_slice(&buf);
    Ok(result)
}

fn decrypt_vault(data: &[u8], key: &[u8; 32]) -> AppResult<Vec<u8>> {
    if data.len() < IV_LENGTH + 1 {
        return Err(AppError::Encryption("Vault data too short".into()));
    }
    let iv: [u8; 16] = data[..IV_LENGTH].try_into()
        .map_err(|_| AppError::Encryption("Invalid IV".into()))?;
    let mut buf = data[IV_LENGTH..].to_vec();
    let cipher = Aes256CbcDec::new(key.into(), &iv.into());
    let decrypted = cipher.decrypt_padded_mut::<aes::cipher::block_padding::Pkcs7>(&mut buf)
        .map_err(|e| AppError::Encryption(format!("Vault decrypt failed: {}", e)))?;
    Ok(decrypted.to_vec())
}

fn read_vault() -> SecureVault {
    let path = vault_path();
    let key = derive_machine_key();
    match fs::read_to_string(&path) {
        Ok(content) => {
            if let Ok(raw) = BASE64.decode(content.trim()) {
                if let Ok(decrypted) = decrypt_vault(&raw, &key) {
                    if let Ok(vault) = serde_json::from_slice::<SecureVault>(&decrypted) {
                        return vault;
                    }
                }
            }
            SecureVault::default()
        }
        Err(_) => SecureVault::default(),
    }
}

fn write_vault(vault: &SecureVault) -> AppResult<()> {
    let path = vault_path();
    if let Some(parent) = path.parent() {
        fs::create_dir_all(parent).map_err(|e| AppError::Internal(format!("Failed to create vault dir: {}", e)))?;
    }
    let json = serde_json::to_string(vault)
        .map_err(|e| AppError::Internal(format!("Failed to serialize vault: {}", e)))?;
    let key = derive_machine_key();
    let encrypted = encrypt_vault(json.as_bytes(), &key)?;
    let encoded = BASE64.encode(&encrypted);
    fs::write(&path, encoded).map_err(|e| AppError::Internal(format!("Failed to write vault: {}", e)))
}

// ===== Public API (same signatures as before) =====

pub fn get_user_encryption_key(user_id: &str) -> AppResult<String> {
    let vault = read_vault();
    vault.encryption_keys.get(user_id).cloned()
        .ok_or_else(|| AppError::Keyring(format!("No encryption key for user {}", user_id)))
}

pub fn set_user_encryption_key(user_id: &str, encryption_key: &str) -> AppResult<()> {
    let mut vault = read_vault();
    vault.encryption_keys.insert(user_id.to_string(), encryption_key.to_string());
    write_vault(&vault)
}

pub fn has_user_encryption_key(user_id: &str) -> bool {
    let vault = read_vault();
    vault.encryption_keys.contains_key(user_id)
}


pub fn get_token() -> AppResult<Option<String>> {
    let vault = read_vault();
    Ok(vault.token)
}

pub fn set_token(token: &str) -> AppResult<()> {
    let mut vault = read_vault();
    vault.token = Some(token.to_string());
    write_vault(&vault)
}

pub fn remove_token() -> AppResult<()> {
    let mut vault = read_vault();
    vault.token = None;
    write_vault(&vault)
}

pub fn set_pin_hash(user_id: &str, pin_hash: &str) -> AppResult<()> {
    let mut vault = read_vault();
    vault.pin_hashes.insert(user_id.to_string(), pin_hash.to_string());
    write_vault(&vault)
}

pub fn get_pin_hash(user_id: &str) -> AppResult<Option<String>> {
    let vault = read_vault();
    Ok(vault.pin_hashes.get(user_id).cloned())
}

pub fn set_last_user_id(user_id: &str) -> AppResult<()> {
    let mut vault = read_vault();
    vault.last_user_id = Some(user_id.to_string());
    write_vault(&vault)
}

pub fn get_last_user_id() -> AppResult<Option<String>> {
    let vault = read_vault();
    Ok(vault.last_user_id)
}