From e6fcb07ed28e0e1d48074138aa508aac3bbf1903 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E9=B2=81=E6=A0=91=E4=BA=BA?= <lu.shuren@um-react.app>
Date: Fri, 6 Sep 2024 00:52:23 +0100
Subject: [PATCH] feat: add qmc2 cli

---
 um_cli/Cargo.toml      |  1 +
 um_cli/src/cmd/mod.rs  |  3 ++
 um_cli/src/cmd/qmc2.rs | 91 ++++++++++++++++++++++++++++++++++++++++++
 um_cli/src/main.rs     |  7 ++--
 4 files changed, 99 insertions(+), 3 deletions(-)
 create mode 100644 um_cli/src/cmd/qmc2.rs

diff --git a/um_cli/Cargo.toml b/um_cli/Cargo.toml
index 9951f2a..2bfb996 100644
--- a/um_cli/Cargo.toml
+++ b/um_cli/Cargo.toml
@@ -8,3 +8,4 @@ anyhow = "1.0.86"
 clap = { version = "4.5.17", features = ["derive"] }
 umc_kuwo = { path = "../um_crypto/kuwo" }
 umc_qmc = { path = "../um_crypto/qmc" }
+umc_utils = { path = "../um_crypto/utils" }
diff --git a/um_cli/src/cmd/mod.rs b/um_cli/src/cmd/mod.rs
index 7ae8d2b..430a8d5 100644
--- a/um_cli/src/cmd/mod.rs
+++ b/um_cli/src/cmd/mod.rs
@@ -1,9 +1,12 @@
 use clap::Subcommand;
 
 pub mod qmc1;
+pub mod qmc2;
 
 #[derive(Subcommand)]
 pub enum Commands {
     #[command(name = "qmc1")]
     QMCv1(qmc1::ArgsQMCv1),
+    #[command(name = "qmc2")]
+    QMCv2(qmc2::ArgsQMCv2),
 }
diff --git a/um_cli/src/cmd/qmc2.rs b/um_cli/src/cmd/qmc2.rs
new file mode 100644
index 0000000..05942b2
--- /dev/null
+++ b/um_cli/src/cmd/qmc2.rs
@@ -0,0 +1,91 @@
+use crate::Cli;
+use anyhow::{bail, Result};
+use clap::Args;
+use std::fs::File;
+use std::io::{BufReader, BufWriter, Read, Seek, SeekFrom, Write};
+use std::path::PathBuf;
+use umc_qmc::{footer, QMCv2Cipher};
+use umc_utils::base64;
+
+/// Decrypt a QMCv1 file
+#[derive(Args)]
+pub struct ArgsQMCv2 {
+    /// Path to output file, e.g. /export/Music/song.flac
+    #[arg(short, long)]
+    output: Option<PathBuf>,
+
+    /// Path to input file, e.g. /export/Music/song.qmcflac
+    #[arg(name = "input")]
+    input: PathBuf,
+
+    /// Override EKey for this file.
+    /// Prefix with "decrypted:" to use base64 encoded raw key.
+    #[arg(short = 'K', long = "ekey")]
+    ekey: Option<String>,
+
+    /// Print info about this file, and do not perform decryption.
+    #[arg(short = 'I', long, action=clap::ArgAction::SetTrue, default_value_t=false)]
+    info_only: bool,
+}
+
+impl ArgsQMCv2 {
+    pub fn run(&self, cli: &Cli) -> Result<i32> {
+        let mut file_input = File::open(&self.input)?;
+        let mut footer_detection_buffer = vec![0u8; footer::INITIAL_DETECTION_LEN];
+        file_input.seek(SeekFrom::End(-(footer::INITIAL_DETECTION_LEN as i64)))?;
+        file_input.read_exact(&mut footer_detection_buffer)?;
+        let input_size = file_input.stream_position()?;
+        file_input.seek(SeekFrom::Start(0))?;
+
+        let (footer_len, ekey) = match footer::from_byte_slice(&footer_detection_buffer) {
+            Ok(Some(metadata)) => {
+                if self.info_only || cli.verbose {
+                    println!("metadata: {:?}", metadata);
+                }
+                (metadata.size, metadata.ekey.or_else(|| self.ekey.clone()))
+            }
+            Ok(None) => {
+                eprintln!("could not find any qmc metadata.");
+                (0usize, self.ekey.clone())
+            }
+            Err(err) => {
+                eprintln!("failed to parse qmc metadata: {}", err);
+                (0usize, self.ekey.clone())
+            }
+        };
+
+        if self.info_only {
+            return Ok(0);
+        }
+
+        let key = match ekey {
+            None => bail!("--ekey is required when embedded ekey is not present."),
+            Some(ekey) => match ekey.strip_suffix("decrypted:") {
+                Some(decrypted) => base64::decode(decrypted)?.into_boxed_slice(),
+                None => umc_qmc::ekey::decrypt(ekey)?,
+            },
+        };
+        let cipher = QMCv2Cipher::new(key)?;
+
+        let mut file_output = match &self.output {
+            None => bail!("--output is required"),
+            Some(output) => BufWriter::new(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;
+        }
+
+        Ok(0)
+    }
+}
diff --git a/um_cli/src/main.rs b/um_cli/src/main.rs
index 5f342f7..517e77c 100644
--- a/um_cli/src/main.rs
+++ b/um_cli/src/main.rs
@@ -16,8 +16,8 @@ pub struct Cli {
     command: Option<Commands>,
 
     /// Be more verbose about what is going on.
-    #[clap(long, short, action=clap::ArgAction::SetTrue)]
-    verbose: Option<bool>,
+    #[clap(long, short, action=clap::ArgAction::SetTrue, default_value_t=false)]
+    verbose: bool,
 
     /// Preferred buffer size when reading file, in bytes.
     /// Default to 4MiB.
@@ -28,6 +28,7 @@ pub struct Cli {
 fn run_command(cli: &Cli) -> Result<i32> {
     match &cli.command {
         Some(Commands::QMCv1(cmd)) => cmd.run(&cli),
+        Some(Commands::QMCv2(cmd)) => cmd.run(&cli),
         None => {
             // https://github.com/clap-rs/clap/issues/3857#issuecomment-1161796261
             todo!("implement a sensible default command, similar to um/cli");
@@ -43,7 +44,7 @@ fn main() {
         -1
     });
     let duration = start.elapsed();
-    if let Some(true) = cli.verbose {
+    if cli.verbose {
         eprintln!("time: {:?}", duration);
     };
     exit(code);