omy-core 0.0.1

//! 生成用于跨实现验证的样本文件。
//!
//! Rust 在此写出各种配置的 omy 文件，供其它实现（如 `docs/research/reference/omy_ref.py`）
//! 独立读取。这是与 Rust 自身解码路径**完全无关**的验证通道：只靠 Rust 自测，
//! 无法发现「编码与解码两侧同样误解了规范」这类偏差。
//!
//! 用法：`cargo run --example gen_samples -- <输出目录>`

// 这是生成测试数据的工具，算术用于构造样本内容而非处理不可信输入，
// 因此局部放宽相关 lint，不影响库代码的严格检查。
#![allow(clippy::arithmetic_side_effects, clippy::integer_division)]

use omy_core::crypto::{Argon2Params, CipherId, Kek};
use omy_core::file::{EncryptOptions, RandomMaterial, encrypt};
use omy_core::header::MIN_CHUNK_SIZE;
use std::io::Write as _;

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let out_dir = std::env::args().nth(1).ok_or("用法: gen_samples <输出目录>")?;
    std::fs::create_dir_all(&out_dir)?;

    // 固定 salt，让 Python 侧能用同样的密码派生出同样的 KEK
    let salt = [0x7Au8; 16];
    let params = Argon2Params::TEST_WEAK;

    struct Case {
        name: &'static str,
        password: &'static str,
        filename: &'static str,
        plaintext: Vec<u8>,
        compress: bool,
        cipher: CipherId,
        chunk_size: u32,
        extra_passwords: Vec<&'static str>,
    }

    let cases = vec![
        Case {
            name: "simple",
            password: "rust-pw-1",
            filename: "hello.txt",
            plaintext: b"Generated by the Rust implementation.".to_vec(),
            compress: false,
            cipher: CipherId::ChaCha20Poly1305,
            chunk_size: MIN_CHUNK_SIZE,
            extra_passwords: vec![],
        },
        Case {
            name: "empty",
            password: "rust-pw-2",
            filename: "nothing.bin",
            plaintext: Vec::new(),
            compress: false,
            cipher: CipherId::ChaCha20Poly1305,
            chunk_size: MIN_CHUNK_SIZE,
            extra_passwords: vec![],
        },
        Case {
            name: "multichunk",
            password: "rust-pw-3",
            filename: "large-pattern.dat",
            // 2.5 块，确保存在非满的末块
            plaintext: (0..(MIN_CHUNK_SIZE as usize * 5 / 2))
                .map(|i| u8::try_from((i * 37 + 11) % 256).unwrap_or(0))
                .collect(),
            compress: false,
            cipher: CipherId::ChaCha20Poly1305,
            chunk_size: MIN_CHUNK_SIZE,
            extra_passwords: vec![],
        },
        Case {
            name: "compressed",
            password: "rust-pw-4",
            filename: "repetitive.log",
            plaintext: "omy compression test line\n".repeat(500).into_bytes(),
            compress: true,
            cipher: CipherId::ChaCha20Poly1305,
            chunk_size: MIN_CHUNK_SIZE,
            extra_passwords: vec![],
        },
        Case {
            name: "aesgcm",
            password: "rust-pw-5",
            filename: "aes-payload.bin",
            plaintext: b"encrypted with AES-256-GCM instead of ChaCha".to_vec(),
            compress: false,
            cipher: CipherId::Aes256Gcm,
            chunk_size: MIN_CHUNK_SIZE,
            extra_passwords: vec![],
        },
        Case {
            name: "multipassword",
            password: "rust-pw-6a",
            filename: "shared-secret.txt",
            plaintext: b"three passwords can open this file".to_vec(),
            compress: false,
            cipher: CipherId::ChaCha20Poly1305,
            chunk_size: MIN_CHUNK_SIZE,
            extra_passwords: vec!["rust-pw-6b", "rust-pw-6c"],
        },
        Case {
            name: "unicode-name",
            password: "rust-pw-7",
            filename: "中文文件名-测试.txt",
            plaintext: "内容也是中文的，用于验证 UTF-8 处理。".as_bytes().to_vec(),
            compress: false,
            cipher: CipherId::ChaCha20Poly1305,
            chunk_size: MIN_CHUNK_SIZE,
            extra_passwords: vec![],
        },
    ];

    let mut manifest = Vec::new();

    for c in &cases {
        let mut pws = vec![c.password];
        pws.extend(c.extra_passwords.iter());

        let keks: Vec<Kek> = pws
            .iter()
            .map(|p| Kek::from_password(p.as_bytes(), &salt, params))
            .collect::<Result<_, _>>()?;

        let opts = EncryptOptions {
            filename: Some(c.filename.to_owned()),
            compress: c.compress,
            zstd_level: 3,
            chunk_size: c.chunk_size,
            cipher: c.cipher,
            argon2: params,
            write_content_hash: true,
            ..EncryptOptions::default()
        };

        let enc = encrypt(&c.plaintext, &keks, &salt, &opts, &RandomMaterial::generate())?;

        let path = std::path::Path::new(&out_dir).join(format!("{}.omy", c.name));
        std::fs::write(&path, &enc.bytes)?;

        // 明文另存一份，供 Python 侧比对
        let plain_path = std::path::Path::new(&out_dir).join(format!("{}.plain", c.name));
        std::fs::write(&plain_path, &c.plaintext)?;

        manifest.push(format!(
            r#"{{"name":"{}","file":"{}.omy","plain":"{}.plain","passwords":[{}],"filename":{},"plaintext_size":{},"compressed":{},"cipher":{},"chunk_size":{},"n_chunks":{},"header_len":{},"total_size":{}}}"#,
            c.name,
            c.name,
            c.name,
            pws.iter().map(|p| format!("\"{p}\"")).collect::<Vec<_>>().join(","),
            serde_json_escape(c.filename),
            c.plaintext.len(),
            c.compress,
            c.cipher as u8,
            c.chunk_size,
            enc.header.n_chunks(),
            enc.header.header_len,
            enc.bytes.len(),
        ));

        println!(
            "wrote {:<16} {:>8} bytes  chunks={} cipher={}",
            format!("{}.omy", c.name),
            enc.bytes.len(),
            enc.header.n_chunks(),
            c.cipher as u8
        );
    }

    let manifest_json = format!(
        "{{\n  \"vault_salt\": \"{}\",\n  \"argon2\": {{\"m_kib\": {}, \"t\": {}, \"p\": {}}},\n  \"cases\": [\n    {}\n  ]\n}}\n",
        hex_encode(&salt),
        params.m_kib,
        params.t,
        params.p,
        manifest.join(",\n    ")
    );
    let mf = std::path::Path::new(&out_dir).join("manifest.json");
    let mut f = std::fs::File::create(&mf)?;
    f.write_all(manifest_json.as_bytes())?;

    println!("\nmanifest written to {}", mf.display());
    Ok(())
}

fn hex_encode(b: &[u8]) -> String {
    b.iter().map(|x| format!("{x:02x}")).collect()
}

/// 最小化的 JSON 字符串转义，够用于文件名。
fn serde_json_escape(s: &str) -> String {
    let mut out = String::with_capacity(s.len() + 2);
    out.push('"');
    for ch in s.chars() {
        match ch {
            '"' => out.push_str("\\\""),
            '\\' => out.push_str("\\\\"),
            '\n' => out.push_str("\\n"),
            '\r' => out.push_str("\\r"),
            '\t' => out.push_str("\\t"),
            c if (c as u32) < 0x20 => out.push_str(&format!("\\u{:04x}", c as u32)),
            c => out.push(c),
        }
    }
    out.push('"');
    out
}