BlueArchiveToolkit/infrastructure/src/patch_ops.rs

//! File-level Patch and UnityFS write operations.

use std::path::{Path, PathBuf};

use bat_assetbundle::{
    patch_unityfs_field, patch_unityfs_string_field, patch_unityfs_text_asset, FieldPatch,
    StringFieldPatch, TextAssetPatch, UnitySerializedReplacementValue,
};
use serde::{Deserialize, Serialize};

use crate::path_security::{
    ensure_safe_file_target, lexical_absolute, read_file_no_symlink, write_file_atomic,
    STATE_FILE_MODE,
};

/// File patch algorithm selected by `patch.apply`.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum PatchApplyKind {
    /// Deterministic binary hunk patch.
    Binary,
    /// RFC 6902 JSON Patch.
    Json,
    /// UTF-8 Text Patch.
    Text,
}

impl PatchApplyKind {
    /// Stable RPC/CLI label.
    pub fn as_str(self) -> &'static str {
        match self {
            Self::Binary => "binary",
            Self::Json => "json",
            Self::Text => "text",
        }
    }
}

/// Parameters for applying one file-level patch.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct PatchApplyParams {
    /// Patch algorithm.
    pub kind: PatchApplyKind,
    /// Source file path.
    pub source_path: PathBuf,
    /// Patch document path.
    pub patch_path: PathBuf,
    /// Target file path written atomically.
    pub target_path: PathBuf,
}

/// Parameters for patching one UnityFS TextAsset object.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct UnityFsTextAssetPatchParams {
    /// Source UnityFS bundle file.
    pub bundle_path: PathBuf,
    /// Serialized file path inside the UnityFS directory table.
    pub serialized_file_path: String,
    /// Unity object path ID.
    pub path_id: i64,
    /// Replacement raw bytes path.
    pub replacement_path: PathBuf,
    /// Target bundle file path written atomically.
    pub target_path: PathBuf,
    /// Optional expected TextAsset name.
    #[serde(default)]
    pub expected_name: Option<String>,
}

/// Parameters for patching one TypeTree string field inside a UnityFS bundle.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct UnityFsStringFieldPatchParams {
    /// Source UnityFS bundle file.
    pub bundle_path: PathBuf,
    /// Serialized file path inside the UnityFS directory table.
    pub serialized_file_path: String,
    /// Unity object path ID.
    pub path_id: i64,
    /// Stable TypeTree field path.
    pub field_path: String,
    /// Replacement string supplied inline.
    #[serde(default)]
    pub replacement_text: Option<String>,
    /// Replacement UTF-8 file path.
    #[serde(default)]
    pub replacement_path: Option<PathBuf>,
    /// Target bundle file path written atomically.
    pub target_path: PathBuf,
    /// Optional expected source string.
    #[serde(default)]
    pub expected_value: Option<String>,
}

/// Parameters for patching one semantic TypeTree field inside a UnityFS bundle.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct UnityFsFieldPatchParams {
    /// Source UnityFS bundle file.
    pub bundle_path: PathBuf,
    /// Serialized file path inside the UnityFS directory table.
    pub serialized_file_path: String,
    /// Unity object path ID.
    pub path_id: i64,
    /// Stable TypeTree field path.
    pub field_path: String,
    /// Replacement value encoded according to the current TypeTree field type.
    pub replacement: UnitySerializedReplacementValue,
    /// Target bundle file path written atomically.
    pub target_path: PathBuf,
    /// Optional expected source value.
    #[serde(default)]
    pub expected_value: Option<UnitySerializedReplacementValue>,
}

/// Result of a file-level patch operation.
#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
pub struct PatchApplyReport {
    /// RPC/CLI command name.
    pub command: &'static str,
    /// Operation status.
    pub status: &'static str,
    /// Human-readable summary.
    pub message: &'static str,
    /// Patch algorithm.
    pub kind: PatchApplyKind,
    /// Absolute source path.
    pub source_path: PathBuf,
    /// Absolute patch path.
    pub patch_path: PathBuf,
    /// Absolute target path.
    pub target_path: PathBuf,
    /// Source byte length.
    pub source_size: u64,
    /// Patch document byte length.
    pub patch_size: u64,
    /// Target byte length.
    pub target_size: u64,
    /// Source BLAKE3 hash.
    pub source_blake3: String,
    /// Patch document BLAKE3 hash.
    pub patch_blake3: String,
    /// Target BLAKE3 hash.
    pub target_blake3: String,
}

/// Result of a UnityFS write operation.
#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
pub struct UnityFsPatchReport {
    /// RPC/CLI command name.
    pub command: &'static str,
    /// Operation status.
    pub status: &'static str,
    /// Human-readable summary.
    pub message: &'static str,
    /// Absolute source UnityFS bundle path.
    pub bundle_path: PathBuf,
    /// Serialized file path inside the UnityFS directory table.
    pub serialized_file_path: String,
    /// Unity object path ID.
    pub path_id: i64,
    /// Optional TypeTree field path for string-field patches.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub field_path: Option<String>,
    /// Absolute target bundle path.
    pub target_path: PathBuf,
    /// Source bundle byte length.
    pub source_size: u64,
    /// Replacement byte length.
    pub replacement_size: u64,
    /// Target bundle byte length.
    pub target_size: u64,
    /// Source bundle BLAKE3 hash.
    pub source_blake3: String,
    /// Replacement BLAKE3 hash.
    pub replacement_blake3: String,
    /// Target bundle BLAKE3 hash.
    pub target_blake3: String,
}

/// Applies a Binary/JSON/Text patch document to one source file.
pub fn apply_patch_file(params: &PatchApplyParams) -> anyhow::Result<PatchApplyReport> {
    let source_path = lexical_absolute(&params.source_path).map_err(anyhow::Error::msg)?;
    let patch_path = lexical_absolute(&params.patch_path).map_err(anyhow::Error::msg)?;
    let target_path = lexical_absolute(&params.target_path).map_err(anyhow::Error::msg)?;
    ensure_target_is_not_input(&target_path, &[&source_path, &patch_path])?;

    let source = read_required_file(&source_path, "patch source")?;
    let patch = read_required_file(&patch_path, "patch document")?;
    let target = match params.kind {
        PatchApplyKind::Binary => bat_patch::binary::apply_patch(&source, &patch)?,
        PatchApplyKind::Json => {
            let source = std::str::from_utf8(&source)
                .map_err(|error| anyhow::anyhow!("JSON patch source 不是 UTF-8：{error}"))?;
            let patch = std::str::from_utf8(&patch)
                .map_err(|error| anyhow::anyhow!("JSON patch document 不是 UTF-8：{error}"))?;
            bat_patch::json::apply_json_patch(source, patch)?.into_bytes()
        }
        PatchApplyKind::Text => bat_patch::text::apply_patch_bytes(&source, &patch)?,
    };
    write_output_file(&target_path, &target, "patch target")?;

    Ok(PatchApplyReport {
        command: "patch.apply",
        status: "patched",
        message: "patch 已应用并原子写入目标文件",
        kind: params.kind,
        source_path,
        patch_path,
        target_path,
        source_size: source.len() as u64,
        patch_size: patch.len() as u64,
        target_size: target.len() as u64,
        source_blake3: blake3_hex(&source),
        patch_blake3: blake3_hex(&patch),
        target_blake3: blake3_hex(&target),
    })
}

/// Patches one UnityFS TextAsset and writes the rebuilt bundle atomically.
pub fn apply_unityfs_text_asset_patch_file(
    params: &UnityFsTextAssetPatchParams,
) -> anyhow::Result<UnityFsPatchReport> {
    let bundle_path = lexical_absolute(&params.bundle_path).map_err(anyhow::Error::msg)?;
    let replacement_path =
        lexical_absolute(&params.replacement_path).map_err(anyhow::Error::msg)?;
    let target_path = lexical_absolute(&params.target_path).map_err(anyhow::Error::msg)?;
    ensure_target_is_not_input(&target_path, &[&bundle_path, &replacement_path])?;

    let bundle = read_required_file(&bundle_path, "UnityFS bundle")?;
    let replacement = read_required_file(&replacement_path, "TextAsset replacement")?;
    let mut patch = TextAssetPatch::new(
        params.serialized_file_path.clone(),
        params.path_id,
        replacement.clone(),
    );
    patch.expected_name = params.expected_name.clone();
    let target = patch_unityfs_text_asset(&bundle, &patch)?;
    write_output_file(&target_path, &target, "UnityFS target")?;

    Ok(UnityFsPatchReport {
        command: "unityfs.patch_text_asset",
        status: "patched",
        message: "UnityFS TextAsset patch 已应用并原子写入目标 bundle",
        bundle_path,
        serialized_file_path: params.serialized_file_path.clone(),
        path_id: params.path_id,
        field_path: None,
        target_path,
        source_size: bundle.len() as u64,
        replacement_size: replacement.len() as u64,
        target_size: target.len() as u64,
        source_blake3: blake3_hex(&bundle),
        replacement_blake3: blake3_hex(&replacement),
        target_blake3: blake3_hex(&target),
    })
}

/// Patches one TypeTree string field and writes the rebuilt UnityFS bundle atomically.
pub fn apply_unityfs_string_field_patch_file(
    params: &UnityFsStringFieldPatchParams,
) -> anyhow::Result<UnityFsPatchReport> {
    let bundle_path = lexical_absolute(&params.bundle_path).map_err(anyhow::Error::msg)?;
    let target_path = lexical_absolute(&params.target_path).map_err(anyhow::Error::msg)?;
    let replacement = replacement_text(params)?;
    let extra_inputs = params
        .replacement_path
        .as_ref()
        .map(|path| lexical_absolute(path).map_err(anyhow::Error::msg))
        .transpose()?;
    let mut inputs = vec![bundle_path.as_path()];
    if let Some(path) = extra_inputs.as_ref() {
        inputs.push(path.as_path());
    }
    ensure_target_is_not_input(&target_path, &inputs)?;

    let bundle = read_required_file(&bundle_path, "UnityFS bundle")?;
    let patch = StringFieldPatch {
        serialized_file_path: params.serialized_file_path.clone(),
        path_id: params.path_id,
        field_path: params.field_path.clone(),
        expected_value: params.expected_value.clone(),
        replacement: replacement.clone(),
    };
    let target = patch_unityfs_string_field(&bundle, &patch)?;
    write_output_file(&target_path, &target, "UnityFS target")?;

    Ok(UnityFsPatchReport {
        command: "unityfs.patch_string_field",
        status: "patched",
        message: "UnityFS TypeTree string field patch 已应用并原子写入目标 bundle",
        bundle_path,
        serialized_file_path: params.serialized_file_path.clone(),
        path_id: params.path_id,
        field_path: Some(params.field_path.clone()),
        target_path,
        source_size: bundle.len() as u64,
        replacement_size: replacement.len() as u64,
        target_size: target.len() as u64,
        source_blake3: blake3_hex(&bundle),
        replacement_blake3: blake3_hex(replacement.as_bytes()),
        target_blake3: blake3_hex(&target),
    })
}

/// Patches one semantic TypeTree field and writes the rebuilt UnityFS bundle atomically.
pub fn apply_unityfs_field_patch_file(
    params: &UnityFsFieldPatchParams,
) -> anyhow::Result<UnityFsPatchReport> {
    let bundle_path = lexical_absolute(&params.bundle_path).map_err(anyhow::Error::msg)?;
    let target_path = lexical_absolute(&params.target_path).map_err(anyhow::Error::msg)?;
    ensure_target_is_not_input(&target_path, &[&bundle_path])?;

    let bundle = read_required_file(&bundle_path, "UnityFS bundle")?;
    let replacement_size = serde_json::to_vec(&params.replacement)
        .map_err(anyhow::Error::from)?
        .len() as u64;
    let patch = FieldPatch {
        serialized_file_path: params.serialized_file_path.clone(),
        path_id: params.path_id,
        field_path: params.field_path.clone(),
        expected_value: params.expected_value.clone(),
        replacement: params.replacement.clone(),
    };
    let target = patch_unityfs_field(&bundle, &patch)?;
    write_output_file(&target_path, &target, "UnityFS target")?;

    Ok(UnityFsPatchReport {
        command: "unityfs.patch_field",
        status: "patched",
        message: "UnityFS TypeTree field patch 已应用并原子写入目标 bundle",
        bundle_path,
        serialized_file_path: params.serialized_file_path.clone(),
        path_id: params.path_id,
        field_path: Some(params.field_path.clone()),
        target_path,
        source_size: bundle.len() as u64,
        replacement_size,
        target_size: target.len() as u64,
        source_blake3: blake3_hex(&bundle),
        replacement_blake3: blake3_hex(&serde_json::to_vec(&params.replacement)?),
        target_blake3: blake3_hex(&target),
    })
}

fn replacement_text(params: &UnityFsStringFieldPatchParams) -> anyhow::Result<String> {
    match (&params.replacement_text, &params.replacement_path) {
        (Some(_), Some(_)) => Err(anyhow::anyhow!(
            "replacement_text 和 replacement_path 只能指定一个"
        )),
        (Some(text), None) => Ok(text.clone()),
        (None, Some(path)) => {
            let path = lexical_absolute(path).map_err(anyhow::Error::msg)?;
            let bytes = read_required_file(&path, "string replacement")?;
            String::from_utf8(bytes)
                .map_err(|error| anyhow::anyhow!("string replacement 不是 UTF-8：{error}"))
        }
        (None, None) => Err(anyhow::anyhow!(
            "必须指定 replacement_text 或 replacement_path"
        )),
    }
}

fn read_required_file(path: &Path, label: &str) -> anyhow::Result<Vec<u8>> {
    read_file_no_symlink(path, label)
        .map_err(anyhow::Error::msg)?
        .ok_or_else(|| anyhow::anyhow!("{label} 不存在：{}", path.display()))
}

fn write_output_file(path: &Path, bytes: &[u8], label: &str) -> anyhow::Result<()> {
    let parent = path
        .parent()
        .ok_or_else(|| anyhow::anyhow!("{label} 缺少父目录：{}", path.display()))?;
    ensure_safe_file_target(parent, path, label).map_err(anyhow::Error::msg)?;
    write_file_atomic(path, bytes, STATE_FILE_MODE, label).map_err(anyhow::Error::msg)
}

fn ensure_target_is_not_input(target: &Path, inputs: &[&Path]) -> anyhow::Result<()> {
    for input in inputs {
        if target == *input {
            return Err(anyhow::anyhow!(
                "target_path 不能与输入文件相同：{}",
                target.display()
            ));
        }
    }
    Ok(())
}

fn blake3_hex(bytes: &[u8]) -> String {
    blake3::hash(bytes).to_hex().to_string()
}

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

    #[test]
    fn patch_apply_rejects_in_place_target() {
        let path = PathBuf::from("/tmp/source.bin");
        let params = PatchApplyParams {
            kind: PatchApplyKind::Binary,
            source_path: path.clone(),
            patch_path: PathBuf::from("/tmp/patch.json"),
            target_path: path,
        };
        let error = apply_patch_file(&params).unwrap_err().to_string();
        assert!(error.contains("target_path 不能与输入文件相同"));
    }

    #[test]
    fn string_field_params_require_one_replacement_source() {
        let params = UnityFsStringFieldPatchParams {
            bundle_path: PathBuf::from("/tmp/source.bundle"),
            serialized_file_path: "CAB".to_string(),
            path_id: 1,
            field_path: "message".to_string(),
            replacement_text: Some("a".to_string()),
            replacement_path: Some(PathBuf::from("/tmp/replacement.txt")),
            target_path: PathBuf::from("/tmp/target.bundle"),
            expected_value: None,
        };
        let error = replacement_text(&params).unwrap_err().to_string();
        assert!(error.contains("只能指定一个"));
    }
}