feat: update osssendfile-rs
This commit is contained in:
@@ -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(¶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<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(¶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<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)
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user