From bb5f88a775196291ce14f827ee74ad920e0f9c7e Mon Sep 17 00:00:00 2001 From: Hatter Jiang Date: Sat, 24 Aug 2024 01:26:23 +0800 Subject: [PATCH] feat: update osssendfile-rs --- osssendfile-rs/Cargo.lock | 99 +++++++++++++++++++ osssendfile-rs/Cargo.toml | 3 + osssendfile-rs/src/main.rs | 198 +++++++++++++++++++++++++++++-------- 3 files changed, 261 insertions(+), 39 deletions(-) diff --git a/osssendfile-rs/Cargo.lock b/osssendfile-rs/Cargo.lock index 04989dc..c1ad951 100644 --- a/osssendfile-rs/Cargo.lock +++ b/osssendfile-rs/Cargo.lock @@ -129,12 +129,27 @@ version = "2.6.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b048fb63fd8b5923fc5aa7b340d8e156aec7ec02f0c78fa8a6ddc2613f6f71de" +[[package]] +name = "block-buffer" +version = "0.10.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3078c7629b62d3f0439517fa394996acacc5cbc91c5a20d8c658e77abd503a71" +dependencies = [ + "generic-array", +] + [[package]] name = "bumpalo" version = "3.16.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "79296716171880943b8470b5f8d03aa55eb2e645a4874bdbb28adb49162e012c" +[[package]] +name = "byteorder" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1fd0f2584146f6f2ef48085050886acf353beff7305ebd1ae69500e27c67f64b" + [[package]] name = "bytes" version = "1.7.1" @@ -247,6 +262,16 @@ dependencies = [ "typenum", ] +[[package]] +name = "digest" +version = "0.10.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9ed9a281f7bc9b7576e61468ba615a66a5c8cfdff42420a70aa82701a3b1e292" +dependencies = [ + "block-buffer", + "crypto-common", +] + [[package]] name = "dirs-next" version = "2.0.0" @@ -775,11 +800,14 @@ name = "osssendfile-rs" version = "0.1.0" dependencies = [ "aes-gcm-stream", + "base64", "clap", + "rand", "reqwest", "rust_util", "serde", "serde_json", + "sha2", "tokio", ] @@ -862,6 +890,15 @@ dependencies = [ "universal-hash", ] +[[package]] +name = "ppv-lite86" +version = "0.2.20" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "77957b295656769bb8ad2b6a6b09d897d94f05c41b069aede1fcdaa675eaea04" +dependencies = [ + "zerocopy", +] + [[package]] name = "proc-macro2" version = "1.0.86" @@ -880,6 +917,36 @@ dependencies = [ "proc-macro2", ] +[[package]] +name = "rand" +version = "0.8.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "34af8d1a0e25924bc5b7c43c079c942339d8f0a8b57c39049bef581b46327404" +dependencies = [ + "libc", + "rand_chacha", + "rand_core", +] + +[[package]] +name = "rand_chacha" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e6c10a63a0fa32252be49d21e7709d4d4baf8d231c2dbce1eaa8141b9b127d88" +dependencies = [ + "ppv-lite86", + "rand_core", +] + +[[package]] +name = "rand_core" +version = "0.6.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ec0be4795e2f6a28069bec0b5ff3e2ac9bafc99e6a9a7dc3547996c5c816922c" +dependencies = [ + "getrandom", +] + [[package]] name = "redox_syscall" version = "0.5.3" @@ -1123,6 +1190,17 @@ dependencies = [ "serde", ] +[[package]] +name = "sha2" +version = "0.10.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "793db75ad2bcafc3ffa7c68b215fee268f537982cd901d132f89c6343f3a3dc8" +dependencies = [ + "cfg-if", + "cpufeatures", + "digest", +] + [[package]] name = "shlex" version = "1.3.0" @@ -1704,6 +1782,27 @@ version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "589f6da84c646204747d1270a2a5661ea66ed1cced2631d546fdfb155959f9ec" +[[package]] +name = "zerocopy" +version = "0.7.35" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1b9b4fd18abc82b8136838da5d50bae7bdea537c574d8dc1a34ed098d6c166f0" +dependencies = [ + "byteorder", + "zerocopy-derive", +] + +[[package]] +name = "zerocopy-derive" +version = "0.7.35" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fa4f8080344d4671fb4e831a13ad1e68092748387dfc4f55e356242fae12ce3e" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "zeroize" version = "1.8.1" diff --git a/osssendfile-rs/Cargo.toml b/osssendfile-rs/Cargo.toml index dab272c..ed27486 100644 --- a/osssendfile-rs/Cargo.toml +++ b/osssendfile-rs/Cargo.toml @@ -5,9 +5,12 @@ edition = "2021" [dependencies] aes-gcm-stream = "0.2" +base64 = "0.22.1" clap = { version = "4.5", features = ["derive"] } +rand = "0.8.5" reqwest = "0.12" rust_util = "0.6.47" serde = { version = "1.0", features = ["derive"] } serde_json = "1.0" +sha2 = "0.10.8" tokio = { version = "1.39", features = ["full"] } diff --git a/osssendfile-rs/src/main.rs b/osssendfile-rs/src/main.rs index ff5730d..4ddb688 100644 --- a/osssendfile-rs/src/main.rs +++ b/osssendfile-rs/src/main.rs @@ -1,12 +1,50 @@ +use aes_gcm_stream::Aes128GcmStreamEncryptor; +use base64::engine::general_purpose::STANDARD; +use base64::Engine; use clap::Parser; -use rust_util::{opt_result, opt_value_result, util_file, XResult}; +use reqwest::{Client, Response}; +use rust_util::util_io::DEFAULT_BUF_SIZE; +use rust_util::{debugging, information, opt_result, opt_value_result, simple_error, success, util_file, XResult}; use serde::{Deserialize, Serialize}; +use serde_json::Value; +use sha2::Digest; +use sha2::Sha256; use std::collections::HashMap; use std::fs; -use std::path::PathBuf; +use std::fs::File; +use std::io::{BufReader, ErrorKind, Read, Write}; +use std::path::{Path, PathBuf}; const OSS_SEND_FILE_CONFIG_FILE: &str = "~/.jssp/config/osssendfile.json"; const CREATE_STS_URL: &str = "https://hatter.ink/oidc/create_sts.json"; +const ADD_DOC_URL: &str = "https://playsecurity.org/doc/addDoc.jsonp"; + +#[derive(Debug, Parser)] +#[command(name = "osssendfile-rs")] +#[command(about = "OSS send file Rust edition", long_about = None)] +struct OssSendFileArgs { + /// Config file, default location: ~/.jssp/config/osssendfile.json + #[arg(long, short = 'c')] + config: Option, + // /// Do not encrypt + // #[arg(long)] + // no_enc: bool, + // /// Do remove source file + // #[arg(long)] + // remove_source_file: bool, + // /// JWK + // #[arg(long, short = 'j')] + // jwk: Option, + /// Upload file path + #[arg(long, short = 'f')] + file: PathBuf, + /// File name, default use local file name + #[arg(long, short = 'F')] + filename: Option, + /// Keywords + #[arg(long, short = 'k')] + keywords: Option, +} #[derive(Clone, Debug, Serialize, Deserialize)] #[serde(rename_all = "camelCase")] @@ -25,53 +63,40 @@ struct OidcConfig { client_secret: String, } -#[derive(Debug, Parser)] -#[command(name = "osssendfile-rs")] -#[command(about = "OSS send file Rust edition", long_about = None)] -struct OssSendFileArgs { - /// Config file, default location: ~/.jssp/config/osssendfile.json - #[arg(long)] - config: Option, - /// Do not encrypt - #[arg(long)] - noenc: bool, - /// Do remove source file - #[arg(long)] - removesourcefile: bool, - /// JWK - #[arg(long)] - jwk: Option, - /// Upload file path - #[arg(long)] - file: PathBuf, - /// File name, default use local file name - #[arg(long)] - filename: Option, - /// Keywords - #[arg(long)] - keywords: Option, +#[derive(Clone, Debug, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +struct Sts { + access_key_id: String, + access_key_secret: String, + expiration: String, + security_token: String, } #[tokio::main] async fn main() -> XResult<()> { let args = OssSendFileArgs::parse(); + + let client = Client::new(); let oss_send_file_config = load_config(&args.config)?; + let sts = request_sts(&client, &oss_send_file_config).await?; - println!("{:?}", args); - println!("{:?}", oss_send_file_config); + println!("{:#?}", sts); - let client = reqwest::Client::new(); + let source_file = args.file.clone(); + let filename = source_file.file_name().unwrap().to_str().unwrap(); + let temp_file = match source_file.parent() { + None => PathBuf::from(format!("{}.tmp", filename)), + Some(parent_source_file) => { + parent_source_file.join(&format!("{}.tmp", filename)) + } + }; - let mut params = HashMap::new(); - params.insert("client_id", &oss_send_file_config.oidc.client_id); - params.insert("client_secret", &oss_send_file_config.oidc.client_secret); - params.insert("sub", &oss_send_file_config.oidc.sub); - let response = client.post(CREATE_STS_URL) - .form(¶ms) - .send() - .await?; + information!("Encrypt file {:?} -> {:?}", source_file, temp_file); + let temp_key = encrypt(&source_file, &temp_file)?; + debugging!("Encryption temp key: {:?}", temp_key); + information!("Encrypt file: {:?} success", temp_file); - // TODO ... + // TODO send file to OSS Ok(()) } @@ -88,3 +113,98 @@ fn load_config(config: &Option) -> XResult { serde_json::from_str(&config), "Parse file: {:?} failed: {}", config_file); Ok(oss_send_config) } + +fn new_aes_key_and_nonce() -> ([u8; 16], Vec) { + let temp_key: [u8; 16] = rand::random(); + let mut sha256 = Sha256::new(); + sha256.update(&temp_key); + let sha256_digest = sha256.finalize(); + let nonce = &sha256_digest.as_slice()[0..12]; + (temp_key, nonce.to_vec()) +} + +fn encrypt(source_file: &PathBuf, dest_file: &PathBuf) -> XResult<[u8; 16]> { + if !source_file.exists() { + return simple_error!("File {:?} not exists.", source_file); + } + if dest_file.exists() { + return simple_error!("File {:?} exists.", dest_file); + } + let (temp_key, nonce) = new_aes_key_and_nonce(); + let mut encryptor = Aes128GcmStreamEncryptor::new(temp_key, &nonce); + + let mut source_file_read = File::open(source_file)?; + let mut dest_file_write = File::create(dest_file)?; + + let mut written = 0u64; + let mut buf: [u8; DEFAULT_BUF_SIZE] = [0u8; DEFAULT_BUF_SIZE]; + loop { + let len = match source_file_read.read(&mut buf) { + Ok(0) => { + let (final_block, tag) = encryptor.finalize(); + dest_file_write.write(&final_block)?; + written += final_block.len() as u64; + dest_file_write.write(&tag)?; + written += tag.len() as u64; + return Ok(temp_key); + } + Ok(len) => len, + Err(ref e) if e.kind() == ErrorKind::Interrupted => continue, + Err(e) => return simple_error!("Encrypt file failed: {}", e), + }; + let encrypted = encryptor.update(&buf[0..len]); + dest_file_write.write(&encrypted)?; + written += encrypted.len() as u64; + } +} + +async fn request_sts(client: &Client, oss_send_file_config: &OssSendFileConfig) -> XResult { + let mut params = HashMap::new(); + params.insert("client_id", &oss_send_file_config.oidc.client_id); + params.insert("client_secret", &oss_send_file_config.oidc.client_secret); + params.insert("sub", &oss_send_file_config.oidc.sub); + let response = client.post(CREATE_STS_URL) + .form(¶ms) + .send() + .await?; + + parse_sts_response(response).await +} + +async fn send_file(client: &Client, token: &str, url: &str, title: &str, keywords: &str, key: [u8; 16]) -> XResult<()> { + let mut params = HashMap::new(); + params.insert("jsonp", "1".to_string()); + params.insert("token", token.to_string()); + params.insert("url", url.to_string()); + params.insert("title", title.to_string()); + params.insert("keywords", keywords.to_string()); + params.insert("encryption", STANDARD.encode(&key)); + + let response = client.post(ADD_DOC_URL) + .form(¶ms) + .send() + .await?; + + if !response.status().is_success() { + return simple_error!("Add doc failed, status: {}", response.status().as_u16()); + } + + let response_bytes = response.bytes().await?.as_ref().to_vec(); + let add_doc_result = String::from_utf8_lossy(&response_bytes); + success!("Add doc success: {}", add_doc_result); + Ok(()) +} + +async fn parse_sts_response(response: Response) -> XResult { + if !response.status().is_success() { + return simple_error!("Create STS failed, status: {}", response.status().as_u16()); + } + + let response_bytes = opt_result!(response.bytes().await, "Read STS failed: {}"); + let response_value: Value = opt_result!(serde_json::from_slice(&response_bytes), "Parse STS response failed: {}"); + + let data_value = opt_value_result!(response_value.get("data"), "Parse STS response no data found."); + let data_value_str = serde_json::to_string(data_value).expect("Should not happen."); + let sts: Sts = opt_result!(serde_json::from_str(&data_value_str), "Parse STS response data failed: {}"); + Ok(sts) +}