feat: update osssendfile-rs

2024-08-24 01:26:23 +08:00
parent e2329d96a6
commit bb5f88a775
3 changed files with 261 additions and 39 deletions
--- 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"
--- 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"] }
--- 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<String>,
+    // /// 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<String>,
+    /// Upload file path
+    #[arg(long, short = 'f')]
+    file: PathBuf,
+    /// File name, default use local file name
+    #[arg(long, short = 'F')]
+    filename: Option<String>,
+    /// Keywords
+    #[arg(long, short = 'k')]
+    keywords: Option<String>,
+}

 #[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<String>,
-    /// Do not encrypt
-    #[arg(long)]
-    noenc: bool,
-    /// Do remove source file
-    #[arg(long)]
-    removesourcefile: bool,
-    /// JWK
-    #[arg(long)]
-    jwk: Option<String>,
-    /// Upload file path
-    #[arg(long)]
-    file: PathBuf,
-    /// File name, default use local file name
-    #[arg(long)]
-    filename: Option<String>,
-    /// Keywords
-    #[arg(long)]
-    keywords: Option<String>,
+#[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(&params)
-        .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<String>) -> XResult<OssSendFileConfig> {
        serde_json::from_str(&config), "Parse file: {:?} failed: {}", config_file);
    Ok(oss_send_config)
 }
+
+fn new_aes_key_and_nonce() -> ([u8; 16], Vec<u8>) {
+    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<Sts> {
+    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(&params)
+        .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(&params)
+        .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<Sts> {
+    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)
+}