Add CxEncryption support (WIP) (tested game: https://vndb.org/v5502 )

src/scripts/kirikiri/archive/xp3/crypt/mod.rs

@@ -0,0 +1,512 @@
+mod cx;
+
+use super::archive::*;
+use crate::ext::io::*;
+use crate::scripts::base::*;
+use crate::types::*;
+use crate::utils::encoding::*;
+use crate::utils::serde_base64bytes::*;
+use anyhow::Result;
+use serde::Deserialize;
+use std::collections::{BTreeMap, HashMap};
+use std::io::{Read, Seek, SeekFrom};
+use std::sync::Arc;
+
+pub trait Crypt: std::fmt::Debug {
+    /// Initializes the cryptographic context for the archive.
+    fn init(&self, _archive: &mut Xp3Archive) -> Result<()> {
+        Ok(())
+    }
+
+    /// Read a entry name from archive index
+    fn read_name<'a>(&self, reader: &mut Box<dyn Read + 'a>) -> Result<(String, u64)> {
+        let name_length = reader.read_u16()?;
+        let name = reader.read_exact_vec(name_length as usize * 2)?;
+        Ok((
+            decode_to_string(Encoding::Utf16LE, &name, true)?,
+            name_length as u64 * 2 + 2,
+        ))
+    }
+
+    /// Decrypts the given stream of data for the specified entry and segment.
+    fn decrypt<'a>(
+        &self,
+        _entry: &Xp3Entry,
+        _cur_seg: &Segment,
+        _stream: Box<dyn Read + 'a>,
+    ) -> Result<Box<dyn ReadDebug + 'a>> {
+        Err(anyhow::anyhow!("This crypt does not support decrypt"))
+    }
+
+    /// Decrypts the given stream of data for the specified entry and segment, with seek support.
+    fn decrypt_with_seek<'a>(
+        &self,
+        _entry: &Xp3Entry,
+        _cur_seg: &Segment,
+        _stream: Box<dyn ReadSeek + 'a>,
+    ) -> Result<Box<dyn ReadSeek + 'a>> {
+        Err(anyhow::anyhow!(
+            "This crypt does not support decrypt with seek"
+        ))
+    }
+
+    /// Returns true if this crypt support decrypt
+    fn decrypt_supported(&self) -> bool {
+        false
+    }
+
+    /// Returns true if this crypt support seek when decrypting
+    fn decrypt_seek_supported(&self) -> bool {
+        false
+    }
+}
+
+#[derive(Clone, Debug, Deserialize)]
+#[serde(rename_all = "PascalCase")]
+pub struct CxSchema {
+    mask: u32,
+    offset: u32,
+    prolog_order: Base64Bytes,
+    odd_branch_order: Base64Bytes,
+    even_branch_order: Base64Bytes,
+    control_block_name: Option<String>,
+    tpm_file_name: Option<String>,
+}
+
+#[derive(Clone, Debug, Deserialize)]
+#[serde(rename_all = "PascalCase", tag = "$type")]
+enum CryptType {
+    NoCrypt,
+    FateCrypt,
+    MizukakeCrypt,
+    HashCrypt,
+    #[serde(rename_all = "PascalCase")]
+    XorCrypt {
+        key: u8,
+    },
+    FlyingShineCrypt,
+    CxEncryption(CxSchema),
+}
+
+#[derive(Clone, Debug, Deserialize)]
+#[serde(rename_all = "PascalCase")]
+pub struct Schema {
+    #[serde(flatten)]
+    crypt: CryptType,
+    title: Option<String>,
+}
+
+impl Schema {
+    pub fn create_crypt(&self, filename: &str) -> Result<Box<dyn Crypt>> {
+        Ok(match &self.crypt {
+            CryptType::NoCrypt => Box::new(NoCrypt::new()),
+            CryptType::FateCrypt => Box::new(FateCrypt::new()),
+            CryptType::MizukakeCrypt => Box::new(MizukakeCrypt::new()),
+            CryptType::HashCrypt => Box::new(HashCrypt::new()),
+            CryptType::XorCrypt { key } => Box::new(XorCrypt::new(*key)),
+            CryptType::FlyingShineCrypt => Box::new(FlyingShineCrypt::new()),
+            CryptType::CxEncryption(schema) => Box::new(cx::CxEncryption::new(&schema, filename)?),
+        })
+    }
+}
+
+include_flate::flate!(static CRYPT_DATA: str from "src/scripts/kirikiri/archive/xp3/crypt.json" with zstd);
+
+lazy_static::lazy_static! {
+    static ref CRYPT_SCHEMA: BTreeMap<String, Schema> = {
+        serde_json::from_str(&CRYPT_DATA).expect("Failed to parse crypt.json")
+    };
+    static ref ALIAS_TABLE: HashMap<String, String> = {
+        let mut table = HashMap::new();
+        for (game, fulltitle) in get_supported_games_with_title() {
+            if let Some(title) = fulltitle {
+                let mut alias_count = 0usize;
+                for part in title.split("|") {
+                    let alias = part.trim();
+                    table.insert(alias.to_string(), game.to_string());
+                    alias_count += 1;
+                }
+                // also insert full title if there are multiple aliases
+                if alias_count > 1 {
+                    table.insert(title.to_string(), game.to_string());
+                }
+            }
+        }
+        table
+    };
+}
+
+/// Get the supported game titles for encrypted xp3 archives.
+pub fn get_supported_games() -> Vec<&'static str> {
+    CRYPT_SCHEMA.keys().map(|s| s.as_str()).collect()
+}
+
+/// Get the supported game titles for encrypted xp3 archives with their full titles.
+pub fn get_supported_games_with_title() -> Vec<(&'static str, Option<&'static str>)> {
+    CRYPT_SCHEMA
+        .iter()
+        .map(|(k, v)| (k.as_str(), v.title.as_deref()))
+        .collect()
+}
+
+pub fn query_crypt_schema(game: &str) -> Option<&'static Schema> {
+    CRYPT_SCHEMA.get(game).or_else(|| {
+        ALIAS_TABLE
+            .get(game)
+            .and_then(|real_game| CRYPT_SCHEMA.get(real_game))
+    })
+}
+
+#[derive(Debug)]
+pub struct NoCrypt {}
+
+impl NoCrypt {
+    pub fn new() -> Self {
+        Self {}
+    }
+}
+
+impl Crypt for NoCrypt {}
+
+macro_rules! seek_impl {
+    ($reader:ident<$t:ident>) => {
+        impl<$t: Read + Seek> Seek for $reader<$t> {
+            fn seek(&mut self, pos: SeekFrom) -> std::io::Result<u64> {
+                let new_pos: i64 = match pos {
+                    SeekFrom::Start(offset) => offset as i64,
+                    SeekFrom::End(offset) => self.seg_size as i64 + offset,
+                    SeekFrom::Current(offset) => self.pos as i64 + offset,
+                };
+                let offset = new_pos - self.pos as i64;
+                if offset != 0 {
+                    self.inner.seek(SeekFrom::Current(offset))?;
+                    self.pos = new_pos as u64;
+                }
+                Ok(self.pos)
+            }
+        }
+    };
+}
+
+macro_rules! seek_reader_impl {
+    ($reader:ident<$t:ident>) => {
+        #[derive(msg_tool_macro::MyDebug)]
+        struct $reader<$t: Read> {
+            #[skip_fmt]
+            inner: $t,
+            /// Start offset of the current xp3 entry.
+            seg_start: u64,
+            seg_size: u64,
+            pos: u64,
+        }
+        impl<$t: Read> $reader<$t> {
+            pub fn new(inner: $t, seg: &Segment) -> Self {
+                Self {
+                    inner,
+                    seg_start: seg.offset_in_file,
+                    seg_size: seg.original_size,
+                    pos: 0,
+                }
+            }
+        }
+        seek_impl!($reader<$t>);
+    };
+}
+
+macro_rules! seek_reader_key_impl {
+    ($reader:ident<$t:ident>, $key:ty) => {
+        #[derive(msg_tool_macro::MyDebug)]
+        #[allow(dead_code)]
+        struct $reader<$t: Read> {
+            #[skip_fmt]
+            inner: $t,
+            /// Start offset of the current xp3 entry.
+            seg_start: u64,
+            seg_size: u64,
+            pos: u64,
+            key: $key,
+        }
+        impl<$t: Read> $reader<$t> {
+            pub fn new(inner: $t, seg: &Segment, key: $key) -> Self {
+                Self {
+                    inner,
+                    seg_start: seg.offset_in_file,
+                    seg_size: seg.original_size,
+                    pos: 0,
+                    key,
+                }
+            }
+        }
+        seek_impl!($reader<$t>);
+    };
+}
+
+macro_rules! seek_crypt_base_impl {
+    ($crypt:ident, $reader:ident) => {
+        #[derive(Debug)]
+        pub struct $crypt {}
+        impl $crypt {
+            pub fn new() -> Self {
+                Self {}
+            }
+        }
+        impl Crypt for $crypt {
+            fn decrypt_supported(&self) -> bool {
+                true
+            }
+            fn decrypt_seek_supported(&self) -> bool {
+                true
+            }
+            fn decrypt<'a>(
+                &self,
+                _entry: &Xp3Entry,
+                cur_seg: &Segment,
+                stream: Box<dyn Read + 'a>,
+            ) -> Result<Box<dyn ReadDebug + 'a>> {
+                Ok(Box::new($reader::new(stream, cur_seg)))
+            }
+            fn decrypt_with_seek<'a>(
+                &self,
+                _entry: &Xp3Entry,
+                cur_seg: &Segment,
+                stream: Box<dyn ReadSeek + 'a>,
+            ) -> Result<Box<dyn ReadSeek + 'a>> {
+                Ok(Box::new($reader::new(stream, cur_seg)))
+            }
+        }
+    };
+}
+
+macro_rules! seek_crypt_impl {
+    ($crypt:ident, $reader:ident<$t:ident>) => {
+        seek_crypt_base_impl!($crypt, $reader);
+        seek_reader_impl!($reader<$t>);
+    };
+}
+
+seek_crypt_impl!(FateCrypt, FateCryptReader<T>);
+
+impl<R: Read> Read for FateCryptReader<R> {
+    fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
+        const XOR1_OFFSET: u64 = 0x13;
+        const XOR3_OFFSET: u64 = 0x2ea29;
+        let readed = self.inner.read(buf)?;
+        for (i, t) in (&mut buf[..readed]).iter_mut().enumerate() {
+            let tpos = self.seg_start + self.pos + i as u64;
+            *t ^= 0x36;
+            if tpos == XOR1_OFFSET {
+                *t ^= 0x1;
+            } else if tpos == XOR3_OFFSET {
+                *t ^= 0x3;
+            }
+        }
+        self.pos += readed as u64;
+        Ok(readed)
+    }
+}
+
+seek_crypt_impl!(MizukakeCrypt, MizukakeCryptReader<T>);
+
+impl<R: Read> Read for MizukakeCryptReader<R> {
+    fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
+        let readed = self.inner.read(buf)?;
+        for (i, t) in (&mut buf[..readed]).iter_mut().enumerate() {
+            let tpos = self.seg_start + self.pos + i as u64;
+            if tpos == 0x103 {
+                *t = (*t).wrapping_sub(1);
+            }
+            *t ^= 0xb6;
+            if tpos == 0x3F82 {
+                *t ^= 1;
+            }
+            if tpos == 0x83 {
+                *t ^= 3;
+            }
+        }
+        self.pos += readed as u64;
+        Ok(readed)
+    }
+}
+
+#[derive(Debug)]
+pub struct HashCrypt {}
+
+impl HashCrypt {
+    pub fn new() -> Self {
+        Self {}
+    }
+}
+
+impl Crypt for HashCrypt {
+    fn decrypt_supported(&self) -> bool {
+        true
+    }
+    fn decrypt_seek_supported(&self) -> bool {
+        true
+    }
+    fn decrypt<'a>(
+        &self,
+        entry: &Xp3Entry,
+        cur_seg: &Segment,
+        stream: Box<dyn Read + 'a>,
+    ) -> Result<Box<dyn ReadDebug + 'a>> {
+        Ok(Box::new(HashCryptReader::new(
+            stream,
+            cur_seg,
+            entry.file_hash as u8,
+        )))
+    }
+    fn decrypt_with_seek<'a>(
+        &self,
+        entry: &Xp3Entry,
+        cur_seg: &Segment,
+        stream: Box<dyn ReadSeek + 'a>,
+    ) -> Result<Box<dyn ReadSeek + 'a>> {
+        Ok(Box::new(HashCryptReader::new(
+            stream,
+            cur_seg,
+            entry.file_hash as u8,
+        )))
+    }
+}
+
+seek_reader_key_impl!(HashCryptReader<T>, u8);
+
+impl<R: Read> Read for HashCryptReader<R> {
+    fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
+        let readed = self.inner.read(buf)?;
+        for t in (&mut buf[..readed]).iter_mut() {
+            *t ^= self.key;
+        }
+        self.pos += readed as u64;
+        Ok(readed)
+    }
+}
+
+#[derive(Debug)]
+pub struct XorCrypt {
+    key: u8,
+}
+
+impl XorCrypt {
+    pub fn new(key: u8) -> Self {
+        Self { key }
+    }
+}
+
+impl Crypt for XorCrypt {
+    fn decrypt_supported(&self) -> bool {
+        true
+    }
+    fn decrypt_seek_supported(&self) -> bool {
+        true
+    }
+    fn decrypt<'a>(
+        &self,
+        _entry: &Xp3Entry,
+        cur_seg: &Segment,
+        stream: Box<dyn Read + 'a>,
+    ) -> Result<Box<dyn ReadDebug + 'a>> {
+        Ok(Box::new(XorCryptReader::new(stream, cur_seg, self.key)))
+    }
+    fn decrypt_with_seek<'a>(
+        &self,
+        _entry: &Xp3Entry,
+        cur_seg: &Segment,
+        stream: Box<dyn ReadSeek + 'a>,
+    ) -> Result<Box<dyn ReadSeek + 'a>> {
+        Ok(Box::new(XorCryptReader::new(stream, cur_seg, self.key)))
+    }
+}
+
+seek_reader_key_impl!(XorCryptReader<T>, u8);
+
+impl<R: Read> Read for XorCryptReader<R> {
+    fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
+        let readed = self.inner.read(buf)?;
+        for t in (&mut buf[..readed]).iter_mut() {
+            *t ^= self.key;
+        }
+        self.pos += readed as u64;
+        Ok(readed)
+    }
+}
+
+#[derive(Debug)]
+pub struct FlyingShineCrypt {}
+
+impl FlyingShineCrypt {
+    pub fn new() -> Self {
+        Self {}
+    }
+
+    fn adjust(&self, hash: u32) -> (u8, u32) {
+        let mut shift = hash & 0xFF;
+        if shift == 0 {
+            shift = 0xF;
+        }
+        let mut key = ((hash >> 8) & 0xFF) as u8;
+        if key == 0 {
+            key = 0xF0;
+        }
+        (key, shift)
+    }
+}
+
+impl Crypt for FlyingShineCrypt {
+    fn decrypt_supported(&self) -> bool {
+        true
+    }
+    fn decrypt_seek_supported(&self) -> bool {
+        true
+    }
+    fn decrypt<'a>(
+        &self,
+        entry: &Xp3Entry,
+        cur_seg: &Segment,
+        stream: Box<dyn Read + 'a>,
+    ) -> Result<Box<dyn ReadDebug + 'a>> {
+        Ok(Box::new(FlyingShineCryptReader::new(
+            stream,
+            cur_seg,
+            self.adjust(entry.file_hash),
+        )))
+    }
+    fn decrypt_with_seek<'a>(
+        &self,
+        entry: &Xp3Entry,
+        cur_seg: &Segment,
+        stream: Box<dyn ReadSeek + 'a>,
+    ) -> Result<Box<dyn ReadSeek + 'a>> {
+        Ok(Box::new(FlyingShineCryptReader::new(
+            stream,
+            cur_seg,
+            self.adjust(entry.file_hash),
+        )))
+    }
+}
+
+seek_reader_key_impl!(FlyingShineCryptReader<T>, (u8, u32));
+
+impl<R: Read> Read for FlyingShineCryptReader<R> {
+    fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
+        let (xor, shift) = self.key;
+        let readed = self.inner.read(buf)?;
+        for t in (&mut buf[..readed]).iter_mut() {
+            *t ^= xor;
+            *t = t.rotate_right(shift);
+        }
+        self.pos += readed as u64;
+        Ok(readed)
+    }
+}
+
+// extended in cx.rs
+seek_reader_key_impl!(CxEncryptionReader<T>, (u32, Arc<cx::CxEncryption>));
+
+#[test]
+fn test_deserialize_crypt() {
+    for (key, schema) in CRYPT_SCHEMA.iter() {
+        println!("Title: {}, Schema: {:?}", key, schema);
+    }
+}