s3-presign 0.0.1

use chrono::{DateTime, Utc};
use hmac::{Hmac, Mac};
use percent_encoding::percent_encode;
use sha2::{Digest, Sha256};
use url::Url;

type HmacSha256 = Hmac<Sha256>;

const LONG_DATETIME_FMT: &str = "%Y%m%dT%H%M%SZ";
const SHORT_DATE_FMT: &str = "%Y%m%d";

const PERCENT_ENCODING_CHARSET: percent_encoding::AsciiSet =
    percent_encoding::CONTROLS.add(b'/').add(b':').add(b'+');
// Safe characters: https://docs.aws.amazon.com/AmazonS3/latest/userguide/object-keys.html
const S3_KEY_PERCENT_ENCODING_CHARSET: percent_encoding::AsciiSet =
    percent_encoding::NON_ALPHANUMERIC
        .remove(b'/')
        .remove(b'-')
        .remove(b'!')
        .remove(b'_')
        .remove(b'.')
        .remove(b'*')
        .remove(b'\'')
        .remove(b'(')
        .remove(b')');

// AWS_ACCESS_KEY_ID,
// AWS_SECRET_ACCESS_KEY
// AWS_SESSION_TOKEN
pub struct Credentials {
    access_key: String,
    secret_key: String,
    // ref: https://docs.aws.amazon.com/STS/latest/APIReference/CommonParameters.html
    session_token: Option<String>,
}

impl Credentials {
    pub fn new(access_key: String, secret_key: String, session_token: Option<String>) -> Self {
        Self {
            access_key,
            secret_key,
            session_token,
        }
    }

    pub fn new_temporary(access_key: &str, secret_key: &str, session_token: &str) -> Self {
        Self {
            access_key: access_key.into(),
            secret_key: secret_key.into(),
            session_token: Some(session_token.into()),
        }
    }
}

pub struct Bucket {
    region: String,
    bucket: String,
}

impl Bucket {
    pub fn new(region: &str, bucket: &str) -> Self {
        Self {
            region: region.into(),
            bucket: bucket.into(),
        }
    }
}

pub fn get(credentials: &Credentials, bucket: &Bucket, key: &str, expires: i64) -> Option<String> {
    let url = format!(
        "https://{}.s3.amazonaws.com/{}",
        bucket.bucket,
        escape_key(key)
    );
    let now = Utc::now();

    presigned_url(
        &credentials.access_key,
        &credentials.secret_key,
        credentials.session_token.as_ref().map(|s| s.as_str()),
        expires as _,
        &url.parse().unwrap(),
        "GET",
        "UNSIGNED-PAYLOAD",
        &bucket.region,
        &now,
    )
}

pub fn put(credentials: &Credentials, bucket: &Bucket, key: &str, expires: i64) -> Option<String> {
    //let url = format!("https://{}.s3.amazonaws.com/{}", bucket.bucket, key);
    let url = format!(
        "https://s3.amazonaws.com/{}/{}",
        bucket.bucket,
        escape_key(key)
    );
    let now = Utc::now();

    presigned_url(
        &credentials.access_key,
        &credentials.secret_key,
        credentials.session_token.as_ref().map(|s| s.as_str()),
        expires as _,
        &url.parse().unwrap(),
        "PUT",
        "UNSIGNED-PAYLOAD",
        &bucket.region,
        &now,
    )
}

fn escape_key(key: &str) -> String {
    if key.contains("%") {
        key.into() // assume esacped
    } else {
        percent_encode(key.as_bytes(), &S3_KEY_PERCENT_ENCODING_CHARSET).to_string()
    }
}

/// Generate pre-signed URL
fn presigned_url(
    access_key: &str,
    secret_key: &str,
    session_token: Option<&str>,
    expiration: u64,
    url: &Url,
    method: &str,
    payload_hash: &str,
    region: &str,
    date_time: &DateTime<Utc>,
) -> Option<String> {
    const SERVICE: &str = "s3";

    let date_time_txt = date_time.format(LONG_DATETIME_FMT).to_string();
    let short_date_time_txt = date_time.format(SHORT_DATE_FMT).to_string();
    let credentials = format!(
        "{}/{}/{}/s3/aws4_request",
        access_key, short_date_time_txt, region
    );
    let mut params = vec![
        (
            "X-Amz-Algorithm".to_string(),
            "AWS4-HMAC-SHA256".to_string(),
        ),
        ("X-Amz-Credential".to_string(), credentials),
        ("X-Amz-Date".to_string(), date_time_txt),
        ("X-Amz-Expires".to_string(), expiration.to_string()),
        ("X-Amz-SignedHeaders".to_string(), "host".to_string()),
    ];
    if let Some(session_token) = session_token {
        params.push((
            "X-Amz-Security-Token".to_string(),
            session_token.to_string(),
        ));
    }

    url.query_pairs().for_each(|(k, v)| {
        params.push((k.to_string(), v.to_string()));
    });

    params.sort();

    let canonical_query_string = params
        .iter()
        .map(|(k, v)| {
            format!(
                "{}={}",
                percent_encode(k.as_bytes(), &PERCENT_ENCODING_CHARSET),
                percent_encode(v.as_bytes(), &PERCENT_ENCODING_CHARSET)
            )
        })
        .collect::<Vec<_>>()
        .join("&");

    let canonical_resource = url.path();
    // FIXME: should handle port
    let canonical_headers = "host:".to_owned()
        + &url
            .host()
            .expect("host is provided in source code; qed")
            .to_string();
    let signed_headers = "host";
    let canonical_request = format!(
        "{}\n{}\n{}\n{}\n\n{}\n{}",
        method.to_uppercase(),
        canonical_resource,
        canonical_query_string,
        canonical_headers,
        signed_headers,
        payload_hash
    );
    let string_to_sign = string_to_sign(&date_time, &region, &canonical_request);
    let signing_key = signing_key(&date_time, secret_key, region, SERVICE)?;

    let mut hmac = HmacSha256::new_from_slice(&signing_key).ok()?;
    hmac.update(string_to_sign.as_bytes());
    let signature = format!("{:x}", hmac.finalize().into_bytes());

    let request_url =
        url.to_string() + "?" + &canonical_query_string + "&X-Amz-Signature=" + &signature;

    Some(request_url)
}

/// Generate the "string to sign" - the value to which the HMAC signing is
/// applied to sign requests.
fn string_to_sign(date_time: &DateTime<Utc>, region: &str, canonical_req: &str) -> String {
    let mut hasher = Sha256::default();
    hasher.update(canonical_req.as_bytes());
    format!(
        "AWS4-HMAC-SHA256\n{timestamp}\n{scope}\n{hash}",
        timestamp = date_time.format(LONG_DATETIME_FMT),
        scope = scope_string(date_time, region),
        hash = format!("{:x}", hasher.finalize())
    )
}

/// Generate the AWS signing key, derived from the secret key, date, region,
/// and service name.
fn signing_key(
    date_time: &DateTime<Utc>,
    secret_key: &str,
    region: &str,
    service: &str,
) -> Option<Vec<u8>> {
    let secret = format!("AWS4{}", secret_key);
    let mut date_hmac = HmacSha256::new_from_slice(secret.as_bytes()).ok()?;
    date_hmac.update(date_time.format(SHORT_DATE_FMT).to_string().as_bytes());
    let mut region_hmac = HmacSha256::new_from_slice(&date_hmac.finalize().into_bytes()).ok()?;
    region_hmac.update(region.to_string().as_bytes());
    let mut service_hmac = HmacSha256::new_from_slice(&region_hmac.finalize().into_bytes()).ok()?;
    service_hmac.update(service.as_bytes());
    let mut signing_hmac =
        HmacSha256::new_from_slice(&service_hmac.finalize().into_bytes()).ok()?;
    signing_hmac.update(b"aws4_request");
    Some(signing_hmac.finalize().into_bytes().to_vec())
}

/// Generate an AWS scope string.
fn scope_string(date_time: &DateTime<Utc>, region: &str) -> String {
    format!(
        "{date}/{region}/s3/aws4_request",
        date = date_time.format(SHORT_DATE_FMT),
        region = region
    )
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_generate() {
        let credentials = Credentials {
            access_key: "ASIAAAAAABBBBBCCCCCDDDDDD".into(),
            secret_key: "AAAAAAA+BBBBBBBB/CCCCCCC/DDDDDDDDDD".into(),
            session_token: Some("xxxxxxxxx".into()),
        };

        let bucket = Bucket {
            region: "us-east-1".into(),
            bucket: "the-bucket".into(),
        };

        let s = put(
            &credentials,
            &bucket,
            "5e4ed04f-1d37-4cef-8210-eea624f2aef5/f2199b3b-7924-4891-863e-800f83765ee0/e.4d346624df83aa0f1514456249042f5d93637e2675f048b9424443d8177f15d9f29a2644fdfb",
            600,
        );
        assert!(s.is_some());
        println!("=> {:?}", s);
    }
}