[kgm] refactor: convert to enum for dispatch

um_cli/src/cmd/kgm.rs

@@ -4,6 +4,7 @@ use std::fs::File;
 use std::io::{Read, Seek, SeekFrom, Write};
 use std::path::PathBuf;
 use umc_kgm::header::Header;
+use umc_kgm::Decipher;
 
 /// Decrypt a KGM/VPR file (Kugou Music)
 #[derive(Args)]
@@ -23,7 +24,7 @@ impl ArgsKGM {
         let mut header = [0u8; 0x40];
         file_input.read_exact(&mut header)?;
         let kgm_header = Header::from_buffer(&mut header)?;
-        let decipher = kgm_header.make_decipher()?;
+        let decipher = Decipher::new(&kgm_header)?;
         file_input.seek(SeekFrom::Start(kgm_header.offset_to_data as u64))?;
 
         let mut offset = 0usize;

um_crypto/kgm/src/header.rs

@@ -1,6 +1,4 @@
-use crate::v2::DecipherV2;
-use crate::v3::DecipherV3;
-use crate::{Decipher, KugouError};
+use crate::KugouError;
 use byteorder::{ByteOrder, LE};
 
 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
@@ -48,25 +46,8 @@ impl Header {
         })
     }
 
-    pub fn make_decipher(&self) -> Result<Box<dyn Decipher>, KugouError> {
-        let slot_key: &[u8] = match self.key_slot {
-            1 => b"l,/'",
-            slot => Err(KugouError::UnsupportedKeySlot(slot))?,
-        };
-
-        let decipher: Box<dyn Decipher> = match self.crypto_version {
-            2 => Box::new(DecipherV2::new(self, slot_key)?),
-            3 => Box::new(DecipherV3::new(self, slot_key)?),
-            version => Err(KugouError::UnsupportedCipherVersion(version))?,
-        };
-
-        let mut test_data = self.decrypt_test_data;
-        decipher.decrypt(&mut test_data, 0);
-        if self.challenge_data != test_data {
-            Err(KugouError::SelfTestFailed)?;
-        }
-
-        Ok(decipher)
+    pub fn get_challenge_data(&self) -> [u8; 0x10] {
+        self.challenge_data
     }
 }
 

um_crypto/kgm/src/lib.rs

@@ -1,7 +1,10 @@
 pub mod header;
 pub mod v2;
-mod v3;
+pub mod v3;
 
+use crate::header::Header;
+use crate::v2::DecipherV2;
+use crate::v3::DecipherV3;
 use thiserror::Error;
 
 #[derive(Debug, Error)]
@@ -22,6 +25,37 @@ pub enum KugouError {
     SelfTestFailed,
 }
 
-pub trait Decipher {
-    fn decrypt(&self, buffer: &mut [u8], offset: usize);
+pub enum Decipher {
+    V2(DecipherV2),
+    V3(DecipherV3),
+}
+
+impl Decipher {
+    pub fn new(header: &Header) -> Result<Self, KugouError> {
+        let slot_key: &[u8] = match header.key_slot {
+            1 => b"l,/'",
+            slot => Err(KugouError::UnsupportedKeySlot(slot))?,
+        };
+
+        let decipher = match header.crypto_version {
+            2 => Decipher::V2(DecipherV2::new(header, slot_key)?),
+            3 => Decipher::V3(DecipherV3::new(header, slot_key)?),
+            version => Err(KugouError::UnsupportedCipherVersion(version))?,
+        };
+
+        let mut test_data = header.decrypt_test_data;
+        decipher.decrypt(&mut test_data, 0);
+        if test_data != header.get_challenge_data() {
+            Err(KugouError::SelfTestFailed)?;
+        }
+
+        Ok(decipher)
+    }
+
+    pub fn decrypt<T: AsMut<[u8]> + ?Sized>(&self, buffer: &mut T, offset: usize) {
+        match self {
+            Decipher::V2(decipher) => decipher.decrypt(buffer, offset),
+            Decipher::V3(decipher) => decipher.decrypt(buffer, offset),
+        }
+    }
 }

um_crypto/kgm/src/v2.rs

@@ -1,5 +1,5 @@
 use crate::header::Header;
-use crate::{Decipher, KugouError};
+use crate::KugouError;
 
 pub struct DecipherV2 {
     key: [u8; 4],
@@ -11,11 +11,9 @@ impl DecipherV2 {
         key.copy_from_slice(slot_key);
         Ok(Self { key })
     }
-}
 
-impl Decipher for DecipherV2 {
-    fn decrypt(&self, buffer: &mut [u8], offset: usize) {
-        for (datum, offset) in buffer.iter_mut().zip(offset..) {
+    pub fn decrypt<T: AsMut<[u8]> + ?Sized>(&self, buffer: &mut T, offset: usize) {
+        for (datum, offset) in buffer.as_mut().iter_mut().zip(offset..) {
             *datum ^= self.key[offset % self.key.len()];
         }
     }

um_crypto/kgm/src/v3.rs

@@ -1,5 +1,5 @@
 use crate::header::Header;
-use crate::{Decipher, KugouError};
+use crate::KugouError;
 
 pub struct DecipherV3 {
     slot_key: [u8; 16],
@@ -30,11 +30,9 @@ impl DecipherV3 {
         let offset_key = (offset as u32).to_ne_bytes();
         offset_key[0] ^ offset_key[1] ^ offset_key[2] ^ offset_key[3]
     }
-}
 
-impl Decipher for DecipherV3 {
-    fn decrypt(&self, buffer: &mut [u8], offset: usize) {
-        for (datum, offset) in buffer.iter_mut().zip(offset..) {
+    pub fn decrypt<T: AsMut<[u8]> + ?Sized>(&self, buffer: &mut T, offset: usize) {
+        for (datum, offset) in buffer.as_mut().iter_mut().zip(offset..) {
             let offset_key = Self::offset_key(offset);
 
             let file_key = self.file_key[offset % self.file_key.len()];