feat: add kgm v5 (kgg) support.

um_crypto/kgm/src/header.rs

@@ -1,39 +1,69 @@
 use crate::KugouError;
-use byteorder::{ByteOrder, LE};
+use byteorder::{ReadBytesExt, LE};
+use std::io::{BufReader, Read};
 
-#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+#[derive(Debug, Clone, PartialEq, Eq)]
 pub struct Header {
     pub magic: [u8; 0x10],
     pub offset_to_data: usize,
     pub crypto_version: u32,
-    pub key_slot: u32,
+    pub key_slot: i32,
     pub decrypt_test_data: [u8; 0x10],
     pub file_key: [u8; 0x10],
 
     challenge_data: [u8; 0x10],
+    pub audio_hash: String,
 }
 
-impl Header {
-    pub fn from_buffer<T>(buffer: T) -> Result<Self, KugouError>
-    where
-        T: AsRef<[u8]>,
-    {
-        let buffer = buffer.as_ref();
-        if buffer.len() < 0x3c {
-            Err(KugouError::HeaderTooSmall(0x3c))?;
-        }
+pub trait HeaderReaderHelper: Read {
+    fn read_u32_le(&mut self) -> Result<u32, KugouError> {
+        self.read_u32::<LE>()
+            .map_err(KugouError::HeaderParseIOError)
+    }
 
+    fn read_i32_le(&mut self) -> Result<i32, KugouError> {
+        self.read_i32::<LE>()
+            .map_err(KugouError::HeaderParseIOError)
+    }
+
+    fn read_buff<T: AsMut<[u8]>>(&mut self, buffer: &mut T) -> Result<(), KugouError> {
+        self.read_exact(buffer.as_mut())
+            .map_err(KugouError::HeaderParseIOError)
+    }
+}
+impl<R: Read + ?Sized> HeaderReaderHelper for R {}
+
+impl Header {
+    pub fn from_reader<T>(reader: &mut T) -> Result<Self, KugouError>
+    where
+        T: Read,
+    {
         let mut magic = [0u8; 0x10];
-        magic.copy_from_slice(&buffer[..0x10]);
+        reader.read_buff(&mut magic)?;
         let challenge_data = get_challenge_data(&magic)?;
 
-        let offset_to_data = LE::read_u32(&buffer[0x10..0x14]) as usize;
-        let crypto_version = LE::read_u32(&buffer[0x14..0x18]);
-        let key_slot = LE::read_u32(&buffer[0x18..0x1C]);
         let mut decrypt_test_data = [0u8; 0x10];
-        decrypt_test_data.copy_from_slice(&buffer[0x1c..0x2c]);
         let mut file_key = [0u8; 0x10];
-        file_key.copy_from_slice(&buffer[0x2c..0x3c]);
+        let mut audio_hash = "".to_string();
+
+        let offset_to_data = reader.read_u32_le()? as usize;
+        let crypto_version = reader.read_u32_le()?;
+        let key_slot = reader.read_i32_le()?;
+        reader.read_buff(&mut decrypt_test_data)?;
+        reader.read_buff(&mut file_key)?;
+
+        if crypto_version == 5 {
+            let mut unused_padding = [0u8; 0x08];
+            reader.read_buff(&mut unused_padding)?; // seek 8 bytes
+            let audio_hash_size = reader.read_u32_le()? as usize;
+            if audio_hash_size != 0x20 {
+                Err(KugouError::HeaderInvalidAudioHash(audio_hash_size))?;
+            }
+
+            let mut audio_hash_bytes = vec![0u8; audio_hash_size];
+            reader.read_buff(&mut audio_hash_bytes)?;
+            audio_hash = String::from_utf8_lossy(&audio_hash_bytes).to_string();
+        }
 
         Ok(Self {
             magic,
@@ -43,9 +73,23 @@ impl Header {
             decrypt_test_data,
             file_key,
             challenge_data,
+            audio_hash,
         })
     }
 
+    pub fn from_buffer<T>(buffer: T) -> Result<Self, KugouError>
+    where
+        T: AsRef<[u8]>,
+    {
+        let buffer = buffer.as_ref();
+        if buffer.len() < 0x40 {
+            Err(KugouError::HeaderTooSmall(0x40))?;
+        }
+
+        let mut reader = BufReader::new(buffer);
+        Self::from_reader(&mut reader)
+    }
+
     pub fn get_challenge_data(&self) -> [u8; 0x10] {
         self.challenge_data
     }
@@ -74,3 +118,50 @@ pub const VPR_HEADER: [u8; 16] = [
 pub const VPR_TEST_DATA: [u8; 16] = [
     0x1D, 0x5A, 0x05, 0x34, 0x0C, 0x41, 0x8D, 0x42, 0x9C, 0x83, 0x92, 0x6C, 0xAE, 0x16, 0xFE, 0x56,
 ];
+
+#[cfg(test)]
+mod tests {
+    use crate::header::Header;
+    use crate::KugouError;
+
+    #[test]
+    fn parse_header_error_too_small() {
+        assert!(matches!(
+            Header::from_buffer(b"invalid file"),
+            Err(KugouError::HeaderTooSmall(_))
+        ));
+    }
+
+    #[test]
+    fn parse_header_error_file_magic() {
+        assert!(matches!(
+            Header::from_buffer(include_bytes!("__fixtures__/kgm_invalid_magic.bin")),
+            Err(KugouError::NotKGMFile)
+        ));
+    }
+
+    #[test]
+    fn parse_header_v2() -> Result<(), KugouError> {
+        let hdr = Header::from_buffer(include_bytes!("__fixtures__/kgm_v2_hdr.bin"))?;
+        assert_eq!(hdr.key_slot, 1);
+        assert_eq!(hdr.crypto_version, 2);
+        Ok(())
+    }
+
+    #[test]
+    fn parse_header_v3() -> Result<(), KugouError> {
+        let hdr = Header::from_buffer(include_bytes!("__fixtures__/kgm_v3_hdr.bin"))?;
+        assert_eq!(hdr.key_slot, 1);
+        assert_eq!(hdr.crypto_version, 3);
+        Ok(())
+    }
+
+    #[test]
+    fn parse_header_v5() -> Result<(), KugouError> {
+        let hdr = Header::from_buffer(include_bytes!("__fixtures__/kgm_v5_hdr.bin"))?;
+        assert_eq!(hdr.key_slot, -1);
+        assert_eq!(hdr.crypto_version, 5);
+        assert_eq!(hdr.audio_hash, "81a26217da847692e7688e0a5ebe9da1");
+        Ok(())
+    }
+}