zero4rs 2.0.0

use reqwest::Url;

use hmac::Mac;
use std::collections::BTreeMap;

use base64::{engine::general_purpose, Engine as _};

const AWS_ACCESS_KEY_ID_FIELD_KEY: &str = "AWSAccessKeyId";
const SIGNATURE_VERSION: &str = "2";
const SIGNATURE_METHOD: &str = "HmacSHA256";
const EXPIRE_SECONDS_FIELD_KEY: &str = "ExpireSeconds";
const TIMESTAMP_FIELD_KEY: &str = "Timestamp";
const SIGNATURE_FIELD_KEY: &str = "Signature";

pub fn pre_signed_url_v2(
    access_key_id: &str,
    secret_access_key: &str,
    expire_seconds: u32,
    method: &str,
    req_url: &str,
    time_format: &str,
) -> String {
    let parsed_url = Url::parse(req_url).unwrap();
    let datetime_str = get_timestamp(time_format);

    println!(
        "\nDATE_TIME_STRING: \n------------------\nFormat={}, Value={}\n",
        time_format, datetime_str
    );

    let mut query = parsed_url
        .query_pairs()
        .into_owned()
        .collect::<BTreeMap<_, _>>();

    // Set new query parameters
    query.insert(
        AWS_ACCESS_KEY_ID_FIELD_KEY.to_owned(),
        access_key_id.to_owned(),
    );

    query.insert("SignatureVersion".to_owned(), SIGNATURE_VERSION.to_owned());
    query.insert("SignatureMethod".to_owned(), SIGNATURE_METHOD.to_owned());

    if !query.contains_key(TIMESTAMP_FIELD_KEY) {
        query.insert(TIMESTAMP_FIELD_KEY.to_owned(), get_timestamp(time_format));
    }

    if expire_seconds > 0 {
        query.insert(
            EXPIRE_SECONDS_FIELD_KEY.to_string(),
            expire_seconds.to_string(),
        );
    }

    // in case this is a retry, ensure no signature present
    query.remove(SIGNATURE_FIELD_KEY);

    let host = match parsed_url.port() {
        Some(port) => format!("{}:{}", parsed_url.host_str().unwrap_or(""), port),
        None => parsed_url.host_str().unwrap_or("").to_string(),
    };

    let path = parsed_url.path();

    let query_keys = query
        .keys()
        .map(|key| key.to_owned())
        .collect::<Vec<String>>();

    let mut sorted_query_keys = query_keys.clone();

    sorted_query_keys.sort();

    let mut query_keys_and_values = Vec::new();

    for key in sorted_query_keys.iter() {
        query_keys_and_values.push(format!(
            "{}={}",
            escape_string(key),
            escape_string(&get_value(&query, key))
        ));
    }

    let query_str = query_keys_and_values.join("&");

    let string_to_sign = format!(
        "{}\n{}\n{}\n{}",
        method.to_uppercase(),
        host,
        path,
        query_str
    );

    let mut mac = hmac::Hmac::<sha2::Sha256>::new_from_slice(secret_access_key.as_bytes()).unwrap();

    mac.update(string_to_sign.as_bytes());

    let signature = escape_string(&general_purpose::STANDARD.encode(mac.finalize().into_bytes()));

    query.insert(SIGNATURE_FIELD_KEY.to_string(), signature.clone());

    let new_url_str = format!(
        "{}://{}{}?{}&{}={}",
        parsed_url.scheme(),
        host,
        path,
        query_str,
        SIGNATURE_FIELD_KEY,
        query.get(SIGNATURE_FIELD_KEY).unwrap()
    );

    println!(
        "PRE_SIGNED_URL_V2: \n------------------\n{}\n",
        string_to_sign
    );

    println!("Signature: \n------------------\n{}\n", signature);

    println!("Signed-Url: \n------------------\n{}\n", new_url_str);

    new_url_str
}

pub fn escape_string(str_val: &str) -> String {
    url::form_urlencoded::byte_serialize(str_val.as_bytes())
        .collect::<String>()
        .replace('+', "%20")
}

pub fn get_value(map: &BTreeMap<String, String>, key: &str) -> String {
    map.get(key).unwrap_or(&"".to_string()).to_string()
}

pub fn get_timestamp(_time_format: &str) -> String {
    let time_format = if !_time_format.ends_with('Z') {
        format!("{}Z", _time_format)
    } else {
        _time_format.to_string()
    };

    let r = chrono::Utc::now().format(&time_format).to_string();

    if time_format.ends_with('Z') && !_time_format.ends_with('Z') {
        r[0..r.len() - 1].to_string()
    } else {
        r
    }
}