src/crypto_pack.rs

use crate::storage::{self, PackedDocument, ThreadRecord};
use aes_gcm::{
    aead::{Aead, Payload},
    Aes256Gcm, KeyInit, Nonce,
};
use anyhow::{bail, Context, Result};
use base64::{engine::general_purpose::STANDARD as B64, Engine as _};
use chrono::Utc;
use pbkdf2::pbkdf2_hmac;
use rand::{rngs::OsRng, RngCore};
use serde::{Deserialize, Serialize};
use serde_json::{json, Value};
use sha2::Sha256;
use std::path::Path;

const FORMAT: &str = "jackailocal-encrypted-pack";
const VERSION: u32 = 1;
const PBKDF2_ITERATIONS: u32 = 210_000;
const AAD: &[u8] = b"JackAILocal encrypted transfer pack v1";

#[derive(Debug, Clone, Deserialize)]
pub struct ExportOptions {
    pub passphrase: String,
    pub include_threads: Option<bool>,
    pub include_documents: Option<bool>,
    pub include_settings: Option<bool>,
}

#[derive(Debug, Clone, Deserialize)]
pub struct ImportOptions {
    pub passphrase: String,
    pub data_base64: String,
    pub mode: Option<String>,
}

#[derive(Debug, Clone, Serialize)]
pub struct ImportResult {
    pub threads_imported: usize,
    pub documents_imported: usize,
    pub settings_imported: bool,
}

#[derive(Debug, Clone, Deserialize, Serialize)]
struct PlainPack {
    product: String,
    format_version: u32,
    exported_at: String,
    threads: Vec<ThreadRecord>,
    documents: Vec<PackedDocument>,
    settings: Option<Value>,
}

#[derive(Debug, Clone, Deserialize, Serialize)]
struct EncryptedEnvelope {
    format: String,
    version: u32,
    cipher: String,
    kdf: String,
    iterations: u32,
    salt_base64: String,
    nonce_base64: String,
    ciphertext_base64: String,
}

pub fn export_pack(root: &Path, options: &ExportOptions) -> Result<Vec<u8>> {
    validate_passphrase(&options.passphrase)?;
    let plain = PlainPack {
        product: "JackAILocal".to_string(),
        format_version: VERSION,
        exported_at: Utc::now().to_rfc3339(),
        threads: if options.include_threads.unwrap_or(true) {
            storage::read_all_threads(root)?
        } else {
            Vec::new()
        },
        documents: if options.include_documents.unwrap_or(true) {
            storage::list_documents_with_content(root)?
        } else {
            Vec::new()
        },
        settings: if options.include_settings.unwrap_or(false) {
            Some(sanitized_settings(storage::read_settings(root)))
        } else {
            None
        },
    };
    let plaintext = serde_json::to_vec(&plain)?;

    let mut salt = [0_u8; 16];
    let mut nonce = [0_u8; 12];
    OsRng.fill_bytes(&mut salt);
    OsRng.fill_bytes(&mut nonce);
    let key = derive_key(options.passphrase.as_bytes(), &salt, PBKDF2_ITERATIONS);
    let cipher = Aes256Gcm::new_from_slice(&key).context("cannot initialize AES-256-GCM")?;
    let ciphertext = cipher
        .encrypt(
            Nonce::from_slice(&nonce),
            Payload {
                msg: &plaintext,
                aad: AAD,
            },
        )
        .map_err(|_| anyhow::anyhow!("pack encryption failed"))?;

    let envelope = EncryptedEnvelope {
        format: FORMAT.to_string(),
        version: VERSION,
        cipher: "AES-256-GCM".to_string(),
        kdf: "PBKDF2-HMAC-SHA256".to_string(),
        iterations: PBKDF2_ITERATIONS,
        salt_base64: B64.encode(salt),
        nonce_base64: B64.encode(nonce),
        ciphertext_base64: B64.encode(ciphertext),
    };
    Ok(serde_json::to_vec_pretty(&envelope)?)
}

pub fn import_pack(root: &Path, options: &ImportOptions) -> Result<ImportResult> {
    validate_passphrase(&options.passphrase)?;
    let envelope_bytes = B64
        .decode(options.data_base64.trim())
        .context("pack data is not valid base64")?;
    let envelope: EncryptedEnvelope =
        serde_json::from_slice(&envelope_bytes).context("pack envelope is not valid JSON")?;
    if envelope.format != FORMAT || envelope.version != VERSION {
        bail!("unsupported JackAILocal pack format or version");
    }
    if envelope.cipher != "AES-256-GCM" || envelope.kdf != "PBKDF2-HMAC-SHA256" {
        bail!("unsupported pack cryptography");
    }
    if envelope.iterations < 100_000 {
        bail!("pack key derivation settings are too weak");
    }
    let salt = B64
        .decode(&envelope.salt_base64)
        .context("invalid pack salt")?;
    let nonce = B64
        .decode(&envelope.nonce_base64)
        .context("invalid pack nonce")?;
    let ciphertext = B64
        .decode(&envelope.ciphertext_base64)
        .context("invalid pack ciphertext")?;
    if salt.len() != 16 || nonce.len() != 12 {
        bail!("invalid pack salt or nonce length");
    }
    let key = derive_key(options.passphrase.as_bytes(), &salt, envelope.iterations);
    let cipher = Aes256Gcm::new_from_slice(&key).context("cannot initialize AES-256-GCM")?;
    let plaintext = cipher
        .decrypt(
            Nonce::from_slice(&nonce),
            Payload {
                msg: &ciphertext,
                aad: AAD,
            },
        )
        .map_err(|_| anyhow::anyhow!("wrong passphrase or corrupted pack"))?;
    let pack: PlainPack =
        serde_json::from_slice(&plaintext).context("decrypted pack is invalid")?;
    if pack.product != "JackAILocal" || pack.format_version != VERSION {
        bail!("decrypted data is not a compatible JackAILocal pack");
    }

    let replace = options.mode.as_deref() == Some("replace");
    let threads_imported = storage::import_threads(root, &pack.threads, replace)?;
    let mut documents_imported = 0;
    for document in pack.documents {
        storage::save_document(root, &document.name, &document.content)?;
        documents_imported += 1;
    }
    let settings_imported = if let Some(settings) = pack.settings {
        storage::write_settings(root, &sanitized_settings(settings))?;
        true
    } else {
        false
    };
    Ok(ImportResult {
        threads_imported,
        documents_imported,
        settings_imported,
    })
}

fn validate_passphrase(passphrase: &str) -> Result<()> {
    if passphrase.chars().count() < 8 {
        bail!("passphrase must contain at least 8 characters");
    }
    Ok(())
}

fn derive_key(passphrase: &[u8], salt: &[u8], iterations: u32) -> [u8; 32] {
    let mut key = [0_u8; 32];
    pbkdf2_hmac::<Sha256>(passphrase, salt, iterations, &mut key);
    key
}

fn sanitized_settings(mut settings: Value) -> Value {
    if !settings.is_object() {
        settings = json!({});
    }
    if let Some(object) = settings.as_object_mut() {
        object.insert("phone_access_enabled".to_string(), json!(false));
        object.remove("phone_access_token");
    }
    settings
}