feat: add kugou db decryption logic

um_cli/src/buffered_decrypt.rs

@@ -0,0 +1,28 @@
+use std::io::{Read, Write};
+
+pub fn buffered_decrypt<T, R, W>(
+    f_in: &mut R,
+    f_out: &mut W,
+    buffer_size: usize,
+    decipher: T,
+) -> anyhow::Result<usize>
+where
+    R: Read,
+    W: Write,
+    T: Fn(&mut [u8], usize),
+{
+    let mut offset = 0usize;
+    let mut buffer = vec![0u8; buffer_size].into_boxed_slice();
+    while let Ok(n) = f_in.read(&mut buffer) {
+        if n == 0 {
+            break;
+        }
+
+        let chunk = &mut buffer[..n];
+        decipher(chunk, offset);
+        f_out.write_all(chunk)?;
+        offset += n;
+    }
+
+    Ok(offset)
+}

um_cli/src/cmd/kgm.rs

@@ -1,3 +1,4 @@
+use crate::buffered_decrypt::buffered_decrypt;
 use crate::Cli;
 use clap::Args;
 use std::fs::File;
@@ -16,10 +17,34 @@ pub struct ArgsKGM {
     /// Path to input file, e.g. /export/Music/song.kgm
     #[arg(name = "input")]
     input: PathBuf,
+
+    /// File mode, one of "kgm" or "db", default to "kgm"
+    #[arg(short, long, default_value = "kgm")]
+    file_mode: String,
 }
 
 impl ArgsKGM {
     pub fn run(&self, cli: &Cli) -> anyhow::Result<i32> {
+        match self.file_mode.as_str() {
+            "kgm" => self.decrypt_kgm_file(cli),
+            "db" => self.decrypt_db_file(),
+            _ => anyhow::bail!("Invalid file mode: {}", self.file_mode),
+        }
+    }
+
+    fn decrypt_db_file(&self) -> anyhow::Result<i32> {
+        let mut file_input = File::open(&self.input)?;
+
+        let mut buffer = Vec::new();
+        file_input.read_to_end(&mut buffer)?;
+        umc_kgm::decrypt_db(&mut buffer)?;
+        let mut file_output = File::create(&self.output)?;
+        file_output.write_all(&buffer)?;
+
+        Ok(0)
+    }
+
+    fn decrypt_kgm_file(&self, cli: &Cli) -> anyhow::Result<i32> {
         let mut file_input = File::open(&self.input)?;
         let mut header = [0u8; 0x40];
         file_input.read_exact(&mut header)?;
@@ -27,18 +52,15 @@ impl ArgsKGM {
         let decipher = Decipher::new(&kgm_header)?;
         file_input.seek(SeekFrom::Start(kgm_header.offset_to_data as u64))?;
 
-        let mut offset = 0usize;
-        let mut buffer = vec![0u8; cli.buffer_size].into_boxed_slice();
         let mut file_output = File::create(&self.output)?;
-        while let Ok(n) = file_input.read(&mut buffer) {
-            if n == 0 {
-                break;
-            }
-
-            decipher.decrypt(&mut buffer[..n], offset);
-            file_output.write_all(&buffer[..n])?;
-            offset += n;
-        }
+        buffered_decrypt(
+            &mut file_input,
+            &mut file_output,
+            cli.buffer_size,
+            |buffer, offset| {
+                decipher.decrypt(buffer, offset);
+            },
+        )?;
 
         Ok(0)
     }

um_cli/src/cmd/ncm.rs

@@ -1,7 +1,8 @@
+use crate::buffered_decrypt::buffered_decrypt;
 use crate::Cli;
 use clap::Args;
 use std::fs::File;
-use std::io::{Read, Seek, SeekFrom, Write};
+use std::io::{Seek, SeekFrom, Write};
 use std::path::PathBuf;
 use umc_ncm::header::NCMFile;
 
@@ -71,18 +72,14 @@ impl ArgsNCM {
             file_input.seek(SeekFrom::Start(ncm.audio_data_offset as u64))?;
             let mut file_output = File::create(output_path)?;
 
-            let mut offset = 0usize;
-            let mut buffer = vec![0u8; cli.buffer_size].into_boxed_slice();
-
-            while let Ok(n) = file_input.read(&mut buffer) {
-                if n == 0 {
-                    break;
-                }
-
-                ncm.decrypt(&mut buffer[..n], offset);
-                file_output.write_all(&buffer[..n])?;
-                offset += n;
-            }
+            buffered_decrypt(
+                &mut file_input,
+                &mut file_output,
+                cli.buffer_size,
+                |buffer, offset| {
+                    ncm.decrypt(buffer, offset);
+                },
+            )?;
         }
 
         Ok(0)

um_cli/src/cmd/qmc1.rs

@@ -1,8 +1,8 @@
+use crate::buffered_decrypt::buffered_decrypt;
 use crate::Cli;
 use anyhow::Result;
 use clap::Args;
 use std::fs::File;
-use std::io::{Read, Write};
 use std::path::PathBuf;
 
 /// Decrypt a QMCv1 file (QQMusic)
@@ -22,17 +22,12 @@ impl ArgsQMCv1 {
         let mut file_input = File::open(&self.input)?;
         let mut file_output = File::create(&self.output)?;
 
-        let mut offset = 0usize;
-        let mut buffer = vec![0u8; cli.buffer_size].into_boxed_slice();
-        while let Ok(n) = file_input.read(&mut buffer) {
-            if n == 0 {
-                break;
-            }
-
-            umc_qmc::v1::decrypt(&mut buffer[..n], offset);
-            file_output.write_all(&buffer[..n])?;
-            offset += n;
-        }
+        buffered_decrypt(
+            &mut file_input,
+            &mut file_output,
+            cli.buffer_size,
+            umc_qmc::v1::decrypt,
+        )?;
 
         Ok(0)
     }

um_cli/src/cmd/qmc2.rs

@@ -1,9 +1,10 @@
+use crate::buffered_decrypt::buffered_decrypt;
 use crate::Cli;
 use anyhow::{bail, Result};
 use clap::Args;
 use std::fs;
 use std::fs::File;
-use std::io::{BufReader, BufWriter, Read, Seek, SeekFrom, Write};
+use std::io::{BufReader, Read, Seek, SeekFrom};
 use std::path::PathBuf;
 use umc_qmc::{footer, QMCv2Cipher};
 use umc_utils::base64;
@@ -97,22 +98,19 @@ impl ArgsQMCv2 {
 
         let mut file_output = match &self.output {
             None => bail!("--output is required"),
-            Some(output) => BufWriter::new(File::create(output)?),
+            Some(output) => File::create(output)?,
         };
 
-        let mut buffer = vec![0u8; cli.buffer_size];
         let reader = BufReader::with_capacity(cli.buffer_size, file_input);
         let mut reader = reader.take(input_size - footer_len as u64);
-        let mut offset = 0usize;
-        while let Ok(n) = reader.read(&mut buffer) {
-            if n == 0 {
-                break;
-            }
-
-            cipher.decrypt(&mut buffer[..n], offset);
-            file_output.write_all(&buffer[..n])?;
-            offset += n;
-        }
+        buffered_decrypt(
+            &mut reader,
+            &mut file_output,
+            cli.buffer_size,
+            |buffer, offset| {
+                cipher.decrypt(buffer, offset);
+            },
+        )?;
 
         Ok(0)
     }

um_cli/src/main.rs

@@ -4,6 +4,7 @@ use clap::Parser;
 use std::process::exit;
 use std::time::Instant;
 
+mod buffered_decrypt;
 mod cmd;
 
 /// um_cli (rust ver.)