rust-scripts/scriptbase
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

feat: update osssendfile-rs, check sha256

Browse Source
This commit is contained in:
Hatter Jiang
2024-08-25 10:32:19 +08:00
parent a1ff0294e1
commit b21964eb41
3 changed files with 62 additions and 19 deletions
Download Patch File Download Diff File Expand all files Collapse all files

73
osssendfile-rs/src/main.rs
View File

@@ -5,6 +5,7 @@
//! aes-gcm-stream = "0.2"
//! base64 = "0.22.1"
//! clap = { version = "4.5", features = ["derive"] }
//! hex = "0.4.3"
//! oss = "0.3.2"
//! rand = "0.8.5"
//! reqwest = "0.12"
@@ -41,7 +42,7 @@ 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(name = "osssendfile-rs", bin_name = "osssendfile.rs")]
#[command(about = "OSS send file Rust edition", long_about = None)]
struct OssSendFileArgs {
/// Config file, default location: ~/.jssp/config/osssendfile.json
@@ -51,8 +52,8 @@ struct OssSendFileArgs {
// #[arg(long)]
// no_enc: bool,
// /// Do remove source file
// #[arg(long)]
// remove_source_file: bool,
#[arg(long)]
remove_source_file: bool,
// /// JWK
// #[arg(long, short = 'j')]
// jwk: Option<String>,
@@ -93,6 +94,16 @@ struct Sts {
security_token: String,
}
#[derive(Clone, Debug, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
struct SendFileResponse {
status: i32,
id: i64,
length: i64,
etag: String,
sha256: String,
}
#[tokio::main]
async fn main() -> XResult<()> {
let args = OssSendFileArgs::parse();
@@ -147,13 +158,13 @@ async fn main() -> XResult<()> {
let object_url = oss_client.generate_signed_get_url(&oss_send_file_config.bucket, &temp_oss_filename, 600);
information!("Send URL to play security: {}", object_url);
send_file(&client,
&oss_send_file_config.token,
&object_url,
&args.filename.clone().unwrap_or_else(|| get_filename(filename)),
args.keywords.as_deref().unwrap_or(""),
temp_key).await?;
information!("Send URL success");
let send_file_response = send_file(&client,
&oss_send_file_config.token,
&object_url,
&args.filename.clone().unwrap_or_else(|| get_filename(filename)),
args.keywords.as_deref().unwrap_or(""),
temp_key).await?;
information!("Send URL success, file id: {}", &send_file_response.id);
information!("Delete object: {}", &temp_oss_filename);
let delete_object_response = oss_client.delete_file(
@@ -164,10 +175,24 @@ async fn main() -> XResult<()> {
information!("Delete object success");
match fs::remove_file(&temp_file) {
Ok(_) => information!("Delete temp file: {:?}", &temp_file),
Ok(_) => information!("Delete temp file success: {:?}", &temp_file),
Err(e) => warning!("Delete temp file failed: {}", e),
}
// check digest
let source_file_sha256 = hex::encode(digest_sha256(&source_file)?);
if source_file_sha256 != send_file_response.sha256 {
failure!("Check send file digest failed: {} vs {}", source_file_sha256, send_file_response.sha256);
return simple_error!("Check send file digest failed.");
}
if args.remove_source_file {
match fs::remove_file(&source_file) {
Ok(_) => information!("Delete source file success: {:?}", &source_file),
Err(e) => warning!("Delete source file failed: {}", e),
}
}
success!("File {} upload success", filename);
Ok(())
}
@@ -221,16 +246,16 @@ fn encrypt(source_file: &PathBuf, dest_file: &PathBuf) -> XResult<[u8; 16]> {
let mut source_file_read = File::open(source_file)?;
let mut dest_file_write = File::create(dest_file)?;
let mut written = 0u64;
// 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;
// written += final_block.len() as u64;
dest_file_write.write(&tag)?;
written += tag.len() as u64;
// written += tag.len() as u64;
return Ok(temp_key);
}
Ok(len) => len,
@@ -239,10 +264,18 @@ fn encrypt(source_file: &PathBuf, dest_file: &PathBuf) -> XResult<[u8; 16]> {
};
let encrypted = encryptor.update(&buf[0..len]);
dest_file_write.write(&encrypted)?;
written += encrypted.len() as u64;
// written += encrypted.len() as u64;
}
}
fn digest_sha256(source_file: &PathBuf) -> XResult<Vec<u8>> {
let mut sha256 = sha2::Sha256::new();
let content = opt_result!(fs::read(&source_file), "Read file: {:?} failed: {}", source_file);
sha256.update(&content);
let digest = sha256.finalize().as_slice().to_vec();
Ok(digest)
}
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);
@@ -256,7 +289,7 @@ async fn request_sts(client: &Client, oss_send_file_config: &OssSendFileConfig)
parse_sts_response(response).await
}
async fn send_file(client: &Client, token: &str, url: &str, title: &str, keywords: &str, key: [u8; 16]) -> XResult<()> {
async fn send_file(client: &Client, token: &str, url: &str, title: &str, keywords: &str, key: [u8; 16]) -> XResult<SendFileResponse> {
let mut params = HashMap::new();
params.insert("jsonp", "1".to_string());
params.insert("token", token.to_string());
@@ -275,9 +308,11 @@ async fn send_file(client: &Client, token: &str, url: &str, title: &str, keyword
}
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(())
let send_file_response: SendFileResponse = opt_result!(
serde_json::from_slice(&response_bytes), "Parse send file response failed: {}");
success!("Add doc id: {} success, sha256: {}", send_file_response.id, send_file_response.sha256);
Ok(send_file_response)
}
async fn parse_sts_response(response: Response) -> XResult<Sts> {