aws_signer 0.0.4

use reqwest::header::{HeaderMap, HeaderValue, AUTHORIZATION};
use reqwest::{Client, Request, Response, Url};
use hmac::{Mac};
use sha2::{Digest, Sha256};
use chrono::Utc;
use std::collections::{HashMap, HashSet};
use std::error::Error;
use std::fmt;
use crate::signature::Signature;


// Define a custom error type
#[derive(Debug)]
struct AwsError(String);

impl fmt::Display for AwsError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "{}", self.0)
    }
}

impl Error for AwsError {}

// Helper function to handle service detection from hostname
fn get_service_from_host(host: &str) -> String {
    let mut host_services: HashMap<&str, &str> = HashMap::new();
    host_services.insert("appstream2", "appstream");
    host_services.insert("cloudhsmv2", "cloudhsm");
    host_services.insert("email", "ses");
    host_services.insert("marketplace", "aws-marketplace");
    host_services.insert("mobile", "AWSMobileHubService");
    host_services.insert("pinpoint", "mobiletargeting");
    host_services.insert("queue", "sqs");
    host_services.insert("git-codecommit", "codecommit");
    host_services.insert("mturk-requester-sandbox", "mturk-requester");
    host_services.insert("personalize-runtime", "personalize");
    let default_service = "s3";
    host_services.get(host).unwrap_or(&default_service).to_string()
}

// Set of headers that should not be signed
fn is_unsignable_header(header_name: &str) -> bool {
    let unsignable_headers: HashSet<&str> = [
        "authorization", "content-type", "content-length", "user-agent",
        "presigned-expires", "expect", "x-amzn-trace-id", "range", "connection",
    ].iter().cloned().collect();

    unsignable_headers.contains(header_name.to_lowercase().as_str())
}

fn encode_rfc3986(input: &str) -> String {
    input.replace("!", "%21")
        .replace("'", "%27")
        .replace("(", "%28")
        .replace(")", "%29")
        .replace("*", "%2A")
}

#[derive(Clone)]
pub struct AwsClient {
    access_key_id: String,
    secret_access_key: String,
    session_token: Option<String>,
    service: Option<String>,
    region: Option<String>,
    cache: HashMap<String, Vec<u8>>,
    retries: usize,
    init_retry_ms: u64,
}

impl AwsClient {
    pub fn new(
        access_key_id: String,
        secret_access_key: String,
        session_token: Option<String>,
        service: Option<String>,
        region: Option<String>,
        cache: Option<HashMap<String, Vec<u8>>>,
        retries: Option<usize>,
        init_retry_ms: Option<u64>,
    ) -> Self {
        if access_key_id.is_empty() {
            panic!("accessKeyId is a required option");
        }
        if secret_access_key.is_empty() {
            panic!("secretAccessKey is a required option");
        }
        let service = service.map(|s| s.to_lowercase()).or_else(|| Some("s3".to_string()));
        AwsClient {
            access_key_id,
            secret_access_key,
            session_token,
            service,
            region,
            cache: cache.unwrap_or_default(),
            retries: retries.unwrap_or(10),
            init_retry_ms: init_retry_ms.unwrap_or(50),
        }
    }

    pub fn default() -> Self {
        AwsClient {
            access_key_id: String::new(),
            secret_access_key: String::new(),
            session_token: None,
            service: None,
            region: None,
            cache: HashMap::new(),
            retries: 10,
            init_retry_ms: 50,
        }
    }

    pub async fn sign(&self, input: Request, init: Option<AwsRequestInit>) -> Result<Request, Box<dyn Error>> {
        let url = input.url().to_string();
        let method = input.method().as_str().to_string();
        let mut headers = input.headers().clone();
        let body = input.body().map(|b| b.as_bytes().unwrap_or_default().to_vec());
        let sign_query = init.as_ref().and_then(|i| i.aws.as_ref()).and_then(|o| o.sign_query).unwrap_or(false);
        let mut signer = AwsV4Signer::new(
            &method,
            &url,
            headers.clone(),
            body.clone(),
            self.access_key_id.clone(),
            self.secret_access_key.clone(),
            sign_query,
            self.session_token.clone(),
            self.service.clone(),
            self.region.clone(),
        );

        signer.sign().await?;

        // Update headers and URL in the request
        headers = signer.headers;
        let url = signer.url;
        let mut req = Request::new(method.parse()?, url);
        *req.headers_mut() = headers;
        if let Some(body) = body {
            *req.body_mut() = Some(body.into());
        }

        Ok(req)
    }

    pub async fn fetch(&self, input: Request, init: Option<AwsRequestInit>) -> Result<Response, Box<dyn Error>> {
        let client = Client::new();
        for i in 0..=self.retries {
            // Clone the request and init to ensure each retry starts fresh
            let signed_request = self.sign(input.try_clone().unwrap(), init.clone()).await?;
            let response = client.execute(signed_request).await;

            match response {
                Ok(res) if res.status().as_u16() < 500 && res.status().as_u16() != 429 => {
                    return Ok(res);
                }
                Ok(_) => (),
                Err(_) => (),
            }

            if i == self.retries {
                break;
            }
            let wait_time = self.init_retry_ms * 2_u64.pow(i as u32);
            tokio::time::sleep(std::time::Duration::from_millis(wait_time)).await;
        }
        Err(Box::new(AwsError("All retries failed".to_string())))
    }
}

#[derive(Debug, Clone)]  // Derive Clone for automatic implementation
pub struct AwsRequestInit {
    aws: Option<AwsOptions>,
}

impl AwsRequestInit {
    pub fn new(aws: Option<AwsOptions>) -> Self {
        AwsRequestInit { aws }
    }

    pub fn default() -> Self {
        AwsRequestInit { aws: None }
    }
}

#[derive(Debug, Clone)]  // Derive Clone for automatic implementation
pub struct AwsOptions {
    access_key_id: Option<String>,
    secret_access_key: Option<String>,
    session_token: Option<String>,
    service: Option<String>,
    region: Option<String>,
    datetime: Option<String>,
    sign_query: Option<bool>,
    append_session_token: Option<bool>,
    all_headers: Option<bool>,
    single_encode: Option<bool>,
}

impl AwsOptions {
    pub fn new(
        access_key_id: Option<String>,
        secret_access_key: Option<String>,
        session_token: Option<String>,
        service: Option<String>,
        region: Option<String>,
        datetime: Option<String>,
        sign_query: Option<bool>,
        append_session_token: Option<bool>,
        all_headers: Option<bool>,
        single_encode: Option<bool>,
    ) -> Self {
        AwsOptions {
            access_key_id,
            secret_access_key,
            session_token,
            service,
            region,
            datetime,
            sign_query,
            append_session_token,
            all_headers,
            single_encode,
        }
    }

    pub fn default() -> Self {
        AwsOptions {
            access_key_id: None,
            secret_access_key: None,
            session_token: None,
            service: None,
            region: None,
            datetime: None,
            sign_query: Some(false),
            append_session_token: Some(false),
            all_headers: Some(false),
            single_encode: Some(true),
        }
    }
}

pub struct AwsV4Signer {
    method: String,
    url: Url,
    headers: HeaderMap,
    body: Option<Vec<u8>>,
    access_key_id: String,
    secret_access_key: String,
    session_token: Option<String>,
    service: String,
    region: String,
    datetime: String,
    sign_query: bool,
    append_session_token: bool,
}

impl AwsV4Signer {
    pub fn new(
        method: &str,
        url: &str,
        headers: HeaderMap,
        body: Option<Vec<u8>>,
        access_key_id: String,
        secret_access_key: String,
        sign_query: bool,
        session_token: Option<String>,
        service: Option<String>,
        region: Option<String>,
    ) -> Self {
        let url = Url::parse(url).expect("Invalid URL");
        let datetime = Utc::now().format("%Y%m%dT%H%M%SZ").to_string();
        let (service, region) = if service.is_some() && region.is_some() {
            (service.unwrap(), region.unwrap())
        } else {
            guess_service_region(&url, &headers)
        };

        AwsV4Signer {
            method: method.to_string(),
            url,
            headers,
            body,
            access_key_id,
            secret_access_key,
            session_token,
            service,
            region,
            datetime,
            sign_query,
            append_session_token: false,
        }
    }

    pub fn default() -> Self {
        AwsV4Signer {
            method: "GET".to_string(),
            url: Url::parse("http://example.com").unwrap(),
            headers: HeaderMap::new(),
            body: None,
            access_key_id: String::new(),
            secret_access_key: String::new(),
            session_token: None,
            service: "s3".to_string(),
            region: "auto".to_string(),
            datetime: Utc::now().format("%Y%m%dT%H%M%SZ").to_string(),
            sign_query: false,
            append_session_token: false,
        }
    }

    pub async fn sign(&mut self) -> Result<(), Box<dyn Error>> {
        if self.sign_query {
            self.url.query_pairs_mut().append_pair("X-Amz-Date", &self.datetime);

            // Add X-Amz-Expires parameter with a default value if not already set
            if !self.url.query_pairs().any(|(k, _)| k == "X-Amz-Expires") {
                self.url.query_pairs_mut().append_pair("X-Amz-Expires", "86400"); // 24 hours
            }

            // Add X-Amz-Algorithm parameter
            self.url.query_pairs_mut().append_pair("X-Amz-Algorithm", "AWS4-HMAC-SHA256");

            // Add X-Amz-Credential parameter
            let credential_scope = format!("{}/{}/{}/aws4_request", &self.datetime[..8], self.region, self.service);
            self.url.query_pairs_mut().append_pair("X-Amz-Credential", &format!("{}/{}", self.access_key_id, credential_scope));

            // Add X-Amz-SignedHeaders parameter
            let signed_headers = "host"; // Example; this should include all relevant headers
            self.url.query_pairs_mut().append_pair("X-Amz-SignedHeaders", signed_headers);

            if let Some(token) = &self.session_token {
                if self.append_session_token {
                    self.url.query_pairs_mut().append_pair("X-Amz-Security-Token", token);
                }
            }
        }

        let signature = Signature::new(
            self.secret_access_key.clone(),
            self.region.clone(),
            self.service.clone(),
        ).signature();
        self.url.query_pairs_mut().append_pair("X-Amz-Signature", &signature);

        if !self.sign_query {
            let auth_header = self.auth_header().await?;
            self.headers.insert(AUTHORIZATION, HeaderValue::from_str(&auth_header)?);
        }

        Ok(())
    }

    async fn auth_header(&self) -> Result<String, Box<dyn Error>> {
        let signature = Signature::new(
            self.secret_access_key.clone(),
            self.region.clone(),
            self.service.clone(),
        ).signature();
        let signed_headers = "host;x-amz-date"; // Example; expand as needed
        let credential_scope = format!("{}/{}/{}/aws4_request", &self.datetime[..8], self.region, self.service);
        let authorization_header = format!(
            "AWS4-HMAC-SHA256 Credential={}/{}, SignedHeaders={}, Signature={}",
            self.access_key_id, credential_scope, signed_headers, signature
        );
        Ok(authorization_header)
    }

    async fn string_to_sign(&self) -> Result<String, Box<dyn Error>> {
        let canonical_request = self.canonical_request().await?;
        let hashed_request = hex::encode(Sha256::digest(canonical_request.as_bytes()));
        let credential_scope = format!("{}/{}/{}/aws4_request", &self.datetime[..8], self.region, self.service);
        Ok(format!(
            "AWS4-HMAC-SHA256\n{}\n{}\n{}",
            self.datetime, credential_scope, hashed_request
        ))
    }

    async fn canonical_request(&self) -> Result<String, Box<dyn Error>> {
        let canonical_uri = self.url.path();
        let canonical_query_string = self.url.query().unwrap_or("");

        let mut canonical_headers = String::new();
        let mut signed_headers = Vec::new();
        for (key, value) in self.headers.iter() {
            let key_str = key.as_str().to_lowercase();
            if !is_unsignable_header(&key_str) {
                let value_str = value.to_str()?.trim().replace("\\s+", " ");
                canonical_headers.push_str(&format!("{}:{}\n", key_str, value_str));
                signed_headers.push(key_str);
            }
        }
        signed_headers.sort();
        let signed_headers = signed_headers.join(";");

        let payload_hash = "UNSIGNED-PAYLOAD"; // Use actual payload hash if needed

        Ok(format!(
            "{}\n{}\n{}\n{}\n{}\n{}",
            self.method, canonical_uri, canonical_query_string, canonical_headers, signed_headers, payload_hash
        ))
    }
}


fn guess_service_region(url: &Url, headers: &HeaderMap) -> (String, String) {
    let host = url.host_str().unwrap_or("");
    let service = get_service_from_host(host);
    let region = if host.ends_with(".r2.cloudflarestorage.com") {
        "auto".to_string()
    } else {
        "auto".to_string() // Default region
    };

    (service, region)
}

#[cfg(test)]
mod tests {
    use super::*;
    use reqwest::{Method, Url};
    use reqwest::header::HeaderMap;

    #[tokio::test]
    async fn test_sign_request() {
        let access_key_id = "test_access_key".to_string();
        let secret_key = "test_secret".to_string();
        let region = "auto".to_string();
        let client = AwsClient::new(
            access_key_id.clone(),
            secret_key.clone(),
            None,
            Some("s3".to_string()),
            Some(region.clone()),
            None,
            None,
            None,
        );

        let url = Url::parse("https://test-bucket.s3.amazonaws.com/test-object").unwrap();
        let request = Request::new(Method::GET, url.clone());

        let signed_request = client.sign(request, None).await.unwrap();
        assert!(signed_request.url().as_str().contains("https://test-bucket.s3.amazonaws.com/test-object"));

        let headers = signed_request.headers();
        assert!(headers.contains_key(AUTHORIZATION));
    }

    #[test]
    fn test_guess_service_region() {
        let url = Url::parse("https://example-bucket.r2.cloudflarestorage.com").unwrap();
        let headers = HeaderMap::new();

        let (service, region) = guess_service_region(&url, &headers);
        println!("Service: {}, Region: {}", service, region);
        assert_eq!(service, "s3");
        assert_eq!(region, "auto");
    }
}