msg_tool\scripts\favorite/
hcb.rs

1//! Favorite HCB script (.hcb)
2use super::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::Write;
10
11#[derive(Debug)]
12/// Favorite HCB script builder
13pub struct HcbScriptBuilder {}
14
15impl HcbScriptBuilder {
16    /// Create a new HcbScriptBuilder
17    pub fn new() -> Self {
18        Self {}
19    }
20}
21
22impl ScriptBuilder for HcbScriptBuilder {
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        Ok(Box::new(HcbScript::new(buf, encoding, config)?))
37    }
38
39    fn extensions(&self) -> &'static [&'static str] {
40        &["hcb"]
41    }
42
43    fn script_type(&self) -> &'static ScriptType {
44        &ScriptType::Favorite
45    }
46}
47
48#[derive(Debug)]
49pub struct HcbScript {
50    data: Data,
51    reader: MemReader,
52    custom_yaml: bool,
53    filter_ascii: bool,
54    encoding: Encoding,
55}
56
57impl HcbScript {
58    pub fn new(buf: Vec<u8>, encoding: Encoding, config: &ExtraConfig) -> Result<Self> {
59        let reader = MemReader::new(buf);
60        let data = Data::disasm(reader.to_ref(), encoding)?;
61        Ok(Self {
62            data,
63            reader,
64            custom_yaml: config.custom_yaml,
65            filter_ascii: config.favorite_hcb_filter_ascii,
66            encoding,
67        })
68    }
69}
70
71impl Script for HcbScript {
72    fn default_output_script_type(&self) -> OutputScriptType {
73        OutputScriptType::Json
74    }
75
76    fn default_format_type(&self) -> FormatOptions {
77        FormatOptions::None
78    }
79
80    fn is_output_supported(&self, _: OutputScriptType) -> bool {
81        true
82    }
83
84    fn custom_output_extension<'a>(&'a self) -> &'a str {
85        if self.custom_yaml { "yaml" } else { "json" }
86    }
87
88    fn extract_messages(&self) -> Result<Vec<Message>> {
89        let messages = self
90            .data
91            .extract_messages(self.filter_ascii)
92            .into_iter()
93            .map(|(speaker, message)| Message::new(message, speaker))
94            .collect();
95        Ok(messages)
96    }
97
98    fn import_messages<'a>(
99        &'a self,
100        messages: Vec<Message>,
101        file: Box<dyn WriteSeek + 'a>,
102        _filename: &str,
103        encoding: Encoding,
104        replacement: Option<&'a ReplacementTable>,
105    ) -> Result<()> {
106        let mut mess = messages.iter();
107        let mut mes = mess.next();
108        let mut patcher =
109            BinaryPatcher::new(self.reader.to_ref(), file, |pos| Ok(pos), |pos| Ok(pos))?;
110        let mut need_pacth_addresses = Vec::new();
111        let mut new_need_patch_addresses = Vec::new();
112        let thread_start_callid = self
113            .data
114            .sys_imports
115            .iter()
116            .position(|s| s == "ThreadStart")
117            .map(|i| i as u16)
118            .unwrap_or(u16::MAX);
119        let mut func_index = 0;
120        let func_len = self.data.functions.len();
121        while func_index < func_len {
122            let func = &self.data.functions[func_index];
123            let mut cur_pos = func.pos + 1;
124            if matches!(func.opcode, 0x02 | 0x06 | 0x07) {
125                need_pacth_addresses.push(cur_pos);
126            }
127            if func.opcode == 0x03 {
128                let syscall_id = if let Some(Operand::W(id)) = func.operands.get(0) {
129                    *id
130                } else {
131                    anyhow::bail!("Invalid syscall operand at function index {}", func_index);
132                };
133                if syscall_id == thread_start_callid {
134                    if func_index == 0 {
135                        anyhow::bail!("ThreadStart syscall cannot be at function index 0");
136                    }
137                    let pre_func = &self.data.functions[func_index - 1];
138                    if pre_func.opcode == 0x0a {
139                        need_pacth_addresses.push(pre_func.pos + 1);
140                    }
141                }
142            }
143            for operand in &func.operands {
144                if let Operand::S(s) = operand {
145                    if self.filter_ascii && s.chars().all(|c| c.is_ascii()) {
146                        continue;
147                    }
148                    let m = match mes {
149                        Some(m) => m,
150                        None => {
151                            return Err(anyhow::anyhow!(
152                                "Not enough messages to import. Missing message: {}",
153                                s
154                            ));
155                        }
156                    };
157                    let mut message = m.message.clone();
158                    if let Some(table) = replacement {
159                        for (k, v) in &table.map {
160                            message = message.replace(k, v);
161                        }
162                    }
163                    patcher.copy_up_to(cur_pos)?;
164                    let ori_len = operand.len(self.encoding)? as u64;
165                    let mut s = encode_string(encoding, &message, true)?;
166                    s.push(0); // null-terminated
167                    let len = s.len();
168                    if len > 255 {
169                        return Err(anyhow::anyhow!(
170                            "Message too long to import in functions section (max 255 bytes): {}",
171                            message
172                        ));
173                    }
174                    patcher.replace_bytes_with_write(ori_len, |writer| {
175                        writer.write_u8(len as u8)?;
176                        writer.write_all(&s)?;
177                        Ok(())
178                    })?;
179                    mes = mess.next();
180                }
181                cur_pos += operand.len(self.encoding)? as u64;
182            }
183            func_index += 1;
184        }
185        func_index = 0;
186        let func_len = self.data.main_script.len();
187        'outer: while func_index < func_len {
188            let func = &self.data.main_script[func_index];
189            let mut cur_pos = func.pos + 1;
190            if matches!(func.opcode, 0x02 | 0x06 | 0x07) {
191                need_pacth_addresses.push(cur_pos);
192            }
193            if func.opcode == 0x03 {
194                let syscall_id = if let Some(Operand::W(id)) = func.operands.get(0) {
195                    *id
196                } else {
197                    anyhow::bail!("Invalid syscall operand at function index {}", func_index);
198                };
199                if syscall_id == thread_start_callid {
200                    if func_index == 0 {
201                        anyhow::bail!("ThreadStart syscall cannot be at function index 0");
202                    }
203                    let pre_func = &self.data.main_script[func_index - 1];
204                    if pre_func.opcode == 0x0a {
205                        need_pacth_addresses.push(pre_func.pos + 1);
206                    }
207                }
208            }
209            for operand in &func.operands {
210                if let Operand::S(s) = operand {
211                    if self.filter_ascii && s.chars().all(|c| c.is_ascii()) {
212                        continue;
213                    }
214                    let m = match mes {
215                        Some(m) => m,
216                        None => {
217                            return Err(anyhow::anyhow!(
218                                "Not enough messages to import. Missing message: {}",
219                                s
220                            ));
221                        }
222                    };
223                    let mut message = m.message.clone();
224                    if let Some(table) = replacement {
225                        for (k, v) in &table.map {
226                            message = message.replace(k, v);
227                        }
228                    }
229                    mes = mess.next();
230                    patcher.copy_up_to(cur_pos)?;
231                    let ori_len = operand.len(self.encoding)? as u64;
232                    let mut s = encode_string(encoding, &message, true)?;
233                    s.push(0); // null-terminated
234                    let len = s.len();
235                    if len > 255 {
236                        if func.opcode != 0x0e {
237                            anyhow::bail!(
238                                "Message too long to import in main script functions section (max 255 bytes): {}",
239                                message
240                            );
241                        }
242                        let cur = message.as_str();
243                        let (mut s, mut remaining) =
244                            truncate_string_with_enter(cur, 254, encoding)?;
245                        s.push(0); // null-terminated
246                        let len = s.len();
247                        patcher.replace_bytes_with_write(ori_len, |writer| {
248                            writer.write_u8(len as u8)?;
249                            writer.write_all(&s)?;
250                            Ok(())
251                        })?;
252                        let mut new_funcs = Vec::new();
253                        func_index += 1;
254                        loop {
255                            let toper = &self.data.main_script[func_index];
256                            new_funcs.push(toper.clone());
257                            if matches!(toper.opcode, 0x02 | 0x06 | 0x07) {
258                                need_pacth_addresses.push(toper.pos + 1);
259                            }
260                            if toper.opcode == 0x03 {
261                                let syscall_id = if let Some(Operand::W(id)) = toper.operands.get(0)
262                                {
263                                    *id
264                                } else {
265                                    anyhow::bail!(
266                                        "Invalid syscall operand at function index {}",
267                                        func_index
268                                    );
269                                };
270                                if syscall_id == thread_start_callid {
271                                    if func_index == 0 {
272                                        anyhow::bail!(
273                                            "ThreadStart syscall cannot be at function index 0"
274                                        );
275                                    }
276                                    let pre_func = &self.data.main_script[func_index - 1];
277                                    if pre_func.opcode == 0x0a {
278                                        need_pacth_addresses.push(pre_func.pos + 1);
279                                    }
280                                }
281                            }
282                            func_index += 1;
283                            // Copy until the next call opcode
284                            if toper.opcode == 0x02 {
285                                break;
286                            }
287                        }
288                        cur_pos = self.data.main_script[func_index].pos;
289                        patcher.copy_up_to(cur_pos)?;
290                        let mut mem = MemWriter::new();
291                        while let Some(remain) = remaining {
292                            let (mut s, rem) = truncate_string_with_enter(remain, 254, encoding)?;
293                            s.push(0); // null-terminated
294                            let len = s.len();
295                            remaining = rem;
296                            mem.write_u8(0x0e)?; // pushstring
297                            mem.write_u8(len as u8)?;
298                            mem.write_all(&s)?;
299                            let mut tindex = 0;
300                            let tlen = new_funcs.len();
301                            while tindex < tlen {
302                                let toper = &new_funcs[tindex];
303                                mem.write_u8(toper.opcode)?;
304                                if matches!(toper.opcode, 0x02 | 0x06 | 0x07) {
305                                    let addr_pos = mem.pos;
306                                    let base_pos = patcher.output.stream_position()?;
307                                    let addr = base_pos + addr_pos as u64;
308                                    let data = toper
309                                        .operands
310                                        .iter()
311                                        .find_map(|operand| {
312                                            if let Operand::D(v) = operand {
313                                                Some(*v)
314                                            } else {
315                                                None
316                                            }
317                                        })
318                                        .ok_or(anyhow::anyhow!(
319                                            "Unexpected operand type in function re-write."
320                                        ))?;
321                                    new_need_patch_addresses.push((addr, data));
322                                }
323                                if toper.opcode == 0x03 {
324                                    let syscall_id =
325                                        if let Some(Operand::W(id)) = toper.operands.get(0) {
326                                            *id
327                                        } else {
328                                            anyhow::bail!(
329                                                "Invalid syscall operand at function index {}",
330                                                func_index
331                                            );
332                                        };
333                                    if syscall_id == thread_start_callid {
334                                        if tindex == 0 {
335                                            anyhow::bail!(
336                                                "ThreadStart syscall cannot be at function index 0"
337                                            );
338                                        }
339                                        let pre_func = &new_funcs[tindex - 1];
340                                        if pre_func.opcode == 0x0a {
341                                            let addr_pos = mem.pos - 5; // 1 for opcode, 4 for operand
342                                            let base_pos = patcher.output.stream_position()?;
343                                            let addr = base_pos + addr_pos as u64;
344                                            let data = pre_func
345                                                .operands
346                                                .get(0)
347                                                .and_then(|operand| {
348                                                    if let Operand::D(v) = operand {
349                                                        Some(*v)
350                                                    } else {
351                                                        None
352                                                    }
353                                                })
354                                                .ok_or(anyhow::anyhow!(
355                                                    "Unexpected operand type in function re-write."
356                                                ))?;
357                                            new_need_patch_addresses.push((addr, data));
358                                        }
359                                    }
360                                }
361                                for operand in &toper.operands {
362                                    match operand {
363                                        Operand::B(v) => mem.write_u8(*v)?,
364                                        Operand::W(v) => mem.write_u16(*v)?,
365                                        Operand::D(v) => mem.write_u32(*v)?,
366                                        Operand::F(v) => mem.write_f32(*v)?,
367                                        _ => {
368                                            return Err(anyhow::anyhow!(
369                                                "Unexpected operand type in function re-write."
370                                            ));
371                                        }
372                                    }
373                                }
374                                tindex += 1;
375                            }
376                        }
377                        let new_data = mem.into_inner();
378                        patcher.replace_bytes(0, &new_data)?;
379                        continue 'outer;
380                    }
381                    patcher.replace_bytes_with_write(ori_len, |writer| {
382                        writer.write_u8(len as u8)?;
383                        writer.write_all(&s)?;
384                        Ok(())
385                    })?;
386                }
387                cur_pos += operand.len(self.encoding)? as u64;
388            }
389            func_index += 1;
390        }
391        patcher.copy_up_to(self.reader.data.len() as u64)?;
392        for addr in need_pacth_addresses {
393            patcher.patch_u32_address(addr)?;
394        }
395        for (addr, data) in new_need_patch_addresses {
396            let new_data = patcher.map_offset(data as u64)? as u32;
397            patcher.output.write_u32_at(addr, new_data)?;
398        }
399        let script_len = self.reader.cpeek_u32_at(0)? as u64;
400        let new_script_len = patcher.map_offset(script_len)?;
401        patcher.patch_u32(0, new_script_len as u32)?;
402        // fix main script data position
403        patcher.patch_u32_address(script_len)?;
404        Ok(())
405    }
406
407    fn custom_export(&self, filename: &std::path::Path, encoding: Encoding) -> Result<()> {
408        let s = if self.custom_yaml {
409            serde_yaml_ng::to_string(&self.data)?
410        } else {
411            serde_json::to_string_pretty(&self.data)?
412        };
413        let e = encode_string(encoding, &s, false)?;
414        let mut writer = crate::utils::files::write_file(filename)?;
415        writer.write_all(&e)?;
416        Ok(())
417    }
418}