tryaudex_core/

learn.rs

use std::collections::BTreeSet;
use std::time::Duration;

use crate::error::{AvError, Result};

/// Maps CloudTrail eventSource (e.g. "s3.amazonaws.com") to IAM service prefix.
/// Public so `watch.rs` can reuse the same mapping table.
pub fn event_source_to_service(source: &str) -> Option<&str> {
    // Support standard, GovCloud (.amazonaws.com), and China (.amazonaws.com.cn) partitions
    let s = source
        .strip_suffix(".amazonaws.com.cn")
        .or_else(|| source.strip_suffix(".amazonaws.com"))?;
    Some(match s {
        "s3" => "s3",
        "lambda" => "lambda",
        "dynamodb" => "dynamodb",
        "ec2" => "ec2",
        "iam" => "iam",
        "sts" => "sts",
        "sqs" => "sqs",
        "sns" => "sns",
        "logs" => "logs",
        "ssm" => "ssm",
        "cloudformation" => "cloudformation",
        "cloudwatch" => "cloudwatch",
        "monitoring" => "cloudwatch",
        "kms" => "kms",
        "secretsmanager" => "secretsmanager",
        "elasticloadbalancing" => "elasticloadbalancing",
        "autoscaling" => "autoscaling",
        "ecr" => "ecr",
        "ecs" => "ecs",
        "eks" => "eks",
        "rds" => "rds",
        "route53" => "route53",
        "cloudfront" => "cloudfront",
        "apigateway" => "apigateway",
        "cognito-idp" | "cognito-identity" => "cognito-idp",
        "events" => "events",
        "states" => "states",
        "kinesis" => "kinesis",
        "firehose" => "firehose",
        "athena" => "athena",
        "glue" => "glue",
        "redshift" => "redshift",
        "elasticache" => "elasticache",
        other => {
            tracing::warn!(
                event_source = %source,
                prefix = %other,
                "Unknown CloudTrail eventSource — using raw prefix as IAM service name. \
                 This may not be a valid IAM service prefix and could produce an unusable policy."
            );
            other
        }
    })
}

/// Maps CloudTrail event names to their actual IAM action names.
/// CloudTrail and IAM use different names for some actions.
fn cloudtrail_to_iam_action(service: &str, event_name: &str) -> String {
    let iam_action = match (service, event_name) {
        // S3 mismatches — CloudTrail uses different names than IAM
        ("s3", "ListBuckets") => "s3:ListAllMyBuckets",
        ("s3", "GetBucketAcl") => "s3:GetBucketAcl",
        ("s3", "HeadBucket") => "s3:ListBucket",
        ("s3", "ListObjects") => "s3:ListBucket",
        ("s3", "ListObjectsV2") => "s3:ListBucket",
        ("s3", "HeadObject") => "s3:GetObject",
        // EC2
        ("ec2", "DescribeInstanceStatus") => "ec2:DescribeInstanceStatus",
        // STS
        ("sts", "GetCallerIdentity") => "sts:GetCallerIdentity",
        // API Gateway — CloudTrail logs management operations but IAM uses HTTP-verb actions
        ("apigateway", "CreateRestApi") => "apigateway:POST",
        ("apigateway", "GetRestApi") => "apigateway:GET",
        ("apigateway", "GetRestApis") => "apigateway:GET",
        ("apigateway", "DeleteRestApi") => "apigateway:DELETE",
        ("apigateway", "UpdateRestApi") => "apigateway:PATCH",
        ("apigateway", "CreateDeployment") => "apigateway:POST",
        ("apigateway", "CreateStage") => "apigateway:POST",
        ("apigateway", "UpdateStage") => "apigateway:PATCH",
        ("apigateway", "DeleteStage") => "apigateway:DELETE",
        ("apigateway", "GetStages") => "apigateway:GET",
        ("apigateway", "GetResources") => "apigateway:GET",
        ("apigateway", "CreateResource") => "apigateway:POST",
        // ECS — some event names differ from IAM actions
        ("ecs", "DescribeClusters") => "ecs:DescribeClusters",
        ("ecs", "DescribeServices") => "ecs:DescribeServices",
        ("ecs", "DescribeTaskDefinition") => "ecs:DescribeTaskDefinition",
        ("ecs", "DescribeTasks") => "ecs:DescribeTasks",
        // CloudWatch Logs — event source is "logs" but some events differ
        ("logs", "CreateLogGroup") => "logs:CreateLogGroup",
        ("logs", "FilterLogEvents") => "logs:FilterLogEvents",
        // Elastic Load Balancing — v2 ALB/NLB operations
        ("elasticloadbalancing", "CreateLoadBalancer") => "elasticloadbalancing:CreateLoadBalancer",
        ("elasticloadbalancing", "DescribeLoadBalancers") => {
            "elasticloadbalancing:DescribeLoadBalancers"
        }
        ("elasticloadbalancing", "CreateTargetGroup") => "elasticloadbalancing:CreateTargetGroup",
        ("elasticloadbalancing", "DescribeTargetGroups") => {
            "elasticloadbalancing:DescribeTargetGroups"
        }
        // Step Functions — event source is "states"
        ("states", "CreateStateMachine") => "states:CreateStateMachine",
        ("states", "StartExecution") => "states:StartExecution",
        ("states", "DescribeExecution") => "states:DescribeExecution",
        // Default: service:eventName is usually correct
        _ => return default_action_for_service(service, event_name),
    };
    iam_action.to_string()
}

/// Provider-specific fallback when the event name isn't in the explicit
/// CloudTrail-to-IAM mapping table. Most AWS services use
/// `service:EventName` directly, but a handful (notably API Gateway and
/// IoT) use HTTP-verb action names that don't match the CloudTrail
/// operation name at all. R6-M6: the old code always returned
/// `service:EventName`, so any un-mapped API Gateway call produced an
/// IAM action like `apigateway:CreateApiKey` that STS rejects
/// (`AccessDenied: apigateway:CreateApiKey is not a valid action`),
/// making the learned policy unusable.
fn default_action_for_service(service: &str, event_name: &str) -> String {
    match service {
        "apigateway" => {
            // API Gateway IAM actions are HTTP verbs — infer from the
            // CloudTrail operation prefix.
            let verb = if event_name.starts_with("Get") || event_name.starts_with("Describe") {
                "GET"
            } else if event_name.starts_with("Create") || event_name.starts_with("Put") {
                "POST"
            } else if event_name.starts_with("Update") {
                "PATCH"
            } else if event_name.starts_with("Delete") {
                "DELETE"
            } else {
                // Unknown verb prefix — widen to `*` rather than emit
                // an invalid action that STS will reject.
                "*"
            };
            format!("apigateway:{}", verb)
        }
        _ => format!("{}:{}", service, event_name),
    }
}

/// Result of learning: observed API calls mapped to IAM actions.
#[derive(Debug, Clone)]
pub struct LearnedPolicy {
    /// Unique IAM actions observed (e.g. "s3:ListBuckets")
    pub actions: BTreeSet<String>,
}

impl LearnedPolicy {
    /// Format as a comma-separated --allow string.
    pub fn to_allow_str(&self) -> String {
        self.actions.iter().cloned().collect::<Vec<_>>().join(",")
    }

    /// Format as a TOML profile block for config.
    pub fn to_profile_toml(&self, name: &str) -> String {
        format!(
            "[profiles.{}]\nallow = \"{}\"\ndescription = \"Learned from running command\"\n",
            name,
            self.to_allow_str()
        )
    }
}

/// Query CloudTrail for API calls made by a specific access key within a time window.
///
/// WARNING: CloudTrail events are regional. This queries only the configured region
/// (or the SDK default). API calls made in other regions will not be discovered.
/// For complete coverage, callers should query multiple regions or use CloudTrail
/// with an organization trail that aggregates events to a single region/S3 bucket.
pub async fn lookup_events(
    access_key_id: &str,
    start_time: chrono::DateTime<chrono::Utc>,
    end_time: chrono::DateTime<chrono::Utc>,
    region: Option<&str>,
) -> Result<LearnedPolicy> {
    let mut loader = aws_config::defaults(aws_config::BehaviorVersion::latest());
    if let Some(r) = region {
        loader = loader.region(aws_config::Region::new(r.to_string()));
    }
    let config = loader.load().await;
    let client = aws_sdk_cloudtrail::Client::new(&config);

    tracing::warn!(
        region = region.unwrap_or("(default)"),
        "CloudTrail learn queries a single region — API calls in other regions will be missed"
    );

    let mut actions = BTreeSet::new();
    let mut next_token: Option<String> = None;

    loop {
        let mut req = client
            .lookup_events()
            .start_time(aws_sdk_cloudtrail::primitives::DateTime::from_secs(
                start_time.timestamp(),
            ))
            .end_time(aws_sdk_cloudtrail::primitives::DateTime::from_secs(
                end_time.timestamp(),
            ))
            .lookup_attributes(
                aws_sdk_cloudtrail::types::LookupAttribute::builder()
                    .attribute_key(aws_sdk_cloudtrail::types::LookupAttributeKey::AccessKeyId)
                    .attribute_value(access_key_id)
                    .build()
                    .map_err(|e| AvError::Sts(format!("CloudTrail lookup build error: {}", e)))?,
            )
            .max_results(50);

        if let Some(ref token) = next_token {
            req = req.next_token(token);
        }

        let resp = req
            .send()
            .await
            .map_err(|e| AvError::Sts(format!("CloudTrail lookup error: {}", e)))?;

        for event in resp.events() {
            let event_name = event.event_name().unwrap_or_default();
            let event_source = event.event_source().unwrap_or_default();

            if let Some(service) = event_source_to_service(event_source) {
                // Skip the AssumeRole call that Audex itself makes
                if service == "sts" && event_name == "AssumeRole" {
                    continue;
                }
                actions.insert(cloudtrail_to_iam_action(service, event_name));
            }
        }

        next_token = resp.next_token().map(|s| s.to_string());
        if next_token.is_none() {
            break;
        }
    }

    Ok(LearnedPolicy { actions })
}

/// Poll CloudTrail until events appear or timeout is reached.
/// CloudTrail typically has a 5-15 minute delay.
///
/// # Limitation — early return on first non-empty result
///
/// This function returns as soon as any events appear. Because CloudTrail
/// propagates events in batches over several minutes, the first non-empty
/// response may be a partial set; later events from the same session window
/// could arrive in subsequent polls.
///
/// For production use where completeness matters, callers should implement a
/// stabilization check — e.g. poll twice more after the first non-empty
/// result and return only when the event count is stable across two
/// consecutive polls.
pub async fn poll_cloudtrail(
    access_key_id: &str,
    start_time: chrono::DateTime<chrono::Utc>,
    end_time: chrono::DateTime<chrono::Utc>,
    region: Option<&str>,
    timeout: Duration,
    poll_interval: Duration,
) -> Result<LearnedPolicy> {
    let deadline = std::time::Instant::now() + timeout;

    loop {
        let result = lookup_events(access_key_id, start_time, end_time, region).await?;

        if !result.actions.is_empty() {
            return Ok(result);
        }

        if std::time::Instant::now() + poll_interval > deadline {
            return Err(AvError::Sts(
                "Timed out waiting for CloudTrail events. CloudTrail can take 5-15 minutes to propagate. \
                 Try again later — CloudTrail events may not be available yet."
                    .to_string(),
            ));
        }

        tokio::time::sleep(poll_interval).await;
    }
}