Files
scriptbase/osssendfile-rs/src/main.rs
2026-03-17 23:57:14 +08:00

614 lines
23 KiB
Rust
Executable File

#!/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<String>,
// /// 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<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")]
struct OssSendFileConfig {
endpoint: String,
bucket: String,
token: String,
oidc: Option<OidcConfig>,
pkcs7: Option<String>,
}
#[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<String>,
}
#[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<Option<BytesMut>, 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::<String>;
let mut pending_remove_file_2 = None::<String>;
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<String>) -> XResult<OssSendFileConfig> {
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::<OssSendFileConfig>(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::<OssSendFileConfig>(&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::<Vec<&str>>();
filename_parts[filename_parts.len() - 1].to_string()
} else {
filename.to_string()
}
}
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;
}
}
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> {
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<Sts> {
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(&params)
.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<String>,
pub data: Option<StsTokenForAssumeRoleByKeyResponse>,
}
async fn request_sts_via_pkcs7(client: &Client, role_arn: &str) -> XResult<Sts> {
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<SendFileResponse> {
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 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> {
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<String> {
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==