fluidattacks-core-git 0.36.0

use anyhow::{bail, Context, Result};
use regex::Regex;
use uuid::Uuid;

pub struct AwsTempCreds {
    pub access_key: String,
    pub secret_key: String,
    pub session_token: String,
    pub region: String,
}

struct Hop<'a> {
    role_arn: &'a str,
    external_id: Option<&'a str>,
}

/// Chained when the caller is not the principal the client's trust policy pins.
const fn chain<'a>(
    arn: &'a str,
    external_id: &'a str,
    via: Option<&'a str>,
) -> (Hop<'a>, Option<Hop<'a>>) {
    let client = Hop {
        role_arn: arn,
        external_id: Some(external_id),
    };

    match via {
        Some(intermediate) => (
            Hop {
                role_arn: intermediate,
                external_id: None,
            },
            Some(client),
        ),
        None => (client, None),
    }
}

pub async fn assume_role(
    arn: &str,
    external_id: &str,
    url: &str,
    via: Option<&str>,
) -> Result<AwsTempCreds> {
    let region = extract_region(url)?;
    let (first, then) = chain(arn, external_id, via);

    let mut creds = assume(&first, &region, None).await?;
    if let Some(hop) = then {
        creds = assume(&hop, &region, Some(&creds)).await?;
    }

    Ok(creds)
}

async fn assume(hop: &Hop<'_>, region: &str, from: Option<&AwsTempCreds>) -> Result<AwsTempCreds> {
    let mut config = aws_config::defaults(aws_config::BehaviorVersion::latest())
        .region(aws_config::Region::new(region.to_owned()));

    if let Some(session) = from {
        config = config.credentials_provider(aws_sdk_sts::config::Credentials::new(
            session.access_key.clone(),
            session.secret_key.clone(),
            Some(session.session_token.clone()),
            None,
            "fluidattacks-core-chain",
        ));
    }

    let sts = aws_sdk_sts::Client::new(&config.load().await);

    let mut request = sts
        .assume_role()
        .role_arn(hop.role_arn)
        .role_session_name(format!("session-{}", Uuid::new_v4()));

    if let Some(external_id) = hop.external_id {
        request = request.external_id(external_id);
    }

    let resp = request
        .send()
        .await
        .with_context(|| format!("STS AssumeRole failed for {}", hop.role_arn))?;

    let creds = resp
        .credentials()
        .ok_or_else(|| anyhow::anyhow!("no credentials in AssumeRole response"))?;

    let access_key = creds.access_key_id().to_owned();
    let secret_key = creds.secret_access_key().to_owned();
    let session_token = creds.session_token().to_owned();

    if access_key.is_empty() || secret_key.is_empty() {
        bail!("empty credentials from STS AssumeRole");
    }

    Ok(AwsTempCreds {
        access_key,
        secret_key,
        session_token,
        region: region.to_owned(),
    })
}

fn extract_region(url: &str) -> Result<String> {
    let re = Regex::new(r"codecommit::([a-z0-9-]+)://")?;
    if let Some(caps) = re.captures(url) {
        return Ok(caps[1].to_string());
    }

    // Try to find region in standard CodeCommit URLs
    let re2 = Regex::new(r"git-codecommit\.([a-z0-9-]+)\.amazonaws\.com")?;
    if let Some(caps) = re2.captures(url) {
        return Ok(caps[1].to_string());
    }

    Ok("us-east-1".to_owned())
}

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

    #[test]
    fn one_call_takes_the_client_role_with_its_external_id() {
        let (first, then) = chain("arn:aws:iam::1:role/client", "an-external-id", None);

        assert_eq!(first.role_arn, "arn:aws:iam::1:role/client");
        assert_eq!(first.external_id, Some("an-external-id"));
        assert!(then.is_none());
    }

    #[test]
    fn a_chain_takes_the_intermediate_role_first_and_without_the_external_id() {
        let (first, then) = chain(
            "arn:aws:iam::1:role/client",
            "an-external-id",
            Some("arn:aws:iam::2:role/intermediate"),
        );

        assert_eq!(first.role_arn, "arn:aws:iam::2:role/intermediate");
        assert_eq!(first.external_id, None);

        let client = then.expect("the client role is still taken");
        assert_eq!(client.role_arn, "arn:aws:iam::1:role/client");
        assert_eq!(client.external_id, Some("an-external-id"));
    }
}