From ea1e20f6e444fe524f1dd1591e57499c7e892d11 Mon Sep 17 00:00:00 2001 From: lifegpc Date: Tue, 29 Jul 2025 23:27:59 +0800 Subject: [PATCH] Add CRX import support --- src/args.rs | 29 ++- src/main.rs | 6 + src/scripts/cat_system/cst.rs | 2 +- src/scripts/circus/image/crx.rs | 323 +++++++++++++++++++++++++++++++- src/types.rs | 8 +- src/utils/img.rs | 41 ++++ 6 files changed, 403 insertions(+), 6 deletions(-) diff --git a/src/args.rs b/src/args.rs index 7d60099..8fad4ec 100644 --- a/src/args.rs +++ b/src/args.rs @@ -16,6 +16,17 @@ fn parse_compression_level(level: &str) -> Result { clap_num::number_range(level, 0, 9) } +#[cfg(feature = "zstd")] +fn parse_zstd_compression_level(level: &str) -> Result { + let lower = level.to_ascii_lowercase(); + if lower == "default" { + return Ok(3); + } else if lower == "best" { + return Ok(22); + } + clap_num::number_range(level, 0, 22) +} + /// Tools for export and import scripts #[derive(Parser, Debug)] #[clap(group = ArgGroup::new("encodingg").multiple(false), group = ArgGroup::new("output_encodingg").multiple(false), group = ArgGroup::new("archive_encodingg").multiple(false), group = ArgGroup::new("artemis_indentg").multiple(false))] @@ -200,12 +211,24 @@ pub struct Arg { /// If not specified, the first language will be used. pub cat_system_cstl_lang: Option, #[cfg(feature = "flate2")] - #[arg(short = 'z', long, global = true, value_name = "LEVEL", value_parser = parse_compression_level)] - /// Zlib compression level. Default is 6. 0 means no compression, 9 means best compression. - pub zlib_compression_level: Option, + #[arg(short = 'z', long, global = true, value_name = "LEVEL", value_parser = parse_compression_level, default_value_t = 6)] + /// Zlib compression level. 0 means no compression, 9 means best compression. + pub zlib_compression_level: u32, #[cfg(feature = "image")] #[arg(short = 'g', long, global = true, value_enum, default_value_t = PngCompressionLevel::Fast)] pub png_compression_level: PngCompressionLevel, + #[cfg(feature = "circus-img")] + #[arg(long, global = true, action = ArgAction::SetTrue)] + /// Keep original BPP when importing Circus CRX images. + pub circus_crx_keep_original_bpp: bool, + #[cfg(feature = "circus-img")] + #[arg(long, global = true, action = ArgAction::SetTrue)] + /// Use zstd compression for Circus CRX images. (CIRCUS Engine don't support this. Hook is required.) + pub circus_crx_zstd: bool, + #[cfg(feature = "zstd")] + #[arg(short = 'Z', long, global = true, value_name = "LEVEL", value_parser = parse_zstd_compression_level, default_value_t = 3)] + /// Zstd compression level. 0 means default compression level (3), 22 means best compression. + pub zstd_compression_level: i32, #[command(subcommand)] /// Command pub command: Command, diff --git a/src/main.rs b/src/main.rs index 8d066ca..8149236 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1459,6 +1459,12 @@ fn main() { zlib_compression_level: arg.zlib_compression_level, #[cfg(feature = "image")] png_compression_level: arg.png_compression_level, + #[cfg(feature = "circus-img")] + circus_crx_keep_original_bpp: arg.circus_crx_keep_original_bpp, + #[cfg(feature = "circus-img")] + circus_crx_zstd: arg.circus_crx_zstd, + #[cfg(feature = "zstd")] + zstd_compression_level: arg.zstd_compression_level, }; match &arg.command { args::Command::Export { input, output } => { diff --git a/src/scripts/cat_system/cst.rs b/src/scripts/cat_system/cst.rs index c61f442..cbc67e5 100644 --- a/src/scripts/cat_system/cst.rs +++ b/src/scripts/cat_system/cst.rs @@ -169,7 +169,7 @@ impl CstScript { data: file, compressed: compressed_size != 0, strings, - compress_level: config.zlib_compression_level.unwrap_or(6), + compress_level: config.zlib_compression_level, }) } } diff --git a/src/scripts/circus/image/crx.rs b/src/scripts/circus/image/crx.rs index 09d3a63..fb827d5 100644 --- a/src/scripts/circus/image/crx.rs +++ b/src/scripts/circus/image/crx.rs @@ -1,6 +1,7 @@ use crate::ext::io::*; use crate::scripts::base::*; use crate::types::*; +use crate::utils::img::*; use crate::utils::struct_pack::*; use anyhow::Result; use msg_tool_macro::*; @@ -82,6 +83,10 @@ pub struct CrxImage { header: Header, color_type: ImageColorType, data: Vec, + compress_level: u32, + keep_original_bpp: bool, + zstd: bool, + zstd_compression_level: i32, } impl std::fmt::Debug for CrxImage { @@ -95,7 +100,7 @@ impl std::fmt::Debug for CrxImage { } impl CrxImage { - pub fn new(data: T, _config: &ExtraConfig) -> Result { + pub fn new(data: T, config: &ExtraConfig) -> Result { let mut reader = data; let mut magic = [0; 4]; reader.read_exact(&mut magic)?; @@ -138,6 +143,10 @@ impl CrxImage { header, color_type, data: uncompessed, + compress_level: config.zlib_compression_level, + keep_original_bpp: config.circus_crx_keep_original_bpp, + zstd: config.circus_crx_zstd, + zstd_compression_level: config.zstd_compression_level, }) } @@ -336,6 +345,201 @@ impl CrxImage { } Ok(()) } + + fn encode_row0( + dst: &mut Vec, + mut dst_p: usize, + src: &[u8], + width: u16, + pixel_size: u8, + y: u16, + ) -> Result { + let pixel_size = pixel_size as usize; + let mut src_p = y as usize * width as usize * pixel_size; + for _ in 0..pixel_size { + dst[dst_p] = src[src_p]; + dst_p += 1; + src_p += 1; + } + for _ in 1..width { + for _ in 0..pixel_size { + dst[dst_p] = src[src_p].wrapping_sub(src[src_p - pixel_size]); + dst_p += 1; + src_p += 1; + } + } + Ok(dst_p) + } + + fn encode_row1( + dst: &mut Vec, + mut dst_p: usize, + src: &[u8], + width: u16, + pixel_size: u8, + y: u16, + ) -> Result { + let pixel_size = pixel_size as usize; + let mut src_p = y as usize * width as usize * pixel_size; + let mut prev_row_p = (y as usize - 1) * width as usize * pixel_size; + for _ in 0..width { + for _ in 0..pixel_size { + dst[dst_p] = src[src_p].wrapping_sub(src[prev_row_p]); + dst_p += 1; + src_p += 1; + prev_row_p += 1; + } + } + Ok(dst_p) + } + + fn encode_row2( + dst: &mut Vec, + mut dst_p: usize, + src: &[u8], + width: u16, + pixel_size: u8, + y: u16, + ) -> Result { + let pixel_size = pixel_size as usize; + let mut src_p = y as usize * width as usize * pixel_size; + let mut prev_row_p = (y as usize - 1) * width as usize * pixel_size; + for _ in 0..pixel_size { + dst[dst_p] = src[src_p]; + dst_p += 1; + src_p += 1; + } + for _ in 1..width { + for _ in 0..pixel_size { + dst[dst_p] = src[src_p].wrapping_sub(src[prev_row_p]); + dst_p += 1; + src_p += 1; + prev_row_p += 1; + } + } + Ok(dst_p) + } + + fn encode_row3( + dst: &mut Vec, + mut dst_p: usize, + src: &[u8], + width: u16, + pixel_size: u8, + y: u16, + ) -> Result { + let pixel_size = pixel_size as usize; + let mut src_p = y as usize * width as usize * pixel_size; + let mut prev_row_p = (y as usize - 1) * width as usize * pixel_size + pixel_size; + for _ in 0..width - 1 { + for _ in 0..pixel_size { + dst[dst_p] = src[src_p].wrapping_sub(src[prev_row_p]); + dst_p += 1; + src_p += 1; + prev_row_p += 1; + } + } + for _ in 0..pixel_size { + dst[dst_p] = src[src_p]; + dst_p += 1; + src_p += 1; + } + Ok(dst_p) + } + + fn encode_row4( + dst: &mut Vec, + mut dst_p: usize, + src: &[u8], + width: u16, + pixel_size: u8, + y: u16, + ) -> Result { + let pixel_size = pixel_size as usize; + let src_p = y as usize * width as usize * pixel_size; + for offset in 0..pixel_size { + let mut src_c = src_p + offset; + let mut remaining = width; + let value = src[src_c]; + src_c += pixel_size; + dst[dst_p] = value; + dst_p += 1; + remaining -= 1; + if remaining == 0 { + continue; + } + let mut count = 0; + loop { + if count as u16 >= remaining || count >= 255 || src[src_c] != value { + break; + } + src_c += pixel_size; + count += 1; + } + if count > 0 { + dst[dst_p] = value; + dst_p += 1; + dst[dst_p] = count; + dst_p += 1; + remaining -= count as u16; + } + while remaining > 0 { + let value = src[src_c]; + src_c += pixel_size; + dst[dst_p] = value; + dst_p += 1; + remaining -= 1; + if remaining == 0 { + break; + } + let mut count = 0; + loop { + if count as u16 >= remaining || count >= 255 || src[src_c] != value { + break; + } + src_c += pixel_size; + count += 1; + } + if count > 0 { + dst[dst_p] = value; + dst_p += 1; + dst[dst_p] = count; + dst_p += 1; + remaining -= count as u16; + } + } + } + Ok(dst_p) + } + + fn encode_image_origin( + src: &[u8], + width: u16, + height: u16, + pixel_size: u8, + row_type: &[u8], + ) -> Result> { + if row_type.len() != height as usize { + return Err(anyhow::anyhow!("Row type length does not match height")); + } + let size = width as usize * height as usize * pixel_size as usize + height as usize; + let mut dst = vec![0; size]; + let mut dst_p = 0; + for y in 0..height { + let data = row_type[y as usize]; + dst[dst_p] = data; + dst_p += 1; + dst_p = match data { + 0 => Self::encode_row0(&mut dst, dst_p, src, width, pixel_size, y)?, + 1 => Self::encode_row1(&mut dst, dst_p, src, width, pixel_size, y)?, + 2 => Self::encode_row2(&mut dst, dst_p, src, width, pixel_size, y)?, + 3 => Self::encode_row3(&mut dst, dst_p, src, width, pixel_size, y)?, + 4 => Self::encode_row4(&mut dst, dst_p, src, width, pixel_size, y)?, + _ => return Err(anyhow::anyhow!("Invalid row type: {}", data)), + }; + } + Ok(dst) + } } impl Script for CrxImage { @@ -386,4 +590,121 @@ impl Script for CrxImage { data, }) } + + fn import_image<'a>( + &'a self, + mut data: ImageData, + mut file: Box, + ) -> Result<()> { + let mut color_type = match data.color_type { + ImageColorType::Bgr => ImageColorType::Bgr, + ImageColorType::Bgra => ImageColorType::Bgra, + ImageColorType::Rgb => { + convert_rgb_to_bgr(&mut data)?; + ImageColorType::Bgr + } + ImageColorType::Rgba => { + convert_rgba_to_bgra(&mut data)?; + ImageColorType::Bgra + } + _ => { + return Err(anyhow::anyhow!( + "Unsupported color type: {:?}", + data.color_type + )); + } + }; + if data.width != self.header.width as u32 { + return Err(anyhow::anyhow!( + "Image width does not match: expected {}, got {}", + self.header.width, + data.width + )); + } + if data.height != self.header.height as u32 { + return Err(anyhow::anyhow!( + "Image height does not match: expected {}, got {}", + self.header.height, + data.height + )); + } + if data.depth != 8 { + return Err(anyhow::anyhow!("Image depth must be 8, got {}", data.depth)); + } + if data.color_type != self.color_type && self.keep_original_bpp { + if self.color_type == ImageColorType::Bgr { + convert_bgra_to_bgr(&mut data)?; + } else if self.color_type == ImageColorType::Bgra { + convert_bgr_to_bgra(&mut data)?; + } else { + return Err(anyhow::anyhow!( + "Unsupported color type for import: {:?}", + self.color_type + )); + } + color_type = self.color_type; + } + let mut new_header = self.header.clone(); + new_header.bpp = match color_type { + ImageColorType::Bgr => 0, + ImageColorType::Bgra => 1, + _ => return Err(anyhow::anyhow!("Unsupported color type: {:?}", color_type)), + }; + new_header.flags |= 0x10; // Force add compressed data length + let pixel_size = color_type.bpp(1) as u8; + if color_type == ImageColorType::Bgra && self.header.mode != 1 { + let alpha_flip = if self.header.mode == 2 { 0 } else { 0xFF }; + for i in (0..data.data.len()).step_by(4) { + let b = data.data[i]; + let g = data.data[i + 1]; + let r = data.data[i + 2]; + let a = data.data[i + 3]; + data.data[i] = a ^ alpha_flip; + data.data[i + 1] = b; + data.data[i + 2] = g; + data.data[i + 3] = r; + } + } + let mut row_type = Vec::with_capacity(self.header.height as usize); + let mut dst = vec![ + 0; + self.header.width as usize + * self.header.height as usize + * self.color_type.bpp(1) as usize + ]; + Self::decode_image( + &mut dst, + &self.data, + self.header.width, + self.header.height, + self.color_type.bpp(1) as u8, + &mut row_type, + )?; + let encoded = Self::encode_image_origin( + &data.data, + new_header.width, + new_header.height, + pixel_size, + &row_type, + )?; + let compressed = if self.zstd { + let mut encoder = zstd::Encoder::new(MemWriter::new(), self.zstd_compression_level)?; + encoder.write_all(&encoded)?; + let compressed_data = encoder.finish()?; + compressed_data.into_inner() + } else { + let mut encoder = flate2::write::ZlibEncoder::new( + MemWriter::new(), + flate2::Compression::new(self.compress_level), + ); + encoder.write_all(&encoded)?; + let compressed_data = encoder.finish()?; + compressed_data.into_inner() + }; + file.write_all(b"CRXG")?; + new_header.pack(&mut file, false, Encoding::Utf8)?; + file.write_u32(compressed.len() as u32)?; + file.write_all(&compressed)?; + Ok(()) + } } diff --git a/src/types.rs b/src/types.rs index 3ae8db9..638a06f 100644 --- a/src/types.rs +++ b/src/types.rs @@ -236,9 +236,15 @@ pub struct ExtraConfig { #[cfg(feature = "cat-system")] pub cat_system_cstl_lang: Option, #[cfg(feature = "flate2")] - pub zlib_compression_level: Option, + pub zlib_compression_level: u32, #[cfg(feature = "image")] pub png_compression_level: PngCompressionLevel, + #[cfg(feature = "circus-img")] + pub circus_crx_keep_original_bpp: bool, + #[cfg(feature = "circus-img")] + pub circus_crx_zstd: bool, + #[cfg(feature = "zstd")] + pub zstd_compression_level: i32, } #[derive(Clone, Copy, Debug, ValueEnum, PartialEq, Eq, PartialOrd, Ord)] diff --git a/src/utils/img.rs b/src/utils/img.rs index dfb5953..645b5b6 100644 --- a/src/utils/img.rs +++ b/src/utils/img.rs @@ -16,6 +16,27 @@ pub fn reverse_alpha_values(data: &mut ImageData) -> Result<()> { Ok(()) } +pub fn convert_bgr_to_bgra(data: &mut ImageData) -> Result<()> { + if data.color_type != ImageColorType::Bgr { + return Err(anyhow::anyhow!("Image is not BGR")); + } + if data.depth != 8 { + return Err(anyhow::anyhow!( + "BGR to BGRA conversion only supports 8-bit depth" + )); + } + let mut new_data = Vec::with_capacity(data.data.len() / 3 * 4); + for chunk in data.data.chunks_exact(3) { + new_data.push(chunk[0]); // B + new_data.push(chunk[1]); // G + new_data.push(chunk[2]); // R + new_data.push(255); // A + } + data.data = new_data; + data.color_type = ImageColorType::Bgra; + Ok(()) +} + pub fn convert_bgr_to_rgb(data: &mut ImageData) -> Result<()> { if data.color_type != ImageColorType::Bgr { return Err(anyhow::anyhow!("Image is not BGR")); @@ -34,6 +55,26 @@ pub fn convert_bgr_to_rgb(data: &mut ImageData) -> Result<()> { Ok(()) } +pub fn convert_bgra_to_bgr(data: &mut ImageData) -> Result<()> { + if data.color_type != ImageColorType::Bgra { + return Err(anyhow::anyhow!("Image is not BGRA")); + } + if data.depth != 8 { + return Err(anyhow::anyhow!( + "BGRA to BGR conversion only supports 8-bit depth" + )); + } + let mut new_data = Vec::with_capacity(data.data.len() / 4 * 3); + for chunk in data.data.chunks_exact(4) { + new_data.push(chunk[0]); // B + new_data.push(chunk[1]); // G + new_data.push(chunk[2]); // R + } + data.data = new_data; + data.color_type = ImageColorType::Bgr; + Ok(()) +} + pub fn convert_bgra_to_rgba(data: &mut ImageData) -> Result<()> { if data.color_type != ImageColorType::Bgra { return Err(anyhow::anyhow!("Image is not BGRA"));