Add willplus support

Cargo.toml

@@ -27,7 +27,7 @@ url = { version = "2.5", optional = true }
 utf16string = "0.2"
 
 [features]
-default = ["bgi", "bgi-arc", "bgi-img", "cat-system", "cat-system-arc", "cat-system-img", "circus", "escude", "escude-arc", "kirikiri", "kirikiri-img", "yaneurao", "yaneurao-itufuru"]
+default = ["bgi", "bgi-arc", "bgi-img", "cat-system", "cat-system-arc", "cat-system-img", "circus", "escude", "escude-arc", "kirikiri", "kirikiri-img", "will-plus", "yaneurao", "yaneurao-itufuru"]
 bgi = []
 bgi-arc = ["bgi", "rand", "utils-bit-stream"]
 bgi-img = ["bgi", "image", "utils-bit-stream"]
@@ -39,6 +39,7 @@ escude = ["int-enum"]
 escude-arc = ["escude", "rand", "utils-bit-stream"]
 kirikiri = ["emote-psb", "fancy-regex", "flate2", "json", "utils-escape"]
 kirikiri-img = ["kirikiri", "emote-psb", "image", "libtlg-rs", "url"]
+will-plus = ["utils-str"]
 yaneurao = []
 yaneurao-itufuru = ["yaneurao"]
 # basic feature
@@ -47,6 +48,7 @@ image = ["png"]
 utils-bit-stream = []
 utils-crc32 = []
 utils-escape = ["fancy-regex"]
+utils-str = []
 
 [target.'cfg(windows)'.dependencies]
 windows-sys = { version = "0", features = ["Win32_Globalization", "Win32_System_Diagnostics_Debug"] }

src/ext/io.rs

@@ -531,6 +531,8 @@ pub trait CPeek {
         Ok(i128::from_be_bytes(buf))
     }
 
+    fn cpeek_cstring(&self) -> Result<CString>;
+
     fn cpeek_cstring_at(&self, offset: usize) -> Result<CString> {
         let mut buf = Vec::new();
         let mut byte = [0u8; 1];
@@ -580,6 +582,13 @@ impl<T: Peek> CPeek for Mutex<T> {
         })?;
         lock.peek_at(offset, buf)
     }
+
+    fn cpeek_cstring(&self) -> Result<CString> {
+        let mut lock = self.lock().map_err(|_| {
+            std::io::Error::new(std::io::ErrorKind::Other, "Failed to lock the mutex")
+        })?;
+        lock.peek_cstring()
+    }
 }
 
 pub trait ReadExt {
@@ -957,6 +966,10 @@ impl MemReader {
     pub fn is_eof(&self) -> bool {
         self.pos >= self.data.len()
     }
+
+    pub fn inner(self) -> Vec<u8> {
+        self.data
+    }
 }
 
 impl<'a> MemReaderRef<'a> {
@@ -1036,6 +1049,10 @@ impl CPeek for MemReader {
     fn cpeek_at(&self, offset: usize, buf: &mut [u8]) -> Result<usize> {
         self.to_ref().cpeek_at(offset, buf)
     }
+
+    fn cpeek_cstring(&self) -> Result<CString> {
+        self.to_ref().cpeek_cstring()
+    }
 }
 
 impl<'a> Read for MemReaderRef<'a> {
@@ -1114,6 +1131,17 @@ impl<'a> CPeek for MemReaderRef<'a> {
         buf[..bytes_to_read].copy_from_slice(&self.data[offset..offset + bytes_to_read]);
         Ok(bytes_to_read)
     }
+
+    fn cpeek_cstring(&self) -> Result<CString> {
+        let mut buf = Vec::new();
+        for &byte in &self.data[self.pos..] {
+            if byte == 0 {
+                break;
+            }
+            buf.push(byte);
+        }
+        CString::new(buf).map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e))
+    }
 }
 
 pub struct MemWriter {
@@ -1129,6 +1157,10 @@ impl MemWriter {
         }
     }
 
+    pub fn from_vec(data: Vec<u8>) -> Self {
+        MemWriter { data, pos: 0 }
+    }
+
     pub fn into_inner(self) -> Vec<u8> {
         self.data
     }
@@ -1218,4 +1250,8 @@ impl CPeek for MemWriter {
     fn cpeek_at(&self, offset: usize, buf: &mut [u8]) -> Result<usize> {
         self.to_ref().cpeek_at(offset, buf)
     }
+
+    fn cpeek_cstring(&self) -> Result<CString> {
+        self.to_ref().cpeek_cstring()
+    }
 }

src/scripts/mod.rs

@@ -9,6 +9,8 @@ pub mod circus;
 pub mod escude;
 #[cfg(feature = "kirikiri")]
 pub mod kirikiri;
+#[cfg(feature = "will-plus")]
+pub mod will_plus;
 #[cfg(feature = "yaneurao")]
 pub mod yaneurao;
 
@@ -62,6 +64,8 @@ lazy_static::lazy_static! {
         Box::new(kirikiri::image::dref::DrefBuilder::new()),
         #[cfg(feature = "kirikiri")]
         Box::new(kirikiri::mdf::MdfBuilder::new()),
+        #[cfg(feature = "will-plus")]
+        Box::new(will_plus::ws2::Ws2ScriptBuilder::new()),
     ];
     pub static ref ALL_EXTS: Vec<String> =
         BUILDER.iter().flat_map(|b| b.extensions()).map(|s| s.to_string()).collect();

src/scripts/will_plus/mod.rs

@@ -0,0 +1 @@
+pub mod ws2;

src/scripts/will_plus/ws2.rs

@@ -0,0 +1,305 @@
+use crate::ext::io::*;
+use crate::scripts::base::*;
+use crate::types::*;
+use crate::utils::encoding::*;
+use crate::utils::str::*;
+use anyhow::Result;
+use std::io::{Seek, SeekFrom, Write};
+
+#[derive(Debug)]
+pub struct Ws2ScriptBuilder {}
+
+impl Ws2ScriptBuilder {
+    pub fn new() -> Self {
+        Ws2ScriptBuilder {}
+    }
+}
+
+impl ScriptBuilder for Ws2ScriptBuilder {
+    fn default_encoding(&self) -> Encoding {
+        Encoding::Cp932
+    }
+
+    fn build_script(
+        &self,
+        buf: Vec<u8>,
+        _filename: &str,
+        encoding: Encoding,
+        _archive_encoding: Encoding,
+        config: &ExtraConfig,
+    ) -> Result<Box<dyn Script>> {
+        Ok(Box::new(Ws2Script::new(buf, encoding, config, false)?))
+    }
+
+    fn extensions(&self) -> &'static [&'static str] {
+        &["ws2"]
+    }
+
+    fn script_type(&self) -> &'static ScriptType {
+        &ScriptType::WillPlusWs2
+    }
+}
+
+trait CustomFn {
+    /// check if the current reader's position matches the given byte slice
+    /// 0xFF in the slice is treated as a wildcard that matches any byte
+    fn equal(&self, other: &[u8]) -> bool;
+    /// Reads a string from the current position, decodes it using the specified encoding,
+    fn get_ws2_string(&self, encoding: Encoding) -> Result<Ws2String>;
+}
+
+impl CustomFn for MemReader {
+    fn equal(&self, other: &[u8]) -> bool {
+        self.to_ref().equal(other)
+    }
+
+    fn get_ws2_string(&self, encoding: Encoding) -> Result<Ws2String> {
+        self.to_ref().get_ws2_string(encoding)
+    }
+}
+
+impl<'a> CustomFn for MemReaderRef<'a> {
+    fn equal(&self, other: &[u8]) -> bool {
+        if self.pos + other.len() > self.data.len() {
+            return false;
+        }
+        for (i, &byte) in other.iter().enumerate() {
+            if self.data[self.pos + i] != byte && byte != 0xFF {
+                return false;
+            }
+        }
+        true
+    }
+
+    fn get_ws2_string(&self, encoding: Encoding) -> Result<Ws2String> {
+        let pos = self.pos;
+        let s = self.cpeek_cstring()?;
+        let decoded = decode_to_string(encoding, s.as_bytes(), true)?;
+        Ok(Ws2String {
+            pos,
+            str: decoded,
+            len: s.as_bytes_with_nul().len(),
+            actor: None,
+        })
+    }
+}
+
+struct EncryptWriter<T: Write + Seek> {
+    writer: T,
+}
+
+impl<T: Write + Seek> EncryptWriter<T> {
+    pub fn new(writer: T) -> Self {
+        EncryptWriter { writer }
+    }
+}
+
+impl<T: Write + Seek> Write for EncryptWriter<T> {
+    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
+        let encrypted: Vec<u8> = buf.iter().map(|&b| b.rotate_left(2)).collect();
+        self.writer.write(&encrypted)
+    }
+
+    fn flush(&mut self) -> std::io::Result<()> {
+        self.writer.flush()
+    }
+}
+
+impl<T: Write + Seek> Seek for EncryptWriter<T> {
+    fn seek(&mut self, pos: SeekFrom) -> std::io::Result<u64> {
+        self.writer.seek(pos)
+    }
+    fn stream_position(&mut self) -> std::io::Result<u64> {
+        self.writer.stream_position()
+    }
+    fn rewind(&mut self) -> std::io::Result<()> {
+        self.writer.rewind()
+    }
+    fn seek_relative(&mut self, offset: i64) -> std::io::Result<()> {
+        self.writer.seek_relative(offset)
+    }
+}
+
+#[derive(Debug)]
+struct Ws2String {
+    pos: usize,
+    str: String,
+    /// Length of the string in bytes, including the null terminator
+    len: usize,
+    actor: Option<Box<Ws2String>>,
+}
+
+#[derive(Debug)]
+pub struct Ws2Script {
+    data: MemReader,
+    strs: Vec<Ws2String>,
+    /// Need encrypt when outputting
+    encrypted: bool,
+}
+
+impl Ws2Script {
+    pub fn new(
+        buf: Vec<u8>,
+        encoding: Encoding,
+        config: &ExtraConfig,
+        decrypted: bool,
+    ) -> Result<Self> {
+        let mut reader = MemReader::new(buf);
+        let mut strs = Vec::new();
+        let mut actor = None;
+        while !reader.is_eof() {
+            if reader.equal(b"\x00\xFF\x0F\x02") {
+                reader.pos += 4;
+                if reader.cpeek_u8()? == 0 {
+                    reader.pos += 1;
+                    continue;
+                }
+                let mut continu = true;
+                while !reader.is_eof() && continu {
+                    reader.pos += 2;
+                    let str = reader.get_ws2_string(encoding)?;
+                    reader.pos += str.len + 4;
+                    while reader.cpeek_u8()? != 0 {
+                        reader.pos += 1;
+                    }
+                    reader.pos += 1;
+                    if reader.cpeek_u8()? == 0xFF {
+                        continu = false;
+                    }
+                    strs.push(str);
+                }
+            }
+            if reader.equal(b"%LC") || reader.equal(b"%LF") {
+                reader.pos += 3;
+                let str = Box::new(reader.get_ws2_string(encoding)?);
+                reader.pos += str.len + 4;
+                actor = Some(str);
+            }
+            if reader.equal(b"char\0") {
+                reader.pos += 5;
+                let mut str = reader.get_ws2_string(encoding)?;
+                reader.pos += str.len + 4;
+                str.actor = actor.take();
+                strs.push(str);
+            }
+            reader.pos += 1;
+        }
+        if !decrypted && strs.is_empty() {
+            let mut data = reader.inner();
+            Self::decrypt(&mut data);
+            return Self::new(data, encoding, config, true);
+        }
+        Ok(Self {
+            data: reader,
+            strs,
+            encrypted: decrypted,
+        })
+    }
+
+    fn decrypt(data: &mut [u8]) {
+        for byte in data.iter_mut() {
+            *byte = (*byte).rotate_right(2);
+        }
+    }
+}
+
+impl Script for Ws2Script {
+    fn default_output_script_type(&self) -> OutputScriptType {
+        OutputScriptType::Json
+    }
+
+    fn default_format_type(&self) -> FormatOptions {
+        FormatOptions::None
+    }
+
+    fn extract_messages(&self) -> Result<Vec<Message>> {
+        let mut messages = Vec::new();
+        for str in &self.strs {
+            let message = Message {
+                message: str.str.trim_end_matches("%K%P").to_string(),
+                name: str.actor.as_ref().map(|a| {
+                    a.str
+                        .trim_start_matches("%LC")
+                        .trim_start_matches("%LF")
+                        .to_string()
+                }),
+            };
+            messages.push(message);
+        }
+        Ok(messages)
+    }
+
+    fn import_messages<'a>(
+        &'a self,
+        messages: Vec<Message>,
+        file: Box<dyn WriteSeek + 'a>,
+        encoding: Encoding,
+        replacement: Option<&'a ReplacementTable>,
+    ) -> Result<()> {
+        let mut mes = messages.iter();
+        let mut m = mes.next();
+        let mut file = if self.encrypted {
+            Box::new(EncryptWriter::new(file))
+        } else {
+            file
+        };
+        file.write_all(&self.data.data)?;
+        for str in &self.strs {
+            let me = match m {
+                Some(m) => m,
+                None => {
+                    return Err(anyhow::anyhow!("No enough messages."));
+                }
+            };
+            if let Some(actor) = &str.actor {
+                let prefix = if actor.str.starts_with("%LC") {
+                    "%LC"
+                } else if actor.str.starts_with("%LF") {
+                    "%LF"
+                } else {
+                    ""
+                };
+                let target_len = actor.len - prefix.len() - 1; // -1 for null terminator
+                let mut name = match me.name.as_ref() {
+                    Some(name) => name.to_owned(),
+                    None => return Err(anyhow::anyhow!("Message without name.")),
+                };
+                if let Some(replacement) = replacement {
+                    for (k, v) in &replacement.map {
+                        name = name.replace(k, v);
+                    }
+                }
+                let mut encoded = encode_string(encoding, &name, true)?;
+                if encoded.len() > target_len {
+                    eprintln!("Warning: Name '{}' is too long, truncating.", name);
+                    crate::COUNTER.inc_warning();
+                    encoded = truncate_string(&name, target_len, encoding, true)?;
+                }
+                encoded.resize(target_len, 0x20); // Fill with spaces
+                file.write_all_at(actor.pos + prefix.len(), &encoded)?;
+            }
+            let suffix = if str.str.ends_with("%K%P") {
+                "%K%P"
+            } else {
+                ""
+            };
+            let target_len = str.len - suffix.len() - 1; // -1 for null terminator
+            let mut message = me.message.clone();
+            if let Some(replacement) = replacement {
+                for (k, v) in &replacement.map {
+                    message = message.replace(k, v);
+                }
+            }
+            let mut encoded = encode_string(encoding, &message, true)?;
+            if encoded.len() > target_len {
+                eprintln!("Warning: Message '{}' is too long, truncating.", message);
+                crate::COUNTER.inc_warning();
+                encoded = truncate_string(&message, target_len, encoding, true)?;
+            }
+            encoded.resize(target_len, 0x20); // Fill with spaces
+            file.write_all_at(str.pos, &encoded)?;
+            m = mes.next();
+        }
+        Ok(())
+    }
+}

src/types.rs

@@ -314,6 +314,9 @@ pub enum ScriptType {
     #[value(alias("itufuru-arc"))]
     /// Yaneurao Itufuru script archive
     YaneuraoItufuruArc,
+    #[cfg(feature = "will-plus")]
+    /// WillPlus ws2 script
+    WillPlusWs2,
 }
 
 #[derive(Debug, Serialize, Deserialize)]

src/utils/mod.rs

@@ -13,4 +13,6 @@ pub mod files;
 pub mod img;
 pub mod macros;
 pub mod name_replacement;
+#[cfg(feature = "utils-str")]
+pub mod str;
 pub mod struct_pack;

src/utils/str.rs

@@ -0,0 +1,19 @@
+use crate::types::*;
+use crate::utils::encoding::*;
+use anyhow::Result;
+use unicode_segmentation::UnicodeSegmentation;
+
+/// Truncate a string to a specified length, encoding it with the given encoding.
+/// Output size may less than or equal to the specified length.
+pub fn truncate_string(s: &str, length: usize, encoding: Encoding, check: bool) -> Result<Vec<u8>> {
+    let vec: Vec<_> = UnicodeSegmentation::graphemes(s, true).collect();
+    let mut result = Vec::new();
+    for graphemes in vec {
+        let data = encode_string(encoding, graphemes, check)?;
+        if result.len() + data.len() > length {
+            break;
+        }
+        result.extend(data);
+    }
+    return Ok(result);
+}