tryaudex_core/

estimate.rs

use std::collections::HashMap;

use serde::{Deserialize, Serialize};

use crate::policy::ScopedPolicy;

/// Cost estimate for a set of IAM actions.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CostEstimate {
    /// Per-service cost breakdown.
    pub services: Vec<ServiceEstimate>,
    /// Total estimated cost for the session.
    pub total_min: f64,
    pub total_max: f64,
    /// Risk level: "low", "medium", "high".
    pub risk_level: String,
    /// Human-readable warnings.
    pub warnings: Vec<String>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ServiceEstimate {
    pub service: String,
    pub actions: Vec<String>,
    pub min_cost: f64,
    pub max_cost: f64,
    pub notes: String,
}

/// Known cost profiles for common AWS services and actions.
/// These are rough estimates based on typical usage patterns.
/// Real costs depend on usage volume, region, and data transfer.
fn cost_profiles() -> HashMap<&'static str, ServiceCostProfile> {
    HashMap::from([
        ("s3", ServiceCostProfile {
            read_actions: &["GetObject", "ListBucket", "ListAllMyBuckets", "GetBucketLocation", "HeadObject"],
            write_actions: &["PutObject", "DeleteObject", "CreateBucket", "CopyObject"],
            read_cost_per_1k: 0.0004,   // $0.0004 per 1,000 GET requests
            write_cost_per_1k: 0.005,    // $0.005 per 1,000 PUT requests
            data_cost_per_gb: 0.023,     // $0.023/GB stored
            typical_requests: 100,
            notes: "S3 pricing: GET $0.0004/1K, PUT $0.005/1K, storage $0.023/GB/mo",
        }),
        ("lambda", ServiceCostProfile {
            read_actions: &["GetFunction", "ListFunctions", "GetFunctionConfiguration"],
            write_actions: &["InvokeFunction", "UpdateFunctionCode", "UpdateFunctionConfiguration", "CreateFunction"],
            read_cost_per_1k: 0.0,
            write_cost_per_1k: 0.20,     // ~$0.20 per 1M invocations = $0.0002/1K
            data_cost_per_gb: 0.0000167, // $0.0000166667/GB-second
            typical_requests: 10,
            notes: "Lambda: $0.20/1M requests + $0.0000166667/GB-s compute",
        }),
        ("dynamodb", ServiceCostProfile {
            read_actions: &["GetItem", "Query", "Scan", "BatchGetItem", "DescribeTable", "ListTables"],
            write_actions: &["PutItem", "UpdateItem", "DeleteItem", "BatchWriteItem"],
            read_cost_per_1k: 0.00025,   // $0.25 per 1M RRU
            write_cost_per_1k: 0.00125,  // $1.25 per 1M WRU
            data_cost_per_gb: 0.25,
            typical_requests: 100,
            notes: "DynamoDB on-demand: $0.25/1M RRU, $1.25/1M WRU",
        }),
        ("ec2", ServiceCostProfile {
            read_actions: &["DescribeInstances", "DescribeSecurityGroups", "DescribeSubnets", "DescribeVpcs", "DescribeImages"],
            write_actions: &["RunInstances", "TerminateInstances", "StartInstances", "StopInstances", "CreateSecurityGroup"],
            read_cost_per_1k: 0.0,
            write_cost_per_1k: 0.0,      // EC2 API calls are free, instances cost
            data_cost_per_gb: 0.0,
            typical_requests: 10,
            notes: "EC2 API calls are free. Instance costs: $0.0116/hr (t3.micro) to $3.84/hr (p3.16xlarge)",
        }),
        ("iam", ServiceCostProfile {
            read_actions: &["GetRole", "GetPolicy", "GetPolicyVersion", "ListRolePolicies", "ListAttachedRolePolicies"],
            write_actions: &["CreateRole", "DeleteRole", "AttachRolePolicy", "DetachRolePolicy", "PutRolePolicy", "DeleteRolePolicy", "PassRole"],
            read_cost_per_1k: 0.0,
            write_cost_per_1k: 0.0,
            data_cost_per_gb: 0.0,
            typical_requests: 10,
            notes: "IAM API calls are free. WARNING: IAM changes affect account security",
        }),
        ("sqs", ServiceCostProfile {
            read_actions: &["ReceiveMessage", "GetQueueAttributes", "ListQueues"],
            write_actions: &["SendMessage", "DeleteMessage", "CreateQueue"],
            read_cost_per_1k: 0.0004,
            write_cost_per_1k: 0.0004,
            data_cost_per_gb: 0.0,
            typical_requests: 100,
            notes: "SQS: $0.40/1M requests (first 1M free)",
        }),
        ("sns", ServiceCostProfile {
            read_actions: &["ListTopics", "GetTopicAttributes"],
            write_actions: &["Publish"],
            read_cost_per_1k: 0.0,
            write_cost_per_1k: 0.0005,
            data_cost_per_gb: 0.0,
            typical_requests: 10,
            notes: "SNS: $0.50/1M publishes",
        }),
        ("ecr", ServiceCostProfile {
            read_actions: &["GetAuthorizationToken", "BatchCheckLayerAvailability", "GetDownloadUrlForLayer", "BatchGetImage", "DescribeRepositories"],
            write_actions: &["PutImage", "InitiateLayerUpload", "UploadLayerPart", "CompleteLayerUpload", "CreateRepository"],
            read_cost_per_1k: 0.0,
            write_cost_per_1k: 0.0,
            data_cost_per_gb: 0.10,
            typical_requests: 10,
            notes: "ECR: $0.10/GB/month storage. Data transfer charges apply",
        }),
        ("logs", ServiceCostProfile {
            read_actions: &["GetLogEvents", "DescribeLogGroups", "DescribeLogStreams", "FilterLogEvents"],
            write_actions: &["PutLogEvents", "CreateLogGroup", "CreateLogStream"],
            read_cost_per_1k: 0.005,
            write_cost_per_1k: 0.0,
            data_cost_per_gb: 0.50,
            typical_requests: 50,
            notes: "CloudWatch Logs: $0.50/GB ingested, $0.005/1K queries",
        }),
        ("sts", ServiceCostProfile {
            read_actions: &["GetCallerIdentity"],
            write_actions: &["AssumeRole"],
            read_cost_per_1k: 0.0,
            write_cost_per_1k: 0.0,
            data_cost_per_gb: 0.0,
            typical_requests: 1,
            notes: "STS API calls are free",
        }),
        ("cloudformation", ServiceCostProfile {
            read_actions: &["DescribeStacks", "ListStacks", "GetTemplate"],
            write_actions: &["CreateStack", "UpdateStack", "DeleteStack"],
            read_cost_per_1k: 0.0,
            write_cost_per_1k: 0.0,
            data_cost_per_gb: 0.0,
            typical_requests: 10,
            notes: "CloudFormation: free for AWS resources, $0.0009/handler operation for third-party",
        }),
    ])
}

struct ServiceCostProfile {
    read_actions: &'static [&'static str],
    write_actions: &'static [&'static str],
    read_cost_per_1k: f64,
    write_cost_per_1k: f64,
    data_cost_per_gb: f64,
    typical_requests: u32,
    notes: &'static str,
}

/// Estimate the potential cost of running a session with the given policy.
pub fn estimate(policy: &ScopedPolicy, ttl_seconds: u64) -> CostEstimate {
    let profiles = cost_profiles();
    let mut services = Vec::new();
    let mut warnings = Vec::new();
    let mut total_min = 0.0;
    let mut total_max = 0.0;

    // Group actions by service
    let mut by_service: HashMap<String, Vec<String>> = HashMap::new();
    for action in &policy.actions {
        by_service
            .entry(action.service.clone())
            .or_default()
            .push(action.action.clone());
    }

    let ttl_hours = ttl_seconds as f64 / 3600.0;

    for (service, actions) in &by_service {
        if let Some(profile) = profiles.get(service.as_str()) {
            let mut read_count = 0;
            let mut write_count = 0;

            for action in actions {
                if action == "*" {
                    // Wildcard — assume mix of read + write
                    read_count += profile.typical_requests;
                    write_count += profile.typical_requests;
                    warnings.push(format!(
                        "{}:* grants all actions — cost depends on actual usage",
                        service
                    ));
                } else if profile
                    .read_actions
                    .iter()
                    .any(|r| action.starts_with(r) || *r == action.as_str())
                {
                    read_count += profile.typical_requests;
                } else if profile
                    .write_actions
                    .iter()
                    .any(|w| action.starts_with(w) || *w == action.as_str())
                {
                    write_count += profile.typical_requests;
                } else {
                    // Unknown action, assume write
                    write_count += profile.typical_requests / 2;
                }
            }

            // Scale by TTL (longer sessions = more potential requests)
            let scale = (ttl_hours * 0.5).max(1.0); // at least 1x
            let scaled_reads = (read_count as f64 * scale) as u32;
            let scaled_writes = (write_count as f64 * scale) as u32;

            let read_cost = (scaled_reads as f64 / 1000.0) * profile.read_cost_per_1k;
            let write_cost = (scaled_writes as f64 / 1000.0) * profile.write_cost_per_1k;
            let min_cost = read_cost + write_cost;
            // Max assumes 10x typical usage
            let max_cost = min_cost * 10.0 + profile.data_cost_per_gb * 0.1;

            total_min += min_cost;
            total_max += max_cost;

            services.push(ServiceEstimate {
                service: service.clone(),
                actions: actions.clone(),
                min_cost,
                max_cost,
                notes: profile.notes.to_string(),
            });

            // High-cost warnings
            if service == "ec2"
                && actions
                    .iter()
                    .any(|a| a.contains("RunInstances") || a == "*")
            {
                warnings.push(
                    "ec2:RunInstances can launch instances costing up to $3.84/hr".to_string(),
                );
            }
        } else {
            services.push(ServiceEstimate {
                service: service.clone(),
                actions: actions.clone(),
                min_cost: 0.0,
                max_cost: 0.0,
                notes: "No cost estimate available for this service".to_string(),
            });
        }
    }

    // Determine risk level
    let risk_level = if total_max > 10.0 || !warnings.is_empty() {
        "high"
    } else if total_max > 1.0 {
        "medium"
    } else {
        "low"
    }
    .to_string();

    // IAM-specific warnings
    if by_service.contains_key("iam") {
        warnings.push("IAM actions can modify account security — review carefully".to_string());
    }

    services.sort_by(|a, b| {
        b.max_cost
            .partial_cmp(&a.max_cost)
            .unwrap_or(std::cmp::Ordering::Equal)
    });

    CostEstimate {
        services,
        total_min,
        total_max,
        risk_level,
        warnings,
    }
}

/// Format a cost estimate as human-readable text.
pub fn format_text(est: &CostEstimate) -> String {
    let mut out = String::new();

    out.push_str(&format!(
        "Estimated cost: ${:.4} — ${:.4}\n",
        est.total_min, est.total_max
    ));
    out.push_str(&format!("Risk level: {}\n\n", est.risk_level));

    for svc in &est.services {
        out.push_str(&format!(
            "  {} (${:.4} — ${:.4})\n",
            svc.service, svc.min_cost, svc.max_cost
        ));
        out.push_str(&format!("    Actions: {}\n", svc.actions.join(", ")));
        out.push_str(&format!("    {}\n", svc.notes));
    }

    if !est.warnings.is_empty() {
        out.push_str("\nWarnings:\n");
        for w in &est.warnings {
            out.push_str(&format!("  ! {}\n", w));
        }
    }

    out
}

/// Format a cost estimate as JSON.
pub fn format_json(est: &CostEstimate) -> String {
    serde_json::to_string_pretty(est).unwrap_or_else(|_| "{}".to_string())
}

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

    #[test]
    fn test_estimate_s3_readonly() {
        let policy = ScopedPolicy::from_allow_str("s3:GetObject,s3:ListBucket").unwrap();
        let est = estimate(&policy, 900);
        assert_eq!(est.services.len(), 1);
        assert_eq!(est.services[0].service, "s3");
        assert!(est.total_min >= 0.0);
        assert_eq!(est.risk_level, "low");
    }

    #[test]
    fn test_estimate_ec2_high_risk() {
        let policy =
            ScopedPolicy::from_allow_str("ec2:RunInstances,ec2:DescribeInstances").unwrap();
        let est = estimate(&policy, 3600);
        assert_eq!(est.risk_level, "high");
        assert!(est.warnings.iter().any(|w| w.contains("RunInstances")));
    }

    #[test]
    fn test_estimate_iam_warning() {
        let policy = ScopedPolicy::from_allow_str("iam:CreateRole").unwrap();
        let est = estimate(&policy, 900);
        assert!(est.warnings.iter().any(|w| w.contains("IAM")));
    }

    #[test]
    fn test_estimate_wildcard() {
        let policy = ScopedPolicy::from_allow_str("s3:*").unwrap();
        let est = estimate(&policy, 900);
        assert!(est.warnings.iter().any(|w| w.contains("s3:*")));
    }

    #[test]
    fn test_estimate_multi_service() {
        let policy =
            ScopedPolicy::from_allow_str("s3:GetObject,lambda:InvokeFunction,dynamodb:Query")
                .unwrap();
        let est = estimate(&policy, 900);
        assert_eq!(est.services.len(), 3);
    }

    #[test]
    fn test_estimate_unknown_service() {
        let policy = ScopedPolicy::from_allow_str("xray:GetTraceSummaries").unwrap();
        let est = estimate(&policy, 900);
        assert_eq!(est.services.len(), 1);
        assert!(est.services[0].notes.contains("No cost estimate"));
    }

    #[test]
    fn test_format_text() {
        let policy = ScopedPolicy::from_allow_str("s3:GetObject").unwrap();
        let est = estimate(&policy, 900);
        let text = format_text(&est);
        assert!(text.contains("Estimated cost"));
        assert!(text.contains("s3"));
    }

    #[test]
    fn test_format_json() {
        let policy = ScopedPolicy::from_allow_str("s3:GetObject").unwrap();
        let est = estimate(&policy, 900);
        let json = format_json(&est);
        assert!(json.contains("total_min"));
        assert!(json.contains("risk_level"));
    }

    #[test]
    fn test_longer_ttl_scales_cost() {
        let policy = ScopedPolicy::from_allow_str("s3:GetObject,s3:PutObject").unwrap();
        let short = estimate(&policy, 900); // 15 min
        let long = estimate(&policy, 14400); // 4 hours
                                             // Longer TTL should result in higher max cost
        assert!(long.total_max >= short.total_max);
    }

    #[test]
    fn test_free_services() {
        let policy = ScopedPolicy::from_allow_str("sts:GetCallerIdentity").unwrap();
        let est = estimate(&policy, 900);
        assert_eq!(est.total_min, 0.0);
    }
}