#!/usr/bin/env runrs //! ```cargo //! [dependencies] //! 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 = { version = "0.12", features = ["stream"] } //! 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"] } //! tokio-util = "0.7.11" //! ``` use aes_gcm_stream::Aes128GcmStreamEncryptor; use base64::engine::general_purpose::STANDARD; use base64::Engine; use clap::Parser; use oss::OSSClient; use reqwest::{Body, Client, Response}; use rust_util::util_io::DEFAULT_BUF_SIZE; use rust_util::{failure, information, opt_result, opt_value_result, simple_error, success, util_file, util_msg, util_size, util_time, warning, XResult}; use serde::{Deserialize, Serialize}; use serde_json::Value; use sha2::Digest; use sha2::Sha256; use std::collections::HashMap; use std::fs::File; use std::io::{ErrorKind, Read, Write}; use std::path::PathBuf; use std::{fs, io}; use std::process::Command; use tokio::fs::File as TokioFile; use tokio_util::bytes::BytesMut; use tokio_util::codec::{Decoder, FramedRead}; const OSS_SEND_FILE_CONFIG_FILE: &str = "~/.jssp/config/osssendfile.json"; const CREATE_STS_URL: &str = "https://global.hatter.ink/oidc/create_sts.json"; const ADD_DOC_URL: &str = "https://play.hatter.me/doc/addDoc.jsonp"; const ASSUME_ROLE_BY_KEY: &str = "https://global.hatter.ink/cloud/alibaba_cloud/assume_role_by_key.json"; const INSTANCE_IDENTITY_TOKEN: &str = "http://100.100.100.200/latest/api/token"; const INSTANCE_IDENTITY_TTL: &str = "X-aliyun-ecs-metadata-token-ttl-seconds"; const INSTANCE_IDENTITY_PKCS7: &str = "http://100.100.100.200/latest/dynamic/instance-identity/pkcs7"; const INSTANCE_IDENTITY_PKCS7_TOKEN: &str = "X-aliyun-ecs-metadata-token"; const ENV_OSS_SEND_FILE_CONFIG: &str = "OSS_SEND_FILE_CONFIG"; const TINY_ENCRYPT_CONFIG: &str = r##"{ "envelops": [ { "type": "piv-p384", "sid": "piv-83-ecdh-p384", "kid": "02c9f887e28c15e7d80ac176fba5ea271acae6ab473fe7414adca8e2566c791084296f8e568e64ab3900a57e906e66dbb3", "desc": "PIV --slot 83", "args": [ "83" ], "publicPart": "04c9f887e28c15e7d80ac176fba5ea271acae6ab473fe7414adca8e2566c791084296f8e568e64ab3900a57e906e66dbb3f99bf4b4b15f1b19e553ad43c2fea1e74eeeae159a16fe76c3fdb3d4df167413f7f43a35a1682648f48e5ec4aa7ad703" }, { "type": "key-p256", "sid": "macbook-m4-air-bio", "kid": "keychain:02126aaa5ef17f0879a42ac3766742e7e913741caf489f17676b106ceb41a78bf1", "desc": "Secure Enclave P256 require Bio @MacBook M4 Air", "publicPart": "04126aaa5ef17f0879a42ac3766742e7e913741caf489f17676b106ceb41a78bf17ecaf6cb06459456fdf37250d674dd2161f3cd2f636d9068d33dbeb435b3e858" }, { "type": "key-p256", "sid": "mbpse", "kid": "keychain:02c408fcb810f7d9ecceb4cca297a93b7d34c336feda4f06ac4553846099a32b38", "desc": "Secure Enclave P256 require Bio @MacBook Intel Pro", "publicPart": "04c408fcb810f7d9ecceb4cca297a93b7d34c336feda4f06ac4553846099a32b381afcb7ff2abb2b235a53cc6eb71894243bf573db0096a1af10a05bddfb70fc66" }, { "type": "pgp-x25519", "sid": "pgp-x25519-yubikey-5n", "kid": "C0FAD5E563B80E819603B0D9FFC2A910806894FD", "desc": "Card serial no. = 0006 16138686", "publicPart": "7FEBAAB0D80CED24730B613F3D86924560EBCF13A838DEBC065F63C69C24C61E" } ], "profiles": { "default": [ "02c9f887e28c15e7d80ac176fba5ea271acae6ab473fe7414adca8e2566c791084296f8e568e64ab3900a57e906e66dbb3", "keychain:02126aaa5ef17f0879a42ac3766742e7e913741caf489f17676b106ceb41a78bf1", "keychain:02c408fcb810f7d9ecceb4cca297a93b7d34c336feda4f06ac4553846099a32b38", "C0FAD5E563B80E819603B0D9FFC2A910806894FD" ] } }"##; #[derive(Debug, Parser)] #[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 #[arg(long, short = 'c')] config: Option, // /// Do not encrypt // #[arg(long)] // no_enc: bool, /// Do remove source file #[arg(long)] remove_source_file: bool, /// Tiny encrypt #[arg(long)] tiny_encrypt: 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")] struct OssSendFileConfig { endpoint: String, bucket: String, token: String, oidc: Option, pkcs7: Option, } #[derive(Clone, Debug, Serialize, Deserialize)] #[serde(rename_all = "camelCase")] struct OidcConfig { sub: String, client_id: String, client_secret: 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, } #[derive(Clone, Debug, Serialize, Deserialize)] #[serde(rename_all = "camelCase")] struct SendFileResponse { status: i32, id: i64, length: i64, etag: String, sha256: String, cdn_url: Option, } #[derive(Copy, Clone, Debug, Eq, PartialEq, Ord, PartialOrd, Hash, Default)] pub struct CustomBytesCodec { read_len: u64, total_len: u64, } impl CustomBytesCodec { pub fn new(total_len: u64) -> CustomBytesCodec { CustomBytesCodec { read_len: 0, total_len, } } } impl Decoder for CustomBytesCodec { type Item = BytesMut; type Error = io::Error; fn decode(&mut self, buf: &mut BytesMut) -> Result, io::Error> { if !buf.is_empty() { let len = buf.len(); self.read_len += len as u64; util_msg::print_lastline(&format!( "Copied {}, total: {}", util_size::get_display_size(self.read_len as i64), if self.total_len == 0 { "-".to_string() } else { util_size::get_display_size(self.total_len as i64) } )); Ok(Some(buf.split_to(len))) } else { Ok(None) } } } #[tokio::main] async fn main() -> XResult<()> { let args = OssSendFileArgs::parse(); let client = Client::new(); let oss_send_file_config = load_config(&args.config)?; information!("Get STS from hatter ink"); let sts = request_sts(&client, &oss_send_file_config).await?; information!("Get STS success"); let mut pending_remove_file = None::; let mut pending_remove_file_2 = None::; let mut source_file = args.file.clone(); let metadata = fs::metadata(&source_file)?; let filename = source_file.file_name().unwrap().to_str().unwrap().to_string(); // zip file when is dir let temp_zip_file = format!("{}.zip", filename); if metadata.is_dir() { if let Ok(_) = fs::metadata(&temp_zip_file) { return simple_error!("File {} exists", temp_zip_file); } zip_dir(&temp_zip_file, source_file.to_str().unwrap())?; information!("Zip {} to {} success", source_file.display(), temp_zip_file); source_file = PathBuf::from(temp_zip_file.clone()); pending_remove_file = Some(temp_zip_file.clone()); } let filename = source_file.file_name().unwrap().to_str().unwrap().to_string(); // encrypt file via tiny encrypt let temp_tiny_encrypt_file = format!("{}.tinyenc", filename); if args.tiny_encrypt { tiny_encrypt_file(source_file.to_str().unwrap())?; information!("Tiny encrypt {} to {} success", source_file.display(), temp_tiny_encrypt_file); source_file = PathBuf::from(temp_tiny_encrypt_file.clone()); pending_remove_file_2 = Some(temp_tiny_encrypt_file.clone()); } let filename = source_file.file_name().unwrap().to_str().unwrap().to_string(); let filename_ext = get_filename_ext(&filename); let temp_file = match source_file.parent() { None => PathBuf::from(format!("{}.tmp", filename)), Some(parent_source_file) => { parent_source_file.join(&format!("{}.tmp", filename)) } }; if let Ok(metadata) = source_file.metadata() { let file_len = metadata.len(); let file_len_display = util_size::get_display_size(file_len as i64); if file_len < 1024 * 1024 { information!("File length: {}", file_len_display); } else if file_len < 10 * 1024 * 1024 { warning!("File is large, length: {}", file_len_display); } else { failure!("File is very large, length: {}", file_len_display); } } information!("Encrypt file {:?} -> {:?}", source_file, temp_file); let temp_key = encrypt(&source_file, &temp_file)?; information!("Encrypt file: {:?} success", &temp_file); let oss_client = OSSClient::new_sts(&oss_send_file_config.endpoint, &sts.access_key_id, &sts.access_key_secret, &sts.security_token); let temp_oss_filename = format!("tempfiles/temp_transfer_{}.{}", util_time::get_current_millis(), filename_ext); let temp_file_len = temp_file.metadata().map(|m| m.len()).unwrap_or(0); information!("Put object {:?} -> {}", &temp_file, &temp_oss_filename); // thanks: https://stackoverflow.com/questions/65814450/how-to-post-a-file-using-reqwest let temp_file_read = opt_result!(TokioFile::open(&temp_file).await, "Read temp file failed: {}"); let temp_file_stream = FramedRead::new(temp_file_read, CustomBytesCodec::new(temp_file_len)); let temp_file_body = Body::wrap_stream(temp_file_stream); let put_object_response = oss_client.put_body( &oss_send_file_config.bucket, &temp_oss_filename, 600, temp_file_body).await?; util_msg::clear_lastline(); if !put_object_response.status().is_success() { return simple_error!("Put object failed, status: {}", put_object_response.status().as_u16()); } information!("Put object success"); 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); 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( &oss_send_file_config.bucket, &temp_oss_filename).await?; if !delete_object_response.status().is_success() { return simple_error!("Delete object failed, status: {}", delete_object_response.status().as_u16()); } information!("Delete object success"); match fs::remove_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), } } if let Some(cnd_url) = send_file_response.cdn_url { success!("CDN URL: {}", cnd_url); } if let Some(pending_remove_file) = pending_remove_file { fs::remove_file(&pending_remove_file)?; success!("Remove pending remove file: {} success", &pending_remove_file); } if let Some(pending_remove_file_2) = pending_remove_file_2 { fs::remove_file(&pending_remove_file_2)?; success!("Remove pending remove file: {} success", &pending_remove_file_2); } success!("File {} upload success", filename); Ok(()) } fn load_config(config: &Option) -> XResult { let config = match config { Some(config) => Some(config.clone()), None => std::env::var(ENV_OSS_SEND_FILE_CONFIG).ok(), }; if let Some(config_str) = &config { if let Ok(oss_send_config) = serde_json::from_str::(config_str) { return Ok(oss_send_config); } if let Ok(config_vec) = STANDARD.decode(config_str) { if let Ok(config_str_ori) = String::from_utf8(config_vec) { if let Ok(oss_send_config) = serde_json::from_str::(&config_str_ori) { return Ok(oss_send_config); } } } } let config_file_opt = util_file::read_config( config, &[OSS_SEND_FILE_CONFIG_FILE.to_string()], ); let config_file = opt_value_result!(config_file_opt, "Config file not found."); let config = opt_result!( fs::read_to_string(&config_file), "Read file: {:?} failed: {}", config_file); let oss_send_config: OssSendFileConfig = opt_result!( serde_json::from_str(&config), "Parse file: {:?} failed: {}", config_file); Ok(oss_send_config) } fn get_filename_ext(filename: &str) -> String { let split_filename_parts: Vec<&str> = filename.split(".").collect(); split_filename_parts[split_filename_parts.len() - 1].to_string() } fn get_filename(filename: &str) -> String { if filename.contains('/') { let filename_parts = filename.split("/").collect::>(); filename_parts[filename_parts.len() - 1].to_string() } else { filename.to_string() } } 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; } } fn digest_sha256(source_file: &PathBuf) -> XResult> { 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 { if let Some(oidc_config) = &oss_send_file_config.oidc { Ok(request_sts_via_oidc(client, oidc_config).await?) } else if let Some(role_arn) = &oss_send_file_config.pkcs7 { Ok(request_sts_via_pkcs7(client, role_arn).await?) } else { simple_error!("OIDC or PKCS#7 config not found.") } } async fn request_sts_via_oidc(client: &Client, oidc_config: &OidcConfig) -> XResult { let mut params = HashMap::new(); params.insert("client_id", &oidc_config.client_id); params.insert("client_secret", &oidc_config.client_secret); params.insert("sub", &oidc_config.sub); let response = client.post(CREATE_STS_URL) .form(¶ms) .send() .await?; parse_sts_response(response).await } #[derive(Debug, Deserialize)] struct StsTokenForAssumeRoleByKeyResponse { // pub mode: String, pub expiration: String, pub access_key_id: String, pub access_key_secret: String, pub sts_token: String, } #[derive(Debug, Deserialize)] struct AssumeRoleByKeyResponse { pub status: i32, // pub message: Option, pub data: Option, } async fn request_sts_via_pkcs7(client: &Client, role_arn: &str) -> XResult { let pkcs7 = fetch_alibaba_cloud_instance_identity_v1( client, "hatter.ink:/cloud/alibaba_cloud/assume_role_by_key.json").await?; let assumed_role_response = client.get(format!("{}?roleArn={}", ASSUME_ROLE_BY_KEY, role_arn)) .header("Authorization", format!("PKCS7 {}", pkcs7)) .send() .await?; if !assumed_role_response.status().is_success() { return simple_error!("Assume role by key failed, status: {}", assumed_role_response.status()); } let assume_role_json = assumed_role_response.text().await?; let assumed_role: AssumeRoleByKeyResponse = serde_json::from_str(&assume_role_json)?; if assumed_role.status != 200 { return simple_error!("Assume role by key failed, response: {:?}", assumed_role); } let data = assumed_role.data.unwrap(); Ok(Sts { access_key_id: data.access_key_id, access_key_secret: data.access_key_secret, expiration: data.expiration, security_token: data.sts_token, }) } 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 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 { 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) } fn tiny_encrypt_file(target: &str) -> XResult<()> { let tiny_encrypt_config = format!("base64:{}", STANDARD.encode(TINY_ENCRYPT_CONFIG)); let mut cmd = Command::new("tiny-encrypt"); cmd.args(["encrypt", "--config", &tiny_encrypt_config, target]); match rust_util::util_cmd::run_command_and_wait(&mut cmd) { Ok(exit_status) => if !exit_status.success() { return simple_error!("Encrypt file {} failed, status: {}", target, exit_status); } Err(e) => return simple_error!("Encrypt file {} failed: {}", target, e), } Ok(()) } fn zip_dir(temp_zip_file: &str, target: &str) -> XResult<()> { let mut cmd = Command::new("zip"); cmd.args(&["-r", temp_zip_file, target]); match rust_util::util_cmd::run_command_and_wait(&mut cmd) { Ok(exit_status) => if !exit_status.success() { return simple_error!("Zip {} to {} failed, status: {}", target, temp_zip_file, exit_status); } Err(e) => return simple_error!("Zip {} to {} failed: {}", target, temp_zip_file, e), } Ok(()) } #[derive(Debug, Serialize)] struct AlibabaCloudInstanceIdentityAudienceMeta { pub ver: i32, pub iat: i64, pub exp: i64, pub aud: String, pub jti: String, } async fn fetch_alibaba_cloud_instance_identity_v1(client: &Client, audience: &str) -> XResult { let instance_identity_token = client.put(INSTANCE_IDENTITY_TOKEN) .header(INSTANCE_IDENTITY_TTL, "60") .send() .await?; if !instance_identity_token.status().is_success() { return simple_error!("Get instance identity token failed, status: {}", instance_identity_token.status()); } let token = opt_result!(instance_identity_token.text().await, "Get instance identity token failed: {}"); let rand_bytes: [u8; 8] = rand::random(); let audience_meta = AlibabaCloudInstanceIdentityAudienceMeta { ver: 1, iat: util_time::get_current_secs() as i64, exp: (util_time::get_current_secs() + 60) as i64, aud: audience.to_string(), jti: format!("jti-{}-{}", util_time::get_current_millis(), hex::encode(rand_bytes)), }; let audience_meta_json = serde_json::to_string(&audience_meta)?; let instance_identity_pkcs7 = client.get( format!("{}?audience={}", INSTANCE_IDENTITY_PKCS7, audience_meta_json)) .header(INSTANCE_IDENTITY_PKCS7_TOKEN, token) .send() .await?; if !instance_identity_pkcs7.status().is_success() { return simple_error!("Get instance identity PKCS#7 failed, status: {}", instance_identity_pkcs7.status()); } let pkcs7 = opt_result!( instance_identity_pkcs7.text().await, "Get instance identity PKCS#7 failed: {}"); Ok(pkcs7) } // @SCRIPT-SIGNATURE-V1: yk-r1.ES256.20260315T233039+08:00.MEQCIFBq1TNXxRp+k11NmKe2 // lF+Q4aMOZHEB9Z6jF4rkJdSAAiALEgUf4c5HIF7P3I2FW6tjyV0I5jbC4oZUdxw0NvL1Yg==