vios_app 0.1.1

//! Vault crypto: v2 headered format (Argon2id + AES-GCM), plus legacy readers.
//!
//! Formats supported (auto-detected):
//! - v2 (current): [MAGIC(4)=b"VLT2"][VER(1)=0x01][KDF(1)=0x02=Argon2id]
//!   [salt_len(1)=16][salt(16)][mem_kib(4 LE)][iters(4 LE)][lanes(1)]
//!   [nonce(12)] [ciphertext||tag]
//! - v1-raw (interim): [nonce(12)][ciphertext||tag], key = SHA256(password)  (no salt)
//! - legacy JSON: {"metadata": {...}, "data": "..."} PBKDF2-HMAC-SHA256
//!
//! AES mode: AES-256-GCM. AAD currently empty (future: bind metadata/filename).
//!
//! Note: we keep v1-raw read support for migrations, but *write* only v2.
//!
use aes_gcm::{
    aead::{Aead, KeyInit},
    Aes256Gcm, Nonce,
};
use argon2::{Algorithm, Argon2, Params, Version};
use base64::engine::general_purpose::STANDARD as B64;
use base64::Engine;
use rand::{rngs::OsRng, RngCore};
use serde::{Deserialize, Serialize};
use sha2::{Digest, Sha256};
use thiserror::Error;
use zeroize::Zeroize;

use std::io::{Read, Write};

const MAGIC_V2: &[u8; 4] = b"VLT2";
const VERSION_V2: u8 = 0x01;
// KDF ids (reserve 0x01 for PBKDF2 if ever needed)
const KDF_ARGON2ID: u8 = 0x02;

/// Default Argon2id params for desktop-class machines:
/// mem = 64 MiB, iters = 3, lanes = 1 (tune later & make configurable)
const ARGON_MEM_KIB_DEFAULT: u32 = 64 * 1024;
const ARGON_ITERS_DEFAULT: u32 = 3;
const ARGON_LANES_DEFAULT: u32 = 1;

#[derive(Debug, Error)]
pub enum VaultCryptoError {
    #[error("invalid header")]
    InvalidHeader,
    #[error("unsupported version")]
    UnsupportedVersion,
    #[error("unsupported kdf")]
    UnsupportedKdf,
    #[error("crypto failure")]
    CryptoFailure,
    #[error("io error: {0}")]
    Io(#[from] std::io::Error),
    #[error("format error: {0}")]
    Format(String),
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct LegacyMetadata {
    pub salt: String,
    pub iv: String,
    pub tag: String,
    #[serde(default)]
    pub iterations: Option<u32>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct LegacyJson {
    pub metadata: LegacyMetadata,
    pub data: String,
}

/// Encrypts plaintext bytes with password into **v2** header+AES-GCM.
/// Returns the full blob to write.
pub fn encrypt_v2(password: &str, plaintext: &[u8]) -> Result<Vec<u8>, VaultCryptoError> {
    // ---- KDF (Argon2id) ----
    let mut salt = [0u8; 16];
    OsRng.fill_bytes(&mut salt);

    let params = Params::new(
        ARGON_MEM_KIB_DEFAULT,
        ARGON_ITERS_DEFAULT,
        ARGON_LANES_DEFAULT,
        None,
    )
    .map_err(|_| VaultCryptoError::CryptoFailure)?;
    let argon2 = Argon2::new(Algorithm::Argon2id, Version::V0x13, params);

    // Argon2 crate likes hashing a password; we need raw key bytes (32B)
    // Workaround: hash the password into 32 bytes with Sha256 as input to Argon2 "password".
    // This avoids very large pass strings sitting around.
    let mut pw_buf = Sha256::digest(password.as_bytes()).to_vec();
    let mut key_material = [0u8; 32];
    argon2
        .hash_password_into(&pw_buf, &salt, &mut key_material)
        .map_err(|_| VaultCryptoError::CryptoFailure)?;

    pw_buf.zeroize();

    // ---- AES-GCM ----
    let cipher =
        Aes256Gcm::new_from_slice(&key_material).map_err(|_| VaultCryptoError::CryptoFailure)?;
    let mut nonce = [0u8; 12];
    OsRng.fill_bytes(&mut nonce);

    let ct = cipher
        .encrypt(Nonce::from_slice(&nonce), plaintext)
        .map_err(|_| VaultCryptoError::CryptoFailure)?;

    // Zeroize key material
    key_material.zeroize();

    // ---- Build header ----
    // [MAGIC(4)][VER(1)][KDF(1)][salt_len(1)=16][salt(16)][mem_kib(4)][iters(4)][lanes(1)][nonce(12)] [ct||tag]
    let mut out = Vec::with_capacity(4 + 1 + 1 + 1 + 16 + 4 + 4 + 1 + 12 + ct.len());
    out.extend_from_slice(MAGIC_V2);
    out.push(VERSION_V2);
    out.push(KDF_ARGON2ID);
    out.push(16u8);
    out.extend_from_slice(&salt);
    out.extend_from_slice(&ARGON_MEM_KIB_DEFAULT.to_le_bytes());
    out.extend_from_slice(&ARGON_ITERS_DEFAULT.to_le_bytes());
    out.push(ARGON_LANES_DEFAULT as u8);
    out.extend_from_slice(&nonce);
    out.extend_from_slice(&ct);

    Ok(out)
}

/// Decrypts any supported format (v2 headered, v1-raw, legacy JSON).
pub fn decrypt_auto(password: &str, blob: &[u8]) -> Result<Vec<u8>, VaultCryptoError> {
    // v2 header?
    if blob.len() >= 4 && &blob[0..4] == MAGIC_V2 {
        return decrypt_v2(password, blob);
    }

    // legacy JSON?
    if looks_like_json(blob) {
        return decrypt_legacy_json(password, blob);
    }

    // v1-raw fallback if >= 12 bytes
    if blob.len() >= 12 {
        return decrypt_v1_raw(password, blob);
    }

    Err(VaultCryptoError::InvalidHeader)
}

fn decrypt_v2(password: &str, blob: &[u8]) -> Result<Vec<u8>, VaultCryptoError> {
    // Parse header
    let mut idx = 0usize;

    if blob.len() < 4 + 1 + 1 + 1 {
        return Err(VaultCryptoError::InvalidHeader);
    }
    if &blob[idx..idx + 4] != MAGIC_V2 {
        return Err(VaultCryptoError::InvalidHeader);
    }
    idx += 4;

    let ver = blob[idx];
    idx += 1;
    if ver != VERSION_V2 {
        return Err(VaultCryptoError::UnsupportedVersion);
    }

    let kdf = blob[idx];
    idx += 1;
    if kdf != KDF_ARGON2ID {
        return Err(VaultCryptoError::UnsupportedKdf);
    }

    let salt_len = blob[idx] as usize;
    idx += 1;
    if salt_len != 16 || blob.len() < idx + salt_len {
        return Err(VaultCryptoError::InvalidHeader);
    }
    let salt = &blob[idx..idx + salt_len];
    idx += salt_len;

    if blob.len() < idx + 4 + 4 + 1 + 12 {
        return Err(VaultCryptoError::InvalidHeader);
    }
    let mem_kib = u32::from_le_bytes(blob[idx..idx + 4].try_into().unwrap());
    idx += 4;
    let iters = u32::from_le_bytes(blob[idx..idx + 4].try_into().unwrap());
    idx += 4;
    let lanes = blob[idx] as u32;
    idx += 1;

    let nonce = &blob[idx..idx + 12];
    idx += 12;

    let ct = &blob[idx..];

    // ---- KDF ----
    let params =
        Params::new(mem_kib, iters, lanes, None).map_err(|_| VaultCryptoError::CryptoFailure)?;
    let argon2 = Argon2::new(Algorithm::Argon2id, Version::V0x13, params);

    let mut pw_buf = Sha256::digest(password.as_bytes()).to_vec();
    let mut key_material = [0u8; 32];
    argon2
        .hash_password_into(&pw_buf, salt, &mut key_material)
        .map_err(|_| VaultCryptoError::CryptoFailure)?;
    pw_buf.zeroize();

    let cipher =
        Aes256Gcm::new_from_slice(&key_material).map_err(|_| VaultCryptoError::CryptoFailure)?;
    let pt = cipher
        .decrypt(Nonce::from_slice(nonce), ct)
        .map_err(|_| VaultCryptoError::CryptoFailure)?;

    key_material.zeroize();
    Ok(pt)
}

/// v1-raw: [nonce(12)][ct||tag], key = SHA256(password).  (read-only for migration)
fn decrypt_v1_raw(password: &str, blob: &[u8]) -> Result<Vec<u8>, VaultCryptoError> {
    if blob.len() < 12 + 16 {
        return Err(VaultCryptoError::Format("v1-raw too short".into()));
    }
    let (nonce, ct) = blob.split_at(12);

    let mut key = Sha256::digest(password.as_bytes()).to_vec();
    let cipher = Aes256Gcm::new_from_slice(&key).map_err(|_| VaultCryptoError::CryptoFailure)?;
    let pt = cipher
        .decrypt(Nonce::from_slice(nonce), ct)
        .map_err(|_| VaultCryptoError::CryptoFailure)?;
    // zeroize
    key.zeroize();
    Ok(pt)
}

/// Legacy JSON decrypt using PBKDF2-HMAC-SHA256.
/// Implemented elsewhere previously; this keeps API parity.
fn decrypt_legacy_json(password: &str, blob: &[u8]) -> Result<Vec<u8>, VaultCryptoError> {
    let s = std::str::from_utf8(blob).map_err(|_| VaultCryptoError::Format("bad utf8".into()))?;
    let parsed: LegacyJson = serde_json::from_str(s)
        .map_err(|e| VaultCryptoError::Format(format!("json parse: {e}")))?;

    let salt = B64
        .decode(parsed.metadata.salt.as_bytes())
        .map_err(|_| VaultCryptoError::Format("bad salt b64".into()))?;
    let iv = B64
        .decode(parsed.metadata.iv.as_bytes())
        .map_err(|_| VaultCryptoError::Format("bad iv b64".into()))?;
    let tag = B64
        .decode(parsed.metadata.tag.as_bytes())
        .map_err(|_| VaultCryptoError::Format("bad tag b64".into()))?;
    let data = B64
        .decode(parsed.data.as_bytes())
        .map_err(|_| VaultCryptoError::Format("bad data b64".into()))?;

    if iv.len() != 12 || tag.len() != 16 {
        return Err(VaultCryptoError::Format("iv/tag wrong len".into()));
    }

    let iters = parsed.metadata.iterations.unwrap_or(100_000);

    // PBKDF2-HMAC-SHA256
    let mut dk = [0u8; 32];
    pbkdf2::pbkdf2_hmac::<sha2::Sha256>(password.as_bytes(), &salt, iters, &mut dk);

    // assemble ct||tag
    let mut ct = Vec::with_capacity(data.len() + tag.len());
    ct.extend_from_slice(&data);
    ct.extend_from_slice(&tag);

    let cipher = Aes256Gcm::new_from_slice(&dk).map_err(|_| VaultCryptoError::CryptoFailure)?;
    let pt = cipher
        .decrypt(Nonce::from_slice(&iv), ct.as_ref())
        .map_err(|_| VaultCryptoError::CryptoFailure)?;

    dk.zeroize();
    Ok(pt)
}

fn looks_like_json(buf: &[u8]) -> bool {
    // cheap check: leading whitespace then '{'
    let first_non_ws = buf.iter().copied().find(|b| !b" \n\r\t".contains(b));
    matches!(first_non_ws, Some(b'{'))
}

/// Convenience: encrypt-and-write file (v2).  Caller zeroizes plaintext.
pub fn write_vault_v2_file(
    path: &std::path::Path,
    password: &str,
    plaintext: &[u8],
) -> Result<(), VaultCryptoError> {
    let blob = encrypt_v2(password, plaintext)?;
    let mut f = std::fs::File::create(path)?;
    f.write_all(&blob)?;
    Ok(())
}

/// Convenience: read-and-decrypt file (auto-detect).
pub fn read_vault_auto_file(
    path: &std::path::Path,
    password: &str,
) -> Result<Vec<u8>, VaultCryptoError> {
    let mut buf = Vec::new();
    std::fs::File::open(path)?.read_to_end(&mut buf)?;
    decrypt_auto(password, &buf)
}

// Re-export pbkdf2 for legacy without making it a public API dependency
mod pbkdf2 {
    pub use pbkdf2::pbkdf2_hmac;
}

// Public aliases to preserve older call sites
pub use encrypt_v2 as encrypt_vault;

use crate::vault::vault_structs::VaultPayload;

/// Read a vault file from disk, decrypt with auto-detect, and parse as VaultPayload (JSON).
pub fn decrypt_vault(path: &str, password: &str) -> Result<VaultPayload, VaultCryptoError> {
    let bytes = read_vault_auto_file(std::path::Path::new(path), password)?;
    let payload: VaultPayload = serde_json::from_slice(&bytes)
        .map_err(|e| VaultCryptoError::Format(format!("payload json: {e}")))?;
    Ok(payload)
}

/// Compatibility wrapper for older bins expecting encrypt_and_save_vault.
pub fn encrypt_and_save_vault(
    out_path: &str,
    password: &str,
    plaintext: &[u8],
) -> Result<(), VaultCryptoError> {
    write_vault_v2_file(std::path::Path::new(out_path), password, plaintext)
}

/// Get password from ENV (VIOS_VAULT_PASSWORD) or prompt (no echo).
pub fn must_get_password() -> String {
    if let Ok(pw) = std::env::var("VIOS_VAULT_PASSWORD") {
        return pw;
    }
    rpassword::prompt_password("Enter password: ").expect("failed to read password")
}