msg_tool\scripts\will_plus/
ws2.rs

1//! WillPlus Script File (.ws2)
2use super::ws2_disasm::*;
3use crate::ext::io::*;
4use crate::scripts::base::*;
5use crate::types::*;
6use crate::utils::encoding::*;
7use crate::utils::str::*;
8use anyhow::Result;
9use std::io::{Seek, SeekFrom, Write};
10
11#[derive(Debug)]
12/// WillPlus Script Builder
13pub struct Ws2ScriptBuilder {}
14
15impl Ws2ScriptBuilder {
16    /// Creates a new instance of `Ws2ScriptBuilder`
17    pub fn new() -> Self {
18        Ws2ScriptBuilder {}
19    }
20}
21
22impl ScriptBuilder for Ws2ScriptBuilder {
23    fn default_encoding(&self) -> Encoding {
24        Encoding::Cp932
25    }
26
27    fn build_script(
28        &self,
29        buf: Vec<u8>,
30        _filename: &str,
31        encoding: Encoding,
32        _archive_encoding: Encoding,
33        config: &ExtraConfig,
34        _archive: Option<&Box<dyn Script>>,
35    ) -> Result<Box<dyn Script>> {
36        if !config.will_plus_ws2_no_disasm {
37            match Ws2DisasmScript::new(&buf, encoding, config, false) {
38                Ok(script) => return Ok(Box::new(script)),
39                Err(e) => {
40                    eprintln!(
41                        "WARNING: Failed to disassemble WS2 script: {}. An another parser is used.",
42                        e
43                    );
44                    crate::COUNTER.inc_warning();
45                }
46            }
47        }
48        Ok(Box::new(Ws2Script::new(buf, encoding, config, false)?))
49    }
50
51    fn extensions(&self) -> &'static [&'static str] {
52        &["ws2"]
53    }
54
55    fn script_type(&self) -> &'static ScriptType {
56        &ScriptType::WillPlusWs2
57    }
58}
59
60trait CustomFn {
61    /// check if the current reader's position matches the given byte slice
62    /// 0xFF in the slice is treated as a wildcard that matches any byte
63    fn equal(&self, other: &[u8]) -> bool;
64    /// Reads a string from the current position, decodes it using the specified encoding,
65    fn get_ws2_string(&self, encoding: Encoding) -> Result<Ws2String>;
66}
67
68impl CustomFn for MemReader {
69    fn equal(&self, other: &[u8]) -> bool {
70        self.to_ref().equal(other)
71    }
72
73    fn get_ws2_string(&self, encoding: Encoding) -> Result<Ws2String> {
74        self.to_ref().get_ws2_string(encoding)
75    }
76}
77
78impl<'a> CustomFn for MemReaderRef<'a> {
79    fn equal(&self, other: &[u8]) -> bool {
80        if self.pos + other.len() > self.data.len() {
81            return false;
82        }
83        for (i, &byte) in other.iter().enumerate() {
84            if self.data[self.pos + i] != byte && byte != 0xFF {
85                return false;
86            }
87        }
88        true
89    }
90
91    fn get_ws2_string(&self, encoding: Encoding) -> Result<Ws2String> {
92        let pos = self.pos;
93        let s = self.cpeek_cstring()?;
94        let decoded = decode_to_string(encoding, s.as_bytes(), true)?;
95        Ok(Ws2String {
96            pos,
97            str: decoded,
98            len: s.as_bytes_with_nul().len(),
99            actor: None,
100        })
101    }
102}
103
104struct EncryptWriter<T: Write + Seek> {
105    writer: T,
106}
107
108impl<T: Write + Seek> EncryptWriter<T> {
109    pub fn new(writer: T) -> Self {
110        EncryptWriter { writer }
111    }
112}
113
114impl<T: Write + Seek> Write for EncryptWriter<T> {
115    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
116        let encrypted: Vec<u8> = buf.iter().map(|&b| b.rotate_left(2)).collect();
117        self.writer.write(&encrypted)
118    }
119
120    fn flush(&mut self) -> std::io::Result<()> {
121        self.writer.flush()
122    }
123}
124
125impl<T: Write + Seek> Seek for EncryptWriter<T> {
126    fn seek(&mut self, pos: SeekFrom) -> std::io::Result<u64> {
127        self.writer.seek(pos)
128    }
129    fn stream_position(&mut self) -> std::io::Result<u64> {
130        self.writer.stream_position()
131    }
132    fn rewind(&mut self) -> std::io::Result<()> {
133        self.writer.rewind()
134    }
135    fn seek_relative(&mut self, offset: i64) -> std::io::Result<()> {
136        self.writer.seek_relative(offset)
137    }
138}
139
140#[derive(Debug)]
141struct Ws2String {
142    pos: usize,
143    str: String,
144    /// Length of the string in bytes, including the null terminator
145    len: usize,
146    actor: Option<Box<Ws2String>>,
147}
148
149#[derive(Debug)]
150/// WillPlus Script (without disassembly)
151pub struct Ws2Script {
152    data: MemReader,
153    strs: Vec<Ws2String>,
154    /// Need encrypt when outputting
155    encrypted: bool,
156}
157
158impl Ws2Script {
159    /// Creates a new `Ws2Script`
160    ///
161    /// * `buf` - The buffer containing the script data
162    /// * `encoding` - The encoding used for the script
163    /// * `config` - Extra configuration options
164    /// * `decrypted` - Whether the script is decrypted or not
165    pub fn new(
166        buf: Vec<u8>,
167        encoding: Encoding,
168        config: &ExtraConfig,
169        decrypted: bool,
170    ) -> Result<Self> {
171        let mut reader = MemReader::new(buf);
172        let mut strs = Vec::new();
173        let mut actor = None;
174        while !reader.is_eof() {
175            if reader.equal(b"\x00\xFF\x0F\x02") {
176                reader.pos += 4;
177                if reader.cpeek_u8()? == 0 {
178                    reader.pos += 1;
179                    continue;
180                }
181                let mut continu = true;
182                while !reader.is_eof() && continu {
183                    reader.pos += 2;
184                    let str = reader.get_ws2_string(encoding)?;
185                    reader.pos += str.len + 4;
186                    while reader.cpeek_u8()? != 0 {
187                        reader.pos += 1;
188                    }
189                    reader.pos += 1;
190                    if reader.cpeek_u8()? == 0xFF {
191                        continu = false;
192                    }
193                    strs.push(str);
194                }
195            }
196            if reader.equal(b"%LC") || reader.equal(b"%LF") {
197                reader.pos += 3;
198                let str = Box::new(reader.get_ws2_string(encoding)?);
199                reader.pos += str.len + 4;
200                actor = Some(str);
201            }
202            if reader.equal(b"char\0") {
203                reader.pos += 5;
204                let mut str = reader.get_ws2_string(encoding)?;
205                reader.pos += str.len + 4;
206                str.actor = actor.take();
207                strs.push(str);
208            }
209            reader.pos += 1;
210        }
211        if !decrypted && strs.is_empty() {
212            let mut data = reader.inner();
213            Self::decrypt(&mut data);
214            return Self::new(data, encoding, config, true);
215        }
216        Ok(Self {
217            data: reader,
218            strs,
219            encrypted: decrypted,
220        })
221    }
222
223    fn decrypt(data: &mut [u8]) {
224        for byte in data.iter_mut() {
225            *byte = (*byte).rotate_right(2);
226        }
227    }
228}
229
230impl Script for Ws2Script {
231    fn default_output_script_type(&self) -> OutputScriptType {
232        OutputScriptType::Json
233    }
234
235    fn default_format_type(&self) -> FormatOptions {
236        FormatOptions::None
237    }
238
239    fn extract_messages(&self) -> Result<Vec<Message>> {
240        let mut messages = Vec::new();
241        for str in &self.strs {
242            let message = Message {
243                message: str.str.trim_end_matches("%K%P").to_string(),
244                name: str.actor.as_ref().map(|a| {
245                    a.str
246                        .trim_start_matches("%LC")
247                        .trim_start_matches("%LF")
248                        .to_string()
249                }),
250            };
251            messages.push(message);
252        }
253        Ok(messages)
254    }
255
256    fn import_messages<'a>(
257        &'a self,
258        messages: Vec<Message>,
259        file: Box<dyn WriteSeek + 'a>,
260        _filename: &str,
261        encoding: Encoding,
262        replacement: Option<&'a ReplacementTable>,
263    ) -> Result<()> {
264        let mut mes = messages.iter();
265        let mut m = mes.next();
266        let mut file = if self.encrypted {
267            Box::new(EncryptWriter::new(file))
268        } else {
269            file
270        };
271        file.write_all(&self.data.data)?;
272        for str in &self.strs {
273            let me = match m {
274                Some(m) => m,
275                None => {
276                    return Err(anyhow::anyhow!("No enough messages."));
277                }
278            };
279            if let Some(actor) = &str.actor {
280                let prefix = if actor.str.starts_with("%LC") {
281                    "%LC"
282                } else if actor.str.starts_with("%LF") {
283                    "%LF"
284                } else {
285                    ""
286                };
287                let target_len = actor.len - prefix.len() - 1; // -1 for null terminator
288                let mut name = match me.name.as_ref() {
289                    Some(name) => name.to_owned(),
290                    None => return Err(anyhow::anyhow!("Message without name.")),
291                };
292                if let Some(replacement) = replacement {
293                    for (k, v) in &replacement.map {
294                        name = name.replace(k, v);
295                    }
296                }
297                let mut encoded = encode_string(encoding, &name, true)?;
298                if encoded.len() > target_len {
299                    eprintln!("Warning: Name '{}' is too long, truncating.", name);
300                    crate::COUNTER.inc_warning();
301                    encoded = truncate_string(&name, target_len, encoding, true)?;
302                }
303                encoded.resize(target_len, 0x20); // Fill with spaces
304                file.write_all_at(actor.pos as u64 + prefix.len() as u64, &encoded)?;
305            }
306            let suffix = if str.str.ends_with("%K%P") {
307                "%K%P"
308            } else {
309                ""
310            };
311            let target_len = str.len - suffix.len() - 1; // -1 for null terminator
312            let mut message = me.message.clone();
313            if let Some(replacement) = replacement {
314                for (k, v) in &replacement.map {
315                    message = message.replace(k, v);
316                }
317            }
318            let mut encoded = encode_string(encoding, &message, true)?;
319            if encoded.len() > target_len {
320                eprintln!("Warning: Message '{}' is too long, truncating.", message);
321                crate::COUNTER.inc_warning();
322                encoded = truncate_string(&message, target_len, encoding, true)?;
323            }
324            encoded.resize(target_len, 0x20); // Fill with spaces
325            file.write_all_at(str.pos as u64, &encoded)?;
326            m = mes.next();
327        }
328        if m.is_some() || mes.next().is_some() {
329            return Err(anyhow::anyhow!("Too many messages provided."));
330        }
331        Ok(())
332    }
333}
334
335#[derive(Debug)]
336/// WillPlus Disassembled Script
337pub struct Ws2DisasmScript {
338    data: MemReader,
339    texts: Vec<Ws2DString>,
340    addresses: Vec<usize>,
341    /// Need encrypt when outputting
342    encrypted: bool,
343    encoding: Encoding,
344}
345
346impl Ws2DisasmScript {
347    /// Creates a new `Ws2DisasmScript`
348    ///
349    /// * `buf` - The buffer containing the script data
350    /// * `encoding` - The encoding used for the script
351    /// * `config` - Extra configuration options
352    /// * `decrypted` - Whether the script is decrypted or not
353    pub fn new(
354        buf: &[u8],
355        encoding: Encoding,
356        config: &ExtraConfig,
357        decrypted: bool,
358    ) -> Result<Self> {
359        match disassmble(&buf) {
360            Ok((addresses, texts)) => {
361                return Ok(Self {
362                    data: MemReader::new(buf.to_vec()),
363                    texts,
364                    addresses,
365                    encrypted: decrypted,
366                    encoding,
367                });
368            }
369            Err(e) => {
370                if decrypted {
371                    return Err(e);
372                } else {
373                    let mut data = buf.to_vec();
374                    Ws2Script::decrypt(&mut data);
375                    return Self::new(&data, encoding, config, true);
376                }
377            }
378        }
379    }
380}
381
382impl Script for Ws2DisasmScript {
383    fn default_output_script_type(&self) -> OutputScriptType {
384        OutputScriptType::Json
385    }
386
387    fn default_format_type(&self) -> FormatOptions {
388        FormatOptions::None
389    }
390
391    fn extract_messages(&self) -> Result<Vec<Message>> {
392        let mut messages = Vec::new();
393        let mut name = None;
394        for text in &self.texts {
395            match text.typ {
396                StringType::Name => {
397                    let text = decode_to_string(self.encoding, text.text.as_bytes(), false)?
398                        .trim_start_matches("%LC")
399                        .trim_start_matches("%LF")
400                        .to_string();
401                    name = Some(text);
402                }
403                StringType::Message => {
404                    let message = decode_to_string(self.encoding, text.text.as_bytes(), false)?
405                        .trim_end_matches("%K%P")
406                        .to_string();
407                    messages.push(Message {
408                        message,
409                        name: name.take(),
410                    });
411                }
412                StringType::Internal => {}
413            }
414        }
415        Ok(messages)
416    }
417
418    fn import_messages<'a>(
419        &'a self,
420        messages: Vec<Message>,
421        file: Box<dyn WriteSeek + 'a>,
422        _filename: &str,
423        encoding: Encoding,
424        replacement: Option<&'a ReplacementTable>,
425    ) -> Result<()> {
426        let mut output = if self.encrypted {
427            Box::new(EncryptWriter::new(file))
428        } else {
429            file
430        };
431        let mut mes = messages.iter();
432        let mut mess = mes.next();
433        {
434            let mut patcher = BinaryPatcher::new(
435                MemReaderRef::new(&self.data.data),
436                &mut output,
437                |s| Ok(s),
438                |s| Ok(s),
439            )?;
440            for s in &self.texts {
441                let mut encoded = match s.typ {
442                    StringType::Name => {
443                        let prefix = if s.text.as_bytes().starts_with(b"%LC") {
444                            "%LC"
445                        } else if s.text.as_bytes().starts_with(b"%LF") {
446                            "%LF"
447                        } else {
448                            ""
449                        };
450                        let m = match mess {
451                            Some(m) => m,
452                            None => {
453                                return Err(anyhow::anyhow!("No enough messages."));
454                            }
455                        };
456                        let mut name = match m.name.as_ref() {
457                            Some(name) => name.to_owned(),
458                            None => return Err(anyhow::anyhow!("Message without name.")),
459                        };
460                        if let Some(replacement) = replacement {
461                            for (k, v) in &replacement.map {
462                                name = name.replace(k, v);
463                            }
464                        }
465                        name = prefix.to_owned() + &name;
466                        encode_string(encoding, &name, false)?
467                    }
468                    StringType::Message => {
469                        let suffix = if s.text.as_bytes().ends_with(b"%K%P") {
470                            "%K%P"
471                        } else {
472                            ""
473                        };
474                        let m = match mess {
475                            Some(m) => m,
476                            None => {
477                                return Err(anyhow::anyhow!("No enough messages."));
478                            }
479                        };
480                        let mut message = m.message.clone();
481                        if let Some(replacement) = replacement {
482                            for (k, v) in &replacement.map {
483                                message = message.replace(k, v);
484                            }
485                        }
486                        mess = mes.next();
487                        message.push_str(suffix);
488                        encode_string(encoding, &message, false)?
489                    }
490                    StringType::Internal => s.text.as_bytes().to_vec(),
491                };
492                encoded.push(0); // Null terminator
493                patcher.copy_up_to(s.offset as u64)?;
494                patcher.replace_bytes(s.len as u64, &encoded)?;
495            }
496            if mess.is_some() || mes.next().is_some() {
497                return Err(anyhow::anyhow!("Too many messages provided."));
498            }
499            patcher.copy_up_to(self.data.data.len() as u64)?;
500            for offset in &self.addresses {
501                patcher.patch_u32_address(*offset as u64)?;
502            }
503        }
504        output.flush()?;
505        Ok(())
506    }
507}