msg_tool\scripts\emote/
psb.rs

1//! Basic Handle for all emote PSB files.
2use super::rle::*;
3use crate::ext::io::*;
4use crate::ext::psb::*;
5use crate::scripts::base::*;
6use crate::types::*;
7use crate::utils::encoding::*;
8use crate::utils::files::*;
9use crate::utils::img::*;
10use anyhow::Result;
11use base64::Engine;
12use block_compression::BC7Settings;
13use clap::ValueEnum;
14use emote_psb::*;
15use libtlg_rs::*;
16use serde::{Deserialize, Serialize};
17use std::collections::HashMap;
18use std::io::{Read, Seek, Write};
19
20#[derive(Debug)]
21pub struct PsbBuilder {}
22
23impl PsbBuilder {
24    pub fn new() -> Self {
25        Self {}
26    }
27}
28
29impl ScriptBuilder for PsbBuilder {
30    fn default_encoding(&self) -> Encoding {
31        Encoding::Utf8
32    }
33
34    fn build_script(
35        &self,
36        buf: Vec<u8>,
37        _filename: &str,
38        encoding: Encoding,
39        _archive_encoding: Encoding,
40        config: &ExtraConfig,
41        _archive: Option<&Box<dyn Script>>,
42    ) -> Result<Box<dyn Script>> {
43        Ok(Box::new(Psb::new(MemReader::new(buf), encoding, config)?))
44    }
45
46    fn build_script_from_reader(
47        &self,
48        reader: Box<dyn ReadSeek>,
49        _filename: &str,
50        encoding: Encoding,
51        _archive_encoding: Encoding,
52        config: &ExtraConfig,
53        _archive: Option<&Box<dyn Script>>,
54    ) -> Result<Box<dyn Script>> {
55        Ok(Box::new(Psb::new(reader, encoding, config)?))
56    }
57
58    fn build_script_from_file(
59        &self,
60        filename: &str,
61        encoding: Encoding,
62        _archive_encoding: Encoding,
63        config: &ExtraConfig,
64        _archive: Option<&Box<dyn Script>>,
65    ) -> Result<Box<dyn Script>> {
66        let file = std::fs::File::open(filename)?;
67        let f = std::io::BufReader::new(file);
68        Ok(Box::new(Psb::new(f, encoding, config)?))
69    }
70
71    fn extensions(&self) -> &'static [&'static str] {
72        &["psb"]
73    }
74
75    fn script_type(&self) -> &'static ScriptType {
76        &ScriptType::EmotePsb
77    }
78
79    fn is_this_format(&self, _filename: &str, buf: &[u8], buf_len: usize) -> Option<u8> {
80        if buf_len >= 4 && buf.starts_with(b"PSB\0") {
81            return Some(10);
82        } else if buf_len >= 4 && buf.starts_with(&[0x04, 0x22, 0x4D, 0x18]) {
83            for i in 4..buf_len - 4 {
84                if buf[i..i + 4] == *b"PSB\0" {
85                    return Some(10);
86                }
87            }
88        }
89        None
90    }
91
92    fn can_create_file(&self) -> bool {
93        true
94    }
95
96    fn create_file<'a>(
97        &'a self,
98        filename: &'a str,
99        writer: Box<dyn WriteSeek + 'a>,
100        encoding: Encoding,
101        file_encoding: Encoding,
102        config: &ExtraConfig,
103    ) -> Result<()> {
104        create_file(filename, writer, encoding, file_encoding, config)
105    }
106}
107
108#[derive(Debug, ValueEnum, Clone, Copy)]
109pub enum BC7Config {
110    /// Ultra fast settings.
111    UltraFast,
112    /// Very fast settings.
113    VeryFast,
114    /// Fast settings.
115    Fast,
116    /// Basic settings.
117    Basic,
118    /// Slow settings.
119    Slow,
120}
121
122impl Default for BC7Config {
123    fn default() -> Self {
124        Self::Basic
125    }
126}
127
128impl Into<BC7Settings> for BC7Config {
129    fn into(self) -> BC7Settings {
130        match self {
131            Self::UltraFast => BC7Settings::alpha_ultrafast(),
132            Self::VeryFast => BC7Settings::alpha_very_fast(),
133            Self::Fast => BC7Settings::alpha_fast(),
134            Self::Basic => BC7Settings::alpha_basic(),
135            Self::Slow => BC7Settings::alpha_slow(),
136        }
137    }
138}
139
140#[derive(Debug)]
141pub struct Psb {
142    psb: VirtualPsbFixed,
143    encoding: Encoding,
144    config: ExtraConfig,
145}
146
147impl Psb {
148    pub fn new<R: Read + Seek>(
149        reader: R,
150        encoding: Encoding,
151        config: &ExtraConfig,
152    ) -> Result<Self> {
153        let psb = PsbReader::open_psb_v2(reader)?.to_psb_fixed();
154        Ok(Self {
155            psb,
156            encoding,
157            config: config.clone(),
158        })
159    }
160
161    fn output_resource(
162        &self,
163        folder_path: &std::path::PathBuf,
164        path: String,
165        data: &[u8],
166    ) -> Result<Resource> {
167        let mut res = Resource {
168            path,
169            ..Default::default()
170        };
171        if self.config.psb_process_tlg && is_valid_tlg(&data) {
172            let tlg = load_tlg(MemReaderRef::new(&data))?;
173            res.tlg = Some(TlgInfo::from_tlg(&tlg, self.encoding));
174            let outtype = self.config.image_type.unwrap_or(ImageOutputType::Png);
175            res.path = {
176                let mut pb = std::path::PathBuf::from(&res.path);
177                pb.set_extension(outtype.as_ref());
178                pb.to_string_lossy().to_string()
179            };
180            let path = folder_path.join(&res.path);
181            make_sure_dir_exists(&path)?;
182            let img = ImageData {
183                width: tlg.width as u32,
184                height: tlg.height as u32,
185                color_type: match tlg.color {
186                    TlgColorType::Bgr24 => ImageColorType::Bgr,
187                    TlgColorType::Bgra32 => ImageColorType::Bgra,
188                    TlgColorType::Grayscale8 => ImageColorType::Grayscale,
189                },
190                depth: 8,
191                data: tlg.data,
192            };
193            encode_img(img, outtype, &path.to_string_lossy(), &self.config)?;
194        } else {
195            let path = folder_path.join(&res.path);
196            make_sure_dir_exists(&path)?;
197            std::fs::write(&path, data)?;
198        }
199        Ok(res)
200    }
201
202    fn output_rle_resource(
203        &self,
204        folder_path: &std::path::PathBuf,
205        path: String,
206        data: &[u8],
207        width: i64,
208        height: i64,
209    ) -> Result<Resource> {
210        let mut res = Resource {
211            path,
212            rle: Some(RLPixelInfo { width, height }),
213            ..Default::default()
214        };
215        let decompressed = rl_decompress(MemReaderRef::new(data), 4, None)?;
216        let outtype = self.config.image_type.unwrap_or(ImageOutputType::Png);
217        res.path = {
218            let mut pb = std::path::PathBuf::from(&res.path);
219            pb.set_extension(outtype.as_ref());
220            pb.to_string_lossy().to_string()
221        };
222        let path = folder_path.join(&res.path);
223        make_sure_dir_exists(&path)?;
224        let img = ImageData {
225            width: width as u32,
226            height: height as u32,
227            color_type: ImageColorType::Bgra,
228            depth: 8,
229            data: decompressed,
230        };
231        encode_img(img, outtype, &path.to_string_lossy(), &self.config)?;
232        Ok(res)
233    }
234
235    fn output_bc7_resource(
236        &self,
237        folder_path: &std::path::PathBuf,
238        path: String,
239        data: &[u8],
240        width: i64,
241        height: i64,
242    ) -> Result<Resource> {
243        let mut res = Resource {
244            path,
245            bc7: Some(BC7PixelInfo { width, height }),
246            ..Default::default()
247        };
248        let dst_size = (width * height * 4) as usize;
249        let mut decompressed_block = vec![0u8; dst_size];
250        let variant = block_compression::CompressionVariant::BC7(BC7Settings::alpha_basic());
251        let blocks_bytes = variant.blocks_byte_size(width as u32, height as u32);
252        if data.len() != blocks_bytes {
253            return Err(anyhow::anyhow!(
254                "BC7 compressed data size {} does not match expected size {} for image size {}x{}",
255                data.len(),
256                blocks_bytes,
257                width,
258                height
259            ));
260        }
261        block_compression::decode::decompress_blocks_as_rgba8(
262            variant,
263            width as u32,
264            height as u32,
265            data,
266            &mut decompressed_block,
267        );
268        let outtype = self.config.image_type.unwrap_or(ImageOutputType::Png);
269        res.path = {
270            let mut pb = std::path::PathBuf::from(&res.path);
271            pb.set_extension(outtype.as_ref());
272            pb.to_string_lossy().to_string()
273        };
274        let path = folder_path.join(&res.path);
275        make_sure_dir_exists(&path)?;
276        let img = ImageData {
277            width: width as u32,
278            height: height as u32,
279            color_type: ImageColorType::Rgba,
280            depth: 8,
281            data: decompressed_block,
282        };
283        encode_img(img, outtype, &path.to_string_lossy(), &self.config)?;
284        Ok(res)
285    }
286}
287
288#[derive(Debug, Deserialize, Serialize)]
289struct TlgInfo {
290    metadata: HashMap<String, String>,
291}
292
293impl TlgInfo {
294    fn from_tlg(tlg: &Tlg, encoding: Encoding) -> Self {
295        let mut metadata = HashMap::new();
296        for (k, v) in &tlg.tags {
297            let k = if let Ok(s) = decode_to_string(encoding, &k, true) {
298                s
299            } else {
300                format!(
301                    "base64:{}",
302                    base64::engine::general_purpose::STANDARD.encode(k)
303                )
304            };
305            let v = if let Ok(s) = decode_to_string(encoding, &v, true) {
306                s
307            } else {
308                format!(
309                    "base64:{}",
310                    base64::engine::general_purpose::STANDARD.encode(v)
311                )
312            };
313            metadata.insert(k, v);
314        }
315        Self { metadata }
316    }
317
318    fn to_tlg_tags(&self, encoding: Encoding) -> Result<HashMap<Vec<u8>, Vec<u8>>> {
319        let mut tags = HashMap::new();
320        for (k, v) in &self.metadata {
321            let k = if k.starts_with("base64:") {
322                base64::engine::general_purpose::STANDARD.decode(&k[7..])?
323            } else {
324                encode_string(encoding, k, false)?
325            };
326            let v = if v.starts_with("base64:") {
327                base64::engine::general_purpose::STANDARD.decode(&v[7..])?
328            } else {
329                encode_string(encoding, v, false)?
330            };
331            tags.insert(k, v);
332        }
333        Ok(tags)
334    }
335}
336
337#[derive(Debug, Deserialize, Serialize)]
338struct RLPixelInfo {
339    width: i64,
340    height: i64,
341}
342
343#[derive(Debug, Deserialize, Serialize)]
344struct BC7PixelInfo {
345    width: i64,
346    height: i64,
347}
348
349#[derive(Debug, Default, Deserialize, Serialize)]
350struct Resource {
351    path: String,
352    #[serde(skip_serializing_if = "Option::is_none")]
353    tlg: Option<TlgInfo>,
354    #[serde(skip_serializing_if = "Option::is_none")]
355    rle: Option<RLPixelInfo>,
356    #[serde(skip_serializing_if = "Option::is_none")]
357    bc7: Option<BC7PixelInfo>,
358}
359
360impl Script for Psb {
361    fn default_output_script_type(&self) -> OutputScriptType {
362        OutputScriptType::Custom
363    }
364
365    fn is_output_supported(&self, output: OutputScriptType) -> bool {
366        matches!(output, OutputScriptType::Custom)
367    }
368
369    fn default_format_type(&self) -> FormatOptions {
370        FormatOptions::None
371    }
372
373    fn custom_output_extension<'a>(&'a self) -> &'a str {
374        "json"
375    }
376
377    fn custom_export(&self, filename: &std::path::Path, encoding: Encoding) -> Result<()> {
378        let mut data = self.psb.to_json();
379        let mut resources = Vec::new();
380        let mut extra_resources = Vec::new();
381        let folder_path = {
382            let mut pb = filename.to_path_buf();
383            pb.set_extension("");
384            pb
385        };
386        for (i, data) in self.psb.resources().iter().enumerate() {
387            let i = i as u64;
388            let res_path = self.psb.root().find_resource_key(i, vec![]);
389            if let Some(path) = &res_path {
390                if path.len() >= 2 && *path.last().unwrap() == "pixel" {
391                    let pb_data = self.psb.root();
392                    let mut pb_data = &pb_data[*path.first().unwrap()];
393                    for p in path.iter().take(path.len() - 1).skip(1) {
394                        pb_data = &pb_data[*p];
395                    }
396                    let width = pb_data["width"].as_i64();
397                    let height = pb_data["height"].as_i64();
398                    let compress = pb_data["compress"].as_str();
399                    let type_ = pb_data["type"].as_str();
400                    if compress.is_some_and(|s| s == "RL") && (width.is_none() || height.is_none())
401                    {
402                        eprintln!(
403                            "Warning: Resource {:?} is marked as RL compressed but width/height is missing (width={:?}, height={:?})",
404                            path, pb_data["width"], pb_data["height"]
405                        );
406                        crate::COUNTER.inc_warning();
407                    }
408                    if type_.is_some_and(|s| s == "BC7") && (width.is_none() || height.is_none()) {
409                        eprintln!(
410                            "Warning: Resource {:?} is marked as BC7 compressed but width/height is missing (width={:?}, height={:?})",
411                            path, pb_data["width"], pb_data["height"]
412                        );
413                        crate::COUNTER.inc_warning();
414                    }
415                    if let (Some(w), Some(h), Some(c)) = (width, height, compress) {
416                        if c == "RL" {
417                            let res_name: Vec<_> = path
418                                .iter()
419                                .take(path.len() - 1)
420                                .map(|s| s.to_string())
421                                .collect();
422                            let res_name = res_name.join("/");
423                            let res_name = sanitize_path(&res_name);
424                            let res =
425                                self.output_rle_resource(&folder_path, res_name, data, w, h)?;
426                            resources.push(res);
427                            continue;
428                        }
429                    }
430                    if let (Some(w), Some(h), Some(t)) = (width, height, type_) {
431                        if t == "BC7" {
432                            let res_name: Vec<_> = path
433                                .iter()
434                                .take(path.len() - 1)
435                                .map(|s| s.to_string())
436                                .collect();
437                            let res_name = res_name.join("/");
438                            let res_name = sanitize_path(&res_name);
439                            let res =
440                                self.output_bc7_resource(&folder_path, res_name, data, w, h)?;
441                            resources.push(res);
442                            continue;
443                        }
444                    }
445                }
446            }
447            let res_name = res_path
448                .map(|s| s.join("/"))
449                .unwrap_or(format!("res_{}", i));
450            let res_name = sanitize_path(&res_name);
451            let res = self.output_resource(&folder_path, res_name, data)?;
452            resources.push(res);
453        }
454        for (i, data) in self.psb.extra().iter().enumerate() {
455            let i = i as u64;
456            let res_name = self
457                .psb
458                .root()
459                .find_resource_key(i, vec![])
460                .map(|s| format!("extra_{}", s.join("/")))
461                .unwrap_or(format!("extra_res_{}", i));
462            let res_name = sanitize_path(&res_name);
463            let res = self.output_resource(&folder_path, res_name, data)?;
464            extra_resources.push(res);
465        }
466        data["resources"] = json::parse(&serde_json::to_string(&resources)?)?;
467        data["extra_resources"] = json::parse(&serde_json::to_string(&extra_resources)?)?;
468        let s = json::stringify_pretty(data, 2);
469        let s = encode_string(encoding, &s, false)?;
470        let mut file = std::fs::File::create(filename)?;
471        file.write_all(&s)?;
472        Ok(())
473    }
474
475    fn custom_import<'a>(
476        &'a self,
477        custom_filename: &'a str,
478        file: Box<dyn WriteSeek + 'a>,
479        encoding: Encoding,
480        output_encoding: Encoding,
481    ) -> Result<()> {
482        create_file(
483            custom_filename,
484            file,
485            encoding,
486            output_encoding,
487            &self.config,
488        )
489    }
490}
491
492fn read_resource(
493    folder_path: &std::path::PathBuf,
494    res: &Resource,
495    encoding: Encoding,
496    cfg: &ExtraConfig,
497) -> Result<Vec<u8>> {
498    if let Some(tlg) = &res.tlg {
499        let path = folder_path.join(&res.path);
500        let imgfmt = ImageOutputType::try_from(path.as_path())?;
501        let mut img = decode_img(imgfmt, &path.to_string_lossy())?;
502        if img.depth != 8 {
503            return Err(anyhow::anyhow!(
504                "Only 8-bit images are supported for TLG conversion"
505            ));
506        }
507        let color_type = match img.color_type {
508            ImageColorType::Bgr => TlgColorType::Bgr24,
509            ImageColorType::Bgra => TlgColorType::Bgra32,
510            ImageColorType::Grayscale => TlgColorType::Grayscale8,
511            ImageColorType::Rgb => {
512                convert_rgb_to_bgr(&mut img)?;
513                TlgColorType::Bgr24
514            }
515            ImageColorType::Rgba => {
516                convert_rgba_to_bgra(&mut img)?;
517                TlgColorType::Bgra32
518            }
519        };
520        let tlg = Tlg {
521            width: img.width,
522            height: img.height,
523            version: 5,
524            color: color_type,
525            data: img.data,
526            tags: tlg.to_tlg_tags(encoding)?,
527        };
528        let mut writer = MemWriter::new();
529        save_tlg(&tlg, &mut writer)?;
530        Ok(writer.into_inner())
531    } else if let Some(rle) = &res.rle {
532        let path = folder_path.join(&res.path);
533        let imgfmt = ImageOutputType::try_from(path.as_path())?;
534        let mut img = decode_img(imgfmt, &path.to_string_lossy())?;
535        if img.depth != 8 {
536            return Err(anyhow::anyhow!(
537                "Only 8-bit images are supported for RLE conversion"
538            ));
539        }
540        if img.color_type == ImageColorType::Rgba {
541            convert_rgba_to_bgra(&mut img)?;
542        } else if img.color_type == ImageColorType::Rgb {
543            convert_rgb_to_bgr(&mut img)?;
544            convert_bgr_to_bgra(&mut img)?;
545        } else if img.color_type == ImageColorType::Bgr {
546            convert_bgr_to_bgra(&mut img)?;
547        }
548        if img.color_type != ImageColorType::Bgra {
549            return Err(anyhow::anyhow!(
550                "Only BGRA images are supported for RLE conversion"
551            ));
552        }
553        if img.width as i64 != rle.width {
554            eprintln!(
555                "Warning: Image width {} does not match RLE width {}",
556                img.width, rle.width
557            );
558            crate::COUNTER.inc_warning();
559        }
560        if img.height as i64 != rle.height {
561            eprintln!(
562                "Warning: Image height {} does not match RLE height {}",
563                img.height, rle.height
564            );
565            crate::COUNTER.inc_warning();
566        }
567        let compressed = rl_compress(MemReaderRef::new(&img.data), 4)?;
568        Ok(compressed)
569    } else if let Some(bc7) = &res.bc7 {
570        let path = folder_path.join(&res.path);
571        let imgfmt = ImageOutputType::try_from(path.as_path())?;
572        let mut img = decode_img(imgfmt, &path.to_string_lossy())?;
573        if img.depth != 8 {
574            return Err(anyhow::anyhow!(
575                "Only 8-bit images are supported for BC7 conversion"
576            ));
577        }
578        if img.width % 4 != 0 || img.height % 4 != 0 {
579            return Err(anyhow::anyhow!(
580                "Image dimensions must be multiples of 4 for BC7 conversion (width={}, height={})",
581                img.width,
582                img.height
583            ));
584        }
585        if bc7.height != img.height as i64 {
586            eprintln!(
587                "Warning: Image height {} does not match BC7 height {}",
588                img.height, bc7.height
589            );
590            crate::COUNTER.inc_warning();
591        }
592        if bc7.width != img.width as i64 {
593            eprintln!(
594                "Warning: Image width {} does not match BC7 width {}",
595                img.width, bc7.width
596            );
597            crate::COUNTER.inc_warning();
598        }
599        convert_to_rgba(&mut img)?;
600        let variant = block_compression::CompressionVariant::BC7(cfg.bc7.into());
601        let dst_size = variant.blocks_byte_size(img.width, img.height);
602        let mut compressed = vec![0u8; dst_size as usize];
603        block_compression::encode::compress_rgba8(
604            variant,
605            &img.data,
606            &mut compressed,
607            img.width,
608            img.height,
609            img.width * 4,
610        );
611        Ok(compressed)
612    } else {
613        let path = folder_path.join(&res.path);
614        Ok(std::fs::read(&path)?)
615    }
616}
617
618fn create_file<'a>(
619    custom_filename: &'a str,
620    mut writer: Box<dyn WriteSeek + 'a>,
621    encoding: Encoding,
622    output_encoding: Encoding,
623    cfg: &ExtraConfig,
624) -> Result<()> {
625    let input = read_file(custom_filename)?;
626    let s = decode_to_string(output_encoding, &input, true)?;
627    let data = json::parse(&s)?;
628    let resources: Vec<Resource> = serde_json::from_str(&data["resources"].dump())?;
629    let extra_resources: Vec<Resource> = serde_json::from_str(&data["extra_resources"].dump())?;
630    let mut psb = VirtualPsbFixed::with_json(&data)?;
631    psb.header_mut().encryption = 0; // We don't support encryption.
632    let folder_path = {
633        let mut pb = std::path::PathBuf::from(custom_filename);
634        pb.set_extension("");
635        pb
636    };
637    for res in resources {
638        let res = read_resource(&folder_path, &res, encoding, cfg)?;
639        psb.resources_mut().push(res);
640    }
641    for res in extra_resources {
642        let res = read_resource(&folder_path, &res, encoding, cfg)?;
643        psb.extra_mut().push(res);
644    }
645    let psb = psb.to_psb(false);
646    psb.finish_v4(&mut writer)
647        .map_err(|e| anyhow::anyhow!("Failed to write PSB file: {:?}", e))?;
648    Ok(())
649}
msg_tool\scripts\emote/psb.rs

msg_tool\scripts\emote/
psb.rs
