Merge pull request #15 from parakeet-rs/refactor/use-src-filter-dst

Use stream API (std::io::read/seek/write)

.gitignore

@@ -22,10 +22,13 @@ venv/
 
 # Decryption key files
 *.bin
-/*.ogg
+*.ogg
 *.flac
 *.mp3
 *.dff
 
+!**/__fixture__/*
+!**/sample/*
+
 keys/
 local/

Cargo.toml

@@ -8,6 +8,12 @@ name = "parakeet_cli"
 
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
 
+[profile.release]
+# https://github.com/johnthagen/min-sized-rust
+lto = true
+strip = true
+codegen-units = 1
+
 [dependencies]
 argh = "0.1.9"
 byteorder = "1.4.3"

src/bin/cli/cli_handle_kugou.rs

@@ -1,9 +1,8 @@
-use std::{fs::File, process};
+use std::{fs::File, io::copy, process};
 
 use argh::FromArgs;
-use parakeet_crypto::{
- interfaces::Decryptor,
- kugou::{self, kgm_crypto::KGMCryptoConfig, kgm_header::KGMHeader},
+use parakeet_crypto::filters::{
+ file_header::KGMHeader, KGMCryptoConfig, KugouDecryptReader, KugouEncryptReader,
 };
 
 use super::{
@@ -47,19 +46,18 @@ pub fn cli_handle_kugou(args: KugouOptions) {
  let mut config = KGMCryptoConfig::default();
 
  if let Some(table) = args.v4_file_key_expansion_table {
- config.v4_file_key_expand_table = table.content;
+ config.v4_file_key_expand_table = table.content.to_vec();
  log.info("(v4) file key expansion table accepted.");
  }
 
  if let Some(table) = args.v4_slot_key_expansion_table {
- config.v4_slot_key_expand_table = table.content;
+ config.v4_slot_key_expand_table = table.content.to_vec();
  log.info("(v4) slot key expansion table accepted.");
  }
 
  // Configure key slots
- config.slot_keys.insert(1, args.slot_key_1.content);
+ config.slot_keys.insert(1, args.slot_key_1.content.to_vec());
 
- let mut kgm = kugou::kgm_decryptor::KGM::new(&config);
  let mut input_file = File::open(args.input_file.path).unwrap();
  let mut output_file = File::create(args.output_file.path).unwrap();
 
@@ -70,14 +68,16 @@ pub fn cli_handle_kugou(args: KugouOptions) {
  };
 
  if let Some(encrypt_header) = args.encrypt_header {
- let mut header = KGMHeader::from_bytes(&encrypt_header.content).unwrap_or_else(|err| {
+ let header = KGMHeader::from_bytes(&encrypt_header.content).unwrap_or_else(|err| {
  log.error(&format!("Could not parse header: {:?}", err));
  process::exit(1)
  });
 
- kgm.encrypt(&mut header, &mut input_file, &mut output_file)
+ let mut kgm_reader = KugouEncryptReader::new(&config, &header, &mut input_file).unwrap();
+ copy(&mut kgm_reader, &mut output_file)
  } else {
- kgm.decrypt(&mut input_file, &mut output_file)
+ let mut kgm_reader = KugouDecryptReader::new(&config, &mut input_file).unwrap();
+ copy(&mut kgm_reader, &mut output_file)
  }
  .unwrap_or_else(|err| {
  log.error(&format!("{operation} failed: {err}"));

src/bin/cli/cli_handle_qmc1.rs

@@ -1,7 +1,7 @@
 use std::{fs::File, process};
 
 use argh::FromArgs;
-use parakeet_crypto::{interfaces::Decryptor, QmcV1};
+use parakeet_crypto::filters::{QMC1Static, QMC1StaticReader};
 
 use super::{
  logger::CliLogger,
@@ -27,27 +27,30 @@ pub struct QMC1Options {
 
 pub fn cli_handle_qmc1(args: QMC1Options) {
  let log = CliLogger::new("QMC1");
- let mut qmc1_static = match QmcV1::new_static(&args.static_key.content) {
- None => {
- log.error("key rejected, invalid length?");
+
+ let qmc1 = match args.static_key.content.len() {
+ 128 => QMC1Static::new(args.static_key.content[..].try_into().unwrap()),
+ 256 => QMC1Static::new_key256(args.static_key.content[..].try_into().unwrap()),
+ _ => {
+ log.error("key rejected -- invalid length");
  return;
  }
- Some(x) => x,
  };
 
  log.info(&format!(
  "Static key accepted (key{})",
  args.static_key.content.len()
  ));
 
- qmc1_static
- .decrypt(
- &mut File::open(args.input_file.path).unwrap(),
- &mut File::create(args.output_file.path).unwrap(),
- )
- .unwrap_or_else(|err| {
- log.error(&format!("Decryption failed: {err}"));
- process::exit(1);
- });
+ let mut src = File::open(args.input_file.path).unwrap();
+ let mut dst = File::create(args.output_file.path).unwrap();
+
+ let mut qmc1_reader = QMC1StaticReader::new(qmc1, &mut src);
+
+ std::io::copy(&mut qmc1_reader, &mut dst).unwrap_or_else(|err| {
+ log.error(&format!("transform failed: {err}"));
+ process::exit(1)
+ });
+
  log.info("Decryption OK.");
 }

src/bin/cli/cli_handle_qmc2.rs

@@ -1,10 +1,7 @@
 use std::{fs::File, process};
 
 use argh::FromArgs;
-use parakeet_crypto::{
- interfaces::Decryptor,
- qmc2::{QMCTailParser, QMC2},
-};
+use parakeet_crypto::filters::{QMC2Map, QMC2Reader, QMCFooterParser, QMC2RC4};
 
 use crate::cli::logger::CliLogger;
 
@@ -38,7 +35,7 @@ pub struct QMC2Options {
 pub fn cli_handle_qmc2(args: QMC2Options) {
  let log = CliLogger::new("QMC2");
 
- let mut parser = QMCTailParser::new(args.seed);
+ let mut parser = QMCFooterParser::new(args.seed);
 
  if let Some(key1) = args.key1 {
  parser.set_key_stage1(key1.content);
@@ -55,15 +52,22 @@ pub fn cli_handle_qmc2(args: QMC2Options) {
  process::exit(1);
  }
 
- let mut qmc2_map = QMC2::new(parser);
- qmc2_map
- .decrypt(
- &mut File::open(args.input_file.path).unwrap(),
- &mut File::create(args.output_file.path).unwrap(),
- )
+ let mut src = File::open(args.input_file.path).unwrap();
+ let mut dst = File::create(args.output_file.path).unwrap();
+
+ let mut qmc2_map = QMC2Map::new_default();
+ let mut qmc2_rc4 = QMC2RC4::new_default();
+
+ let mut qmc2_reader = QMC2Reader::new(&mut parser, &mut qmc2_map, &mut qmc2_rc4, &mut src)
  .unwrap_or_else(|err| {
- log.error(&format!("Decryption failed: {err}"));
- process::exit(1);
+ log.error(&format!("init qmc2 reader failed: {err}"));
+ process::exit(1)
  });
+
+ std::io::copy(&mut qmc2_reader, &mut dst).unwrap_or_else(|err| {
+ log.error(&format!("transform failed: {err}"));
+ process::exit(1)
+ });
+
  log.info("Decryption OK.");
 }

src/bin/cli/cli_handle_xmly.rs

@@ -1,20 +1,29 @@
 use std::{fs::File, process};
 
 use argh::FromArgs;
-use parakeet_crypto::{interfaces::Decryptor, ximalaya};
+use parakeet_crypto::filters::{XimalayaCrypto, XimalayaReader, SCRAMBLE_HEADER_LEN};
 
 use super::{
  logger::CliLogger,
  utils::{CliBinaryContent, CliFilePath},
 };
 
 /// Handle x2m/x3m encryption/decryption.
-#[derive(Debug, Eq, PartialEq, FromArgs)]
+#[derive(Debug, PartialEq, FromArgs)]
 #[argh(subcommand, name = "xmly")]
 pub struct XimalayaOptions {
  /// scramble table (u16 x 1024 items, little-endian)
+ /// when size mismatch, revert to generator.
  #[argh(option)]
- scramble_table: CliBinaryContent,
+ scramble_table: Option<CliBinaryContent>,
+
+ /// initial value (scramble-table gen)
+ #[argh(option)]
+ mul_init: Option<f64>,
+
+ /// step value (scramble-table gen)
+ #[argh(option)]
+ mul_step: Option<f64>,
 
  /// X2M/X3M key.
  /// 4-bytes = X2M
@@ -39,43 +48,51 @@ pub struct XimalayaOptions {
 pub fn cli_handle_xmly(args: XimalayaOptions) {
  let log = CliLogger::new("XMLY");
 
- if args.scramble_table.content.len() != 2048 {
- log.error(&format!(
- "expecting scramble-table to have a size of 2048, got {} instead.",
- args.scramble_table.content.len()
- ));
- process::exit(1);
- }
+ let xmly_crypto = if let Some(scramble_table_arg) = args.scramble_table {
+ let scramble_table_bin = scramble_table_arg.content;
+ if scramble_table_bin.len() != SCRAMBLE_HEADER_LEN * 2 {
+ log.error(&format!(
+ "expecting scramble-table to have a size of {} but got {} instead.",
+ SCRAMBLE_HEADER_LEN * 2,
+ scramble_table_bin.len()
+ ));
+ process::exit(1);
+ }
 
- let mut scramble_table = [0usize; 1024];
- for (i, item) in scramble_table.iter_mut().enumerate() {
- let mut buffer = [0u8; 2];
- buffer.copy_from_slice(&args.scramble_table.content[i * 2..i * 2 + 2]);
- *item = u16::from_le_bytes(buffer) as usize;
- }
+ log.info("using scramble-table for xmly.");
 
- let operation = if args.encrypt {
- "Encryption"
+ let mut scramble_table = [0u16; SCRAMBLE_HEADER_LEN];
+ let mut buffer = [0u8; 2];
+ for (i, item) in scramble_table.iter_mut().enumerate() {
+ buffer.copy_from_slice(&scramble_table_bin[i * 2..i * 2 + 2]);
+ *item = u16::from_le_bytes(buffer);
+ }
+ XimalayaCrypto::new(&args.key.content, &scramble_table)
+ } else if args.mul_init.is_some() && args.mul_step.is_some() {
+ XimalayaCrypto::new_from_param(
+ &args.key.content,
+ args.mul_init.unwrap(),
+ args.mul_step.unwrap(),
+ )
  } else {
- "Decryption"
+ log.error("you should specify (--scramble-table) or (--mul-init, --mul-step).");
+ process::exit(1);
+ };
+
+ let operation = match args.encrypt {
+ true => "Encryption",
+ false => "Decryption",
  };
 
- let mut xmly =
- ximalaya::new_from_key(&args.key.content[..], &scramble_table).unwrap_or_else(|err| {
- log.error(&format!(
- "Create encryptor/decryptor using key failed: {err}"
- ));
- process::exit(1)
- });
  let mut input_file = File::open(args.input_file.path).unwrap();
  let mut output_file = File::create(args.output_file.path).unwrap();
 
- if args.encrypt {
- xmly.encrypt(&mut input_file, &mut output_file)
- } else {
-  xmly.decrypt(&mut input_file, &mut output_file)
- }
- .unwrap_or_else(|err| {
+ let mut xmly_reader = match args.encrypt {
+ true => XimalayaReader::new_encrypt(xmly_crypto, &mut input_file),
+  false => XimalayaReader::new_decrypt(xmly_crypto, &mut input_file),
+ };
+
+ std::io::copy(&mut xmly_reader, &mut output_file).unwrap_or_else(|err| {
  log.error(&format!("{operation} failed: {err}"));
  process::exit(1)
  });

src/bin/cli/commands.rs

@@ -6,13 +6,13 @@ use super::{
 };
 
 /// Test CLI tool for parakeet_crypto.
-#[derive(FromArgs, Eq, PartialEq, Debug)]
+#[derive(FromArgs, PartialEq, Debug)]
 pub struct ParakeetCLIArgRoot {
  #[argh(subcommand)]
  pub command: ParakeetCryptoName,
 }
 
-#[derive(FromArgs, Eq, PartialEq, Debug)]
+#[derive(FromArgs, PartialEq, Debug)]
 #[argh(subcommand)]
 pub enum ParakeetCryptoName {
  ModuleQMC1(QMC1Options),

src/bin/cli/utils.rs

@@ -11,7 +11,7 @@ pub fn read_key_from_parameter(value: &str) -> Option<Box<[u8]>> {
  let content = Base64.decode(value).unwrap();
  Some(content.into())
  } else if let Some(value) = value.strip_prefix("hex:") {
- let content = hex::decode(value).unwrap();
+ let content = hex::decode(value.replace(' ', "")).unwrap();
  Some(content.into())
  } else if let Some(value) = value.strip_prefix("raw:") {
  let content = value.as_bytes();

src/filters/kugou/base.rs

@@ -0,0 +1,15 @@
+use std::collections::HashMap;
+
+#[derive(Debug, Default, Clone)]
+pub struct KGMCryptoConfig {
+ pub slot_keys: HashMap<u32, Vec<u8>>,
+ pub v4_slot_key_expand_table: Vec<u8>,
+ pub v4_file_key_expand_table: Vec<u8>,
+}
+
+pub trait KGMCrypto {
+ fn configure(&mut self, config: &KGMCryptoConfig, slot_key: &[u8], file_key: &[u8]);
+
+ fn decrypt(&mut self, offset: u64, buffer: &mut [u8]);
+ fn encrypt(&mut self, offset: u64, buffer: &mut [u8]);
+}

src/filters/kugou/crypto/mod.rs

@@ -0,0 +1,9 @@
+mod utils;
+
+mod type2;
+mod type3;
+mod type4;
+
+pub(super) use type2::KGMCryptoType2;
+pub(super) use type3::KGMCryptoType3;
+pub(super) use type4::KGMCryptoType4;