scriptbase/osssendfile-rs/src/main.rs

#!/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 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";

#[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,
    // /// 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: OidcConfig,
}

#[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 source_file = args.file.clone();
    let filename = source_file.file_name().unwrap().to_str().unwrap();
    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);
    }

    success!("File {} upload success", filename);
    Ok(())
}

fn load_config(config: &Option<String>) -> XResult<OssSendFileConfig> {
    let config_file_opt = util_file::read_config(
        config.clone(),
        &[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> {
    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<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)
}

// @SCRIPT-SIGNATURE-V1: yk-r1.ES256.20251010T002421+08:00.MEQCIDpVktuwmcLXHLxEDQYk
// C22TommHKEEhA6lpSqVGtIXqAiBMe7RqBUIfXYUd/+4Nu960nGuQq0m1Z6xcE527UOH0Zg==