oth_rvault 0.4.0

use crate::{
    keys::{check_key_conflicts, is_encryptable},
    typedetect::{FileFormat, detect_format},
    types::{ProcessHandling, TypedValue, ValueType},
    vaultfunc,
};

/// Encrypt all (or a filtered subset of) scalar values in a JSON or YAML file.
///
/// - `input`:    path to the input file
/// - `output`:   path to write the result, or `"stdout"` to print to stdout
/// - `password`: vault password
/// - `keys`:     dot-notation key paths to encrypt; empty means all encryptable values
/// - `level`:    if `Some(n)`, encrypt every node exactly at depth `n` (1 = top-level).
///               Nodes in branches shallower than `n` are left untouched.
///               `Some(0)` encrypts the entire YAML file as a single ansible-vault blob,
///               producing output compatible with `ansible-vault encrypt`.
///               Mutually exclusive with a non-empty `keys` slice.
pub fn encrypt(
    input: &str,
    output: &str,
    password: &str,
    keys: &[String],
    level: Option<u32>,
) -> Result<(), Box<dyn std::error::Error>> {
    if let Some(0) = level {
        let content = std::fs::read_to_string(input)
            .map_err(|e| format!("error reading '{}': {}", input, e))?;
        let encrypted = vaultfunc::encrypt(&TypedValue::String(content), password)?;
        if output == "stdout" {
            print!("{}", encrypted);
        } else {
            std::fs::write(output, &encrypted)
                .map_err(|e| format!("error writing '{}': {}", output, e))?;
        }
        return Ok(());
    }
    check_key_conflicts(keys)?;
    let format = detect_format(input)?;
    let computed_keys;
    let effective_keys: &[String] = if let Some(lvl) = level {
        computed_keys = match format {
            FileFormat::Json => crate::json::collect_paths_at_level(input, lvl)?,
            FileFormat::Yaml => crate::yaml::collect_paths_at_level(input, lvl)?,
            FileFormat::Vault => return Err("input file is already vault-encrypted".into()),
        };
        &computed_keys
    } else {
        keys
    };
    let mut processor = make_encrypt_processor(password);
    match format {
        FileFormat::Json => crate::json::process_file(input, output, effective_keys, &mut processor),
        FileFormat::Yaml => crate::yaml::process_file(input, output, effective_keys, &mut processor),
        FileFormat::Vault => Err("input file is already vault-encrypted".into()),
    }
}

pub(crate) fn make_encrypt_processor(
    password: &str,
) -> impl FnMut(
    TypedValue,
    ValueType,
    &str,
) -> Result<(TypedValue, ValueType, ProcessHandling), Box<dyn std::error::Error>>
+ '_ {
    move |typed, vt, key_path| {
        if !is_encryptable(vt) {
            return Ok((typed, vt, ProcessHandling::Skip));
        }
        match vaultfunc::encrypt(&typed, password) {
            Ok(encrypted) => Ok((
                TypedValue::String(encrypted),
                ValueType::String,
                ProcessHandling::Process,
            )),
            Err(e) => {
                eprintln!("error encrypting key '{}': {}", key_path, e);
                Ok((typed, vt, ProcessHandling::Skip))
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::decrypt::decrypt;

    fn resources(filename: &str) -> String {
        format!(
            "{}/resources/tests/{}",
            env!("CARGO_MANIFEST_DIR"),
            filename
        )
    }

    // --- sub-node tests ---

    #[test]
    fn test_encrypt_subnode_object_json_roundtrip() {
        let input = resources("example.json");
        let encrypted = tempfile::NamedTempFile::new().unwrap();
        let decrypted = tempfile::NamedTempFile::new().unwrap();
        let keys = vec!["second".to_string()];

        encrypt(&input, encrypted.path().to_str().unwrap(), "test999", &keys, None).unwrap();
        decrypt(encrypted.path().to_str().unwrap(), decrypted.path().to_str().unwrap(), "test999", &keys).unwrap();

        let got: serde_json::Value = serde_json::from_str(&std::fs::read_to_string(decrypted.path()).unwrap()).unwrap();
        let expected: serde_json::Value = serde_json::from_str(&std::fs::read_to_string(&input).unwrap()).unwrap();
        assert_eq!(got, expected, "JSON object sub-node roundtrip failed");
    }

    #[test]
    fn test_encrypt_subnode_object_json_only_target_encrypted() {
        let input = resources("example.json");
        let encrypted = tempfile::NamedTempFile::new().unwrap();
        let keys = vec!["second".to_string()];

        encrypt(&input, encrypted.path().to_str().unwrap(), "test999", &keys, None).unwrap();

        let enc: serde_json::Value = serde_json::from_str(&std::fs::read_to_string(encrypted.path()).unwrap()).unwrap();

        // "second" must now be a vault-encrypted string
        assert!(enc["second"].is_string(), "second should be a string after sub-node encrypt");
        assert!(enc["second"].as_str().unwrap().contains("$ANSIBLE_VAULT"), "second should be vault-encrypted");

        // other top-level keys must remain objects
        assert!(enc["first"].is_object(), "first should still be an object");
        assert!(enc["third"].is_object(), "third should still be an object");
        assert!(enc["fourth"].is_object(), "fourth should still be an object");
    }

    #[test]
    fn test_encrypt_subnode_array_json_roundtrip() {
        let input = resources("example.json");
        let encrypted = tempfile::NamedTempFile::new().unwrap();
        let decrypted = tempfile::NamedTempFile::new().unwrap();
        let keys = vec!["fourth.list".to_string()];

        encrypt(&input, encrypted.path().to_str().unwrap(), "test999", &keys, None).unwrap();
        decrypt(encrypted.path().to_str().unwrap(), decrypted.path().to_str().unwrap(), "test999", &keys).unwrap();

        let got: serde_json::Value = serde_json::from_str(&std::fs::read_to_string(decrypted.path()).unwrap()).unwrap();
        let expected: serde_json::Value = serde_json::from_str(&std::fs::read_to_string(&input).unwrap()).unwrap();
        assert_eq!(got, expected, "JSON array sub-node roundtrip failed");
    }

    #[test]
    fn test_encrypt_subnode_array_json_only_target_encrypted() {
        let input = resources("example.json");
        let encrypted = tempfile::NamedTempFile::new().unwrap();
        let keys = vec!["fourth.list".to_string()];

        encrypt(&input, encrypted.path().to_str().unwrap(), "test999", &keys, None).unwrap();

        let enc: serde_json::Value = serde_json::from_str(&std::fs::read_to_string(encrypted.path()).unwrap()).unwrap();

        // "fourth.list" must now be a vault-encrypted string
        assert!(enc["fourth"]["list"].is_string(), "fourth.list should be a string after sub-node encrypt");
        assert!(enc["fourth"]["list"].as_str().unwrap().contains("$ANSIBLE_VAULT"), "fourth.list should be vault-encrypted");

        // sibling keys within "fourth" must remain untouched — "fourth" is still an object
        assert!(enc["fourth"].is_object());
    }

    #[test]
    fn test_encrypt_subnode_object_yaml_roundtrip() {
        let input = resources("example.yaml");
        let encrypted = tempfile::NamedTempFile::new().unwrap();
        let decrypted = tempfile::NamedTempFile::new().unwrap();
        let keys = vec!["second".to_string()];

        encrypt(&input, encrypted.path().to_str().unwrap(), "test999", &keys, None).unwrap();
        decrypt(encrypted.path().to_str().unwrap(), decrypted.path().to_str().unwrap(), "test999", &keys).unwrap();

        let got: serde_yaml::Value = serde_yaml::from_str(&std::fs::read_to_string(decrypted.path()).unwrap()).unwrap();
        let expected: serde_yaml::Value = serde_yaml::from_str(&std::fs::read_to_string(&input).unwrap()).unwrap();
        assert_eq!(got, expected, "YAML object sub-node roundtrip failed");
    }

    #[test]
    fn test_encrypt_subnode_object_yaml_only_target_encrypted() {
        let input = resources("example.yaml");
        let encrypted = tempfile::NamedTempFile::new().unwrap();
        let keys = vec!["second".to_string()];

        encrypt(&input, encrypted.path().to_str().unwrap(), "test999", &keys, None).unwrap();

        let enc: serde_yaml::Value = serde_yaml::from_str(&std::fs::read_to_string(encrypted.path()).unwrap()).unwrap();

        // "second" must now be a vault-encrypted string
        let second = &enc["second"];
        assert!(second.is_string(), "second should be a string after sub-node encrypt");
        assert!(second.as_str().unwrap().contains("$ANSIBLE_VAULT"), "second should be vault-encrypted");

        // other top-level keys must remain mappings
        assert!(enc["first"].is_mapping(), "first should still be a mapping");
        assert!(enc["third"].is_mapping(), "third should still be a mapping");
        assert!(enc["fourth"].is_mapping(), "fourth should still be a mapping");
    }

    #[test]
    fn test_encrypt_subnode_array_yaml_roundtrip() {
        let input = resources("example.yaml");
        let encrypted = tempfile::NamedTempFile::new().unwrap();
        let decrypted = tempfile::NamedTempFile::new().unwrap();
        let keys = vec!["fourth.list".to_string()];

        encrypt(&input, encrypted.path().to_str().unwrap(), "test999", &keys, None).unwrap();
        decrypt(encrypted.path().to_str().unwrap(), decrypted.path().to_str().unwrap(), "test999", &keys).unwrap();

        let got: serde_yaml::Value = serde_yaml::from_str(&std::fs::read_to_string(decrypted.path()).unwrap()).unwrap();
        let expected: serde_yaml::Value = serde_yaml::from_str(&std::fs::read_to_string(&input).unwrap()).unwrap();
        assert_eq!(got, expected, "YAML array sub-node roundtrip failed");
    }

    #[test]
    fn test_key_conflict_rejected() {
        let input = resources("example.json");
        let out = tempfile::NamedTempFile::new().unwrap();
        let keys = vec!["second".to_string(), "second.a".to_string()];

        let result = encrypt(&input, out.path().to_str().unwrap(), "test999", &keys, None);
        assert!(result.is_err(), "conflicting keys should produce an error");
        let msg = result.unwrap_err().to_string();
        assert!(msg.contains("second"), "error should name the conflicting key");
    }

    // --- level tests ---

    #[test]
    fn test_encrypt_level1_json_roundtrip() {
        // level 1: all top-level keys become encrypted sub-node blobs
        let input = resources("example.json");
        let encrypted = tempfile::NamedTempFile::new().unwrap();
        let decrypted = tempfile::NamedTempFile::new().unwrap();

        encrypt(&input, encrypted.path().to_str().unwrap(), "test999", &[], Some(1)).unwrap();

        // after encryption every top-level value must be a vault string
        let enc: serde_json::Value = serde_json::from_str(&std::fs::read_to_string(encrypted.path()).unwrap()).unwrap();
        for key in ["first", "second", "third", "fourth"] {
            let v = &enc[key];
            assert!(v.is_string() && v.as_str().unwrap().contains("$ANSIBLE_VAULT"), "{key} should be vault-encrypted");
        }

        // collect the keys that were encrypted so decrypt can restore them
        let keys: Vec<String> = ["first", "second", "third", "fourth"].iter().map(|s| s.to_string()).collect();
        decrypt(encrypted.path().to_str().unwrap(), decrypted.path().to_str().unwrap(), "test999", &keys).unwrap();

        let got: serde_json::Value = serde_json::from_str(&std::fs::read_to_string(decrypted.path()).unwrap()).unwrap();
        let expected: serde_json::Value = serde_json::from_str(&std::fs::read_to_string(&input).unwrap()).unwrap();
        assert_eq!(got, expected, "JSON level-1 roundtrip failed");
    }

    #[test]
    fn test_encrypt_level1_yaml_roundtrip() {
        let input = resources("example.yaml");
        let encrypted = tempfile::NamedTempFile::new().unwrap();
        let decrypted = tempfile::NamedTempFile::new().unwrap();

        encrypt(&input, encrypted.path().to_str().unwrap(), "test999", &[], Some(1)).unwrap();

        let enc: serde_yaml::Value = serde_yaml::from_str(&std::fs::read_to_string(encrypted.path()).unwrap()).unwrap();
        for key in ["first", "second", "third", "fourth"] {
            let v = &enc[key];
            assert!(v.is_string() && v.as_str().unwrap().contains("$ANSIBLE_VAULT"), "{key} should be vault-encrypted");
        }

        let keys: Vec<String> = ["first", "second", "third", "fourth"].iter().map(|s| s.to_string()).collect();
        decrypt(encrypted.path().to_str().unwrap(), decrypted.path().to_str().unwrap(), "test999", &keys).unwrap();

        let got: serde_yaml::Value = serde_yaml::from_str(&std::fs::read_to_string(decrypted.path()).unwrap()).unwrap();
        let expected: serde_yaml::Value = serde_yaml::from_str(&std::fs::read_to_string(&input).unwrap()).unwrap();
        assert_eq!(got, expected, "YAML level-1 roundtrip failed");
    }

    #[test]
    fn test_encrypt_level2_json_roundtrip() {
        // level 2: second-level scalars are leaf-encrypted, second-level objects become sub-node blobs
        let input = resources("example.json");
        let encrypted = tempfile::NamedTempFile::new().unwrap();
        let decrypted = tempfile::NamedTempFile::new().unwrap();

        encrypt(&input, encrypted.path().to_str().unwrap(), "test999", &[], Some(2)).unwrap();
        decrypt(encrypted.path().to_str().unwrap(), decrypted.path().to_str().unwrap(), "test999", &[]).unwrap();

        let got: serde_json::Value = serde_json::from_str(&std::fs::read_to_string(decrypted.path()).unwrap()).unwrap();
        let expected: serde_json::Value = serde_json::from_str(&std::fs::read_to_string(&input).unwrap()).unwrap();
        assert_eq!(got, expected, "JSON level-2 roundtrip failed");
    }

    #[test]
    fn test_encrypt_level2_yaml_roundtrip() {
        let input = resources("example.yaml");
        let encrypted = tempfile::NamedTempFile::new().unwrap();
        let decrypted = tempfile::NamedTempFile::new().unwrap();

        encrypt(&input, encrypted.path().to_str().unwrap(), "test999", &[], Some(2)).unwrap();
        decrypt(encrypted.path().to_str().unwrap(), decrypted.path().to_str().unwrap(), "test999", &[]).unwrap();

        let got: serde_yaml::Value = serde_yaml::from_str(&std::fs::read_to_string(decrypted.path()).unwrap()).unwrap();
        let expected: serde_yaml::Value = serde_yaml::from_str(&std::fs::read_to_string(&input).unwrap()).unwrap();
        assert_eq!(got, expected, "YAML level-2 roundtrip failed");
    }

    #[test]
    fn test_encrypt_level2_json_structure() {
        // level 2: second.b and second.a are objects → sub-node blobs; first.z etc. are scalars → leaf-encrypted
        let input = resources("example.json");
        let encrypted = tempfile::NamedTempFile::new().unwrap();

        encrypt(&input, encrypted.path().to_str().unwrap(), "test999", &[], Some(2)).unwrap();

        let enc: serde_json::Value = serde_json::from_str(&std::fs::read_to_string(encrypted.path()).unwrap()).unwrap();

        // first-level objects must still be objects (not encrypted)
        assert!(enc["first"].is_object(), "first should still be an object at level 2");
        assert!(enc["second"].is_object(), "second should still be an object at level 2");

        // second-level scalars must be vault-encrypted strings
        assert!(enc["first"]["z"].as_str().unwrap_or("").contains("$ANSIBLE_VAULT"), "first.z should be encrypted");
        assert!(enc["first"]["a"].as_str().unwrap_or("").contains("$ANSIBLE_VAULT"), "first.a should be encrypted");

        // second-level objects must also be vault-encrypted strings (sub-node blobs)
        assert!(enc["second"]["b"].as_str().unwrap_or("").contains("$ANSIBLE_VAULT"), "second.b should be encrypted as sub-node");
        assert!(enc["second"]["a"].as_str().unwrap_or("").contains("$ANSIBLE_VAULT"), "second.a should be encrypted as sub-node");

        // fourth.list (array at level 2) must be a vault-encrypted string
        assert!(enc["fourth"]["list"].as_str().unwrap_or("").contains("$ANSIBLE_VAULT"), "fourth.list should be encrypted as sub-node");
    }

    #[test]
    fn test_encrypt_level3_shallow_nodes_skipped_json() {
        // level 3: first.z is a scalar at depth 2 (doesn't reach depth 3) → stays unencrypted
        let input = resources("example.json");
        let encrypted = tempfile::NamedTempFile::new().unwrap();

        encrypt(&input, encrypted.path().to_str().unwrap(), "test999", &[], Some(3)).unwrap();

        let enc: serde_json::Value = serde_json::from_str(&std::fs::read_to_string(encrypted.path()).unwrap()).unwrap();

        // first.z is a scalar at depth 2, doesn't reach level 3 → stays as plain string
        assert_eq!(enc["first"]["z"], serde_json::Value::String("last".to_string()), "first.z should be untouched at level 3");

        // but second.b.3 is a scalar at depth 3 → should be encrypted
        assert!(enc["second"]["b"]["3"].as_str().unwrap_or("").contains("$ANSIBLE_VAULT"), "second.b.3 should be encrypted at level 3");
    }

    #[test]
    fn test_encrypt_level3_shallow_nodes_skipped_yaml() {
        let input = resources("example.yaml");
        let encrypted = tempfile::NamedTempFile::new().unwrap();

        encrypt(&input, encrypted.path().to_str().unwrap(), "test999", &[], Some(3)).unwrap();

        let enc: serde_yaml::Value = serde_yaml::from_str(&std::fs::read_to_string(encrypted.path()).unwrap()).unwrap();

        // first.z is a scalar at depth 2, doesn't reach level 3 → stays as plain string
        assert_eq!(enc["first"]["z"], serde_yaml::Value::String("last".to_string()), "first.z should be untouched at level 3");

        // second.b.3 (numeric key) is a scalar at depth 3 → should be encrypted
        let key_3 = serde_yaml::Value::Number(serde_yaml::Number::from(3u64));
        let val = enc["second"]["b"].get(&key_3);
        assert!(val.and_then(|v| v.as_str()).unwrap_or("").contains("$ANSIBLE_VAULT"), "second.b.3 should be encrypted at level 3");
    }

    // --- existing roundtrip tests ---

    #[test]
    fn test_encrypt_decrypt_roundtrip_yaml() {
        let input = resources("example.yaml");
        let encrypted = tempfile::NamedTempFile::new().unwrap();
        let decrypted = tempfile::NamedTempFile::new().unwrap();

        encrypt(&input, encrypted.path().to_str().unwrap(), "test999", &[], None).unwrap();
        decrypt(
            encrypted.path().to_str().unwrap(),
            decrypted.path().to_str().unwrap(),
            "test999",
            &[],
        )
        .unwrap();

        let got: serde_yaml::Value = serde_yaml::from_str(
            &std::fs::read_to_string(decrypted.path()).unwrap(),
        )
        .unwrap();
        let expected: serde_yaml::Value =
            serde_yaml::from_str(&std::fs::read_to_string(&input).unwrap()).unwrap();
        assert_eq!(got, expected, "YAML roundtrip: encrypt+decrypt doesn't match input");
    }

    #[test]
    fn test_encrypt_decrypt_roundtrip_yaml_filtered() {
        let input = resources("example.yaml");
        let encrypted = tempfile::NamedTempFile::new().unwrap();
        let decrypted = tempfile::NamedTempFile::new().unwrap();
        let filter = vec!["third.carrot".to_string()];

        encrypt(
            &input,
            encrypted.path().to_str().unwrap(),
            "test999",
            &filter,
            None,
        )
        .unwrap();
        decrypt(
            encrypted.path().to_str().unwrap(),
            decrypted.path().to_str().unwrap(),
            "test999",
            &filter,
        )
        .unwrap();

        let got: serde_yaml::Value = serde_yaml::from_str(
            &std::fs::read_to_string(decrypted.path()).unwrap(),
        )
        .unwrap();
        let expected: serde_yaml::Value =
            serde_yaml::from_str(&std::fs::read_to_string(&input).unwrap()).unwrap();
        assert_eq!(
            got, expected,
            "filtered YAML roundtrip: encrypt+decrypt doesn't match input"
        );
    }

    #[test]
    fn test_encrypt_decrypt_roundtrip_json() {
        let input = resources("example.json");
        let encrypted = tempfile::NamedTempFile::new().unwrap();
        let decrypted = tempfile::NamedTempFile::new().unwrap();

        encrypt(&input, encrypted.path().to_str().unwrap(), "test999", &[], None).unwrap();
        decrypt(
            encrypted.path().to_str().unwrap(),
            decrypted.path().to_str().unwrap(),
            "test999",
            &[],
        )
        .unwrap();

        let got: serde_json::Value = serde_json::from_str(
            &std::fs::read_to_string(decrypted.path()).unwrap(),
        )
        .unwrap();
        let expected: serde_json::Value =
            serde_json::from_str(&std::fs::read_to_string(&input).unwrap()).unwrap();
        assert_eq!(got, expected, "JSON roundtrip: encrypt+decrypt doesn't match input");
    }
}