rag-module 0.6.7

Enterprise RAG module with chat context storage, vector search, session management, and model downloading. Rust implementation with Node.js compatibility.
//! IAMService - Estate resource permission management
//! Integrates with AWS IAM, Azure RBAC, and GCP IAM for unified permission checking

use anyhow::{Result, anyhow};
use serde::{Serialize, Deserialize};
use std::collections::HashMap;
use regex::Regex;

use crate::config::IAMConfig;

/// IAM context structure
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct IAMContext {
    pub user_id: String,
    pub username: Option<String>,
    pub roles: Vec<String>,
    pub permissions: Vec<String>,
    pub account_id: Option<String>,
    pub cloud_provider: String,
}

/// Resource permission check result
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PermissionResult {
    pub resource_id: String,
    pub action: String,
    pub allowed: bool,
    pub reason: Option<String>,
    pub risk_level: String,
}

/// AWS estate processing result
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct EstateProcessingResult {
    pub processed_resources: usize,
    pub document_ids: Vec<String>,
    pub accounts_processed: Vec<String>,
    pub services_found: Vec<String>,
}

/// IAMService implementation
pub struct IAMService {
    config: IAMConfig,
    permission_cache: HashMap<String, (PermissionResult, chrono::DateTime<chrono::Utc>)>,
    resource_patterns: HashMap<String, Regex>,
}

impl IAMService {
    /// Create a new IAMService
    pub async fn new(config: &IAMConfig) -> Result<Self> {
        let mut resource_patterns = HashMap::new();
        
        // AWS ARN pattern
        resource_patterns.insert(
            "aws".to_string(),
            Regex::new(r"^arn:aws:([^:]+):([^:]*):([^:]+):(.+)$")?,
        );
        
        // Azure resource pattern
        resource_patterns.insert(
            "azure".to_string(),
            Regex::new(r"^/subscriptions/([^/]+)/resourceGroups/([^/]+)/providers/([^/]+)/(.+)$")?,
        );
        
        // GCP resource pattern
        resource_patterns.insert(
            "gcp".to_string(),
            Regex::new(r"^projects/([^/]+)/([^/]+)/(.+)$")?,
        );
        
        Ok(Self {
            config: config.clone(),
            permission_cache: HashMap::new(),
            resource_patterns,
        })
    }
    
    /// Check if user has access to a specific resource
    pub async fn check_resource_access(
        &self,
        iam_context: &IAMContext,
        resource_id: &str,
        action: &str,
    ) -> Result<PermissionResult> {
        // Check cache first
        let cache_key = format!("{}:{}:{}", iam_context.user_id, resource_id, action);
        if let Some((cached_result, timestamp)) = self.permission_cache.get(&cache_key) {
            if chrono::Utc::now().signed_duration_since(*timestamp).num_seconds() < 300 {
                return Ok(cached_result.clone());
            }
        }
        
        // Determine cloud provider from resource ID
        let cloud_provider = self.detect_cloud_provider(resource_id)?;
        
        let result = match cloud_provider.as_str() {
            "aws" => self.check_aws_permissions(iam_context, resource_id, action).await?,
            "azure" => self.check_azure_permissions(iam_context, resource_id, action).await?,
            "gcp" => self.check_gcp_permissions(iam_context, resource_id, action).await?,
            _ => PermissionResult {
                resource_id: resource_id.to_string(),
                action: action.to_string(),
                allowed: false,
                reason: Some("Unknown cloud provider".to_string()),
                risk_level: "unknown".to_string(),
            },
        };
        
        Ok(result)
    }
    
    /// Process AWS estate data and create documents
    pub async fn process_estate_data(&self, data: serde_json::Value) -> Result<Vec<String>> {
        let mut document_ids = Vec::new();
        
        // Parse the AWS estate JSON structure
        if let Some(accounts) = data.as_array() {
            for account_data in accounts {
                if let Some(account_obj) = account_data.as_object() {
                    let account_id = account_obj.get("account_id")
                        .and_then(|v| v.as_str())
                        .unwrap_or("unknown");
                    
                    let account_name = account_obj.get("account_name")
                        .and_then(|v| v.as_str())
                        .unwrap_or("Unknown Account");
                    
                    // Process services
                    if let Some(services) = account_obj.get("services").and_then(|s| s.as_object()) {
                        for (service_name, service_data) in services {
                            let doc_id = self.process_service_data(
                                account_id,
                                account_name,
                                service_name,
                                service_data,
                            ).await?;
                            println!("   Created document ID: {}", doc_id);
                            document_ids.push(doc_id);
                        }
                    }
                }
            }
        }
        
        Ok(document_ids)
    }
    
    /// Get IAM context from user data
    pub async fn create_iam_context(
        &self,
        user_id: &str,
        user_data: &serde_json::Value,
    ) -> Result<IAMContext> {
        let username = user_data.get("username")
            .and_then(|u| u.as_str())
            .map(|s| s.to_string());
        
        let account_id = user_data.get("account_id")
            .and_then(|a| a.as_str())
            .map(|s| s.to_string());
        
        // Extract roles and permissions
        let mut roles = Vec::new();
        let mut permissions = Vec::new();
        
        if let Some(services) = user_data.get("services").and_then(|s| s.as_object()) {
            for (service_name, service_data) in services {
                if let Some(user_perms) = service_data.get("user_permissions").and_then(|p| p.as_object()) {
                    roles.push(format!("{}_{}", service_name, user_perms.get("permission_level")
                        .and_then(|l| l.as_str())
                        .unwrap_or("unknown")));
                    
                    // Extract AWS actions as permissions
                    self.extract_permissions_from_service(service_data, &mut permissions);
                }
            }
        }
        
        Ok(IAMContext {
            user_id: user_id.to_string(),
            username,
            roles,
            permissions,
            account_id,
            cloud_provider: "aws".to_string(), // Default to AWS
        })
    }
    
    // Private helper methods
    
    fn detect_cloud_provider(&self, resource_id: &str) -> Result<String> {
        for (provider, pattern) in &self.resource_patterns {
            if pattern.is_match(resource_id) {
                return Ok(provider.clone());
            }
        }
        
        // Fallback detection based on prefixes
        if resource_id.starts_with("arn:aws:") {
            Ok("aws".to_string())
        } else if resource_id.starts_with("/subscriptions/") {
            Ok("azure".to_string())
        } else if resource_id.starts_with("projects/") {
            Ok("gcp".to_string())
        } else {
            Err(anyhow!("Could not detect cloud provider for resource: {}", resource_id))
        }
    }
    
    async fn check_aws_permissions(
        &self,
        iam_context: &IAMContext,
        resource_id: &str,
        action: &str,
    ) -> Result<PermissionResult> {
        // Simplified AWS permission check
        let allowed = match action {
            "read" => iam_context.permissions.iter().any(|p| p.contains("Describe") || p.contains("Get") || p.contains("List")),
            "write" => iam_context.permissions.iter().any(|p| p.contains("Create") || p.contains("Update") || p.contains("Put")),
            "delete" => iam_context.permissions.iter().any(|p| p.contains("Delete") || p.contains("Terminate")),
            "admin" => iam_context.roles.iter().any(|r| r.contains("Full") || r.contains("admin")),
            _ => false,
        };
        
        let risk_level = match action {
            "delete" | "terminate" => "high",
            "write" | "create" => "medium",
            "read" => "low",
            _ => "unknown",
        };
        
        Ok(PermissionResult {
            resource_id: resource_id.to_string(),
            action: action.to_string(),
            allowed,
            reason: if allowed { 
                Some("Permission granted based on IAM analysis".to_string()) 
            } else { 
                Some("Insufficient permissions".to_string()) 
            },
            risk_level: risk_level.to_string(),
        })
    }
    
    async fn check_azure_permissions(
        &self,
        _iam_context: &IAMContext,
        resource_id: &str,
        action: &str,
    ) -> Result<PermissionResult> {
        // Placeholder Azure implementation
        Ok(PermissionResult {
            resource_id: resource_id.to_string(),
            action: action.to_string(),
            allowed: false,
            reason: Some("Azure permissions not implemented".to_string()),
            risk_level: "unknown".to_string(),
        })
    }
    
    async fn check_gcp_permissions(
        &self,
        _iam_context: &IAMContext,
        resource_id: &str,
        action: &str,
    ) -> Result<PermissionResult> {
        // Placeholder GCP implementation
        Ok(PermissionResult {
            resource_id: resource_id.to_string(),
            action: action.to_string(),
            allowed: false,
            reason: Some("GCP permissions not implemented".to_string()),
            risk_level: "unknown".to_string(),
        })
    }
    
    async fn process_service_data(
        &self,
        account_id: &str,
        account_name: &str,
        service_name: &str,
        service_data: &serde_json::Value,
    ) -> Result<String> {
        // Create a document ID
        let doc_id = format!("aws-{}-{}-{}", account_id, service_name, chrono::Utc::now().timestamp());
        
        // Extract service information
        let display_name = service_data
            .get("service_metadata")
            .and_then(|m| m.get("display_name"))
            .and_then(|n| n.as_str())
            .unwrap_or(service_name);
        
        let category = service_data
            .get("service_metadata")
            .and_then(|m| m.get("category"))
            .and_then(|c| c.as_str())
            .unwrap_or("unknown");
        
        // Build content for the document
        let mut content_parts = Vec::new();
        content_parts.push(format!("AWS {} service in account {} ({})", display_name, account_name, account_id));
        
        // Add resource information
        if let Some(resources) = service_data.get("resources").and_then(|r| r.as_array()) {
            content_parts.push(format!("Contains {} resources:", resources.len()));
            for resource in resources.iter().take(5) { // Limit to first 5 for content
                if let Some(resource_id) = resource.get("resource_id").and_then(|r| r.as_str()) {
                    content_parts.push(format!("- {}", resource_id));
                }
            }
        }
        
        // Add permission information
        if let Some(user_perms) = service_data.get("user_permissions") {
            if let Some(perm_level) = user_perms.get("permission_level").and_then(|p| p.as_str()) {
                content_parts.push(format!("User permission level: {}", perm_level));
            }
        }
        
        let content = content_parts.join("\n");
        
        // Create metadata
        let mut metadata = HashMap::new();
        metadata.insert("cloud".to_string(), serde_json::Value::String("aws".to_string()));
        metadata.insert("service".to_string(), serde_json::Value::String(service_name.to_string()));
        metadata.insert("account_id".to_string(), serde_json::Value::String(account_id.to_string()));
        metadata.insert("account_name".to_string(), serde_json::Value::String(account_name.to_string()));
        metadata.insert("category".to_string(), serde_json::Value::String(category.to_string()));
        
        // This would normally create a document in the vector store
        // For now, just return the document ID
        Ok(doc_id)
    }
    
    fn extract_permissions_from_service(
        &self,
        service_data: &serde_json::Value,
        permissions: &mut Vec<String>,
    ) {
        if let Some(user_perms) = service_data.get("user_permissions").and_then(|p| p.as_object()) {
            for (action_name, action_data) in user_perms {
                if action_name == "permission_level" {
                    continue;
                }
                
                if let Some(aws_actions) = action_data.get("aws_actions").and_then(|a| a.as_array()) {
                    for action in aws_actions {
                        if let Some(action_str) = action.as_str() {
                            permissions.push(action_str.to_string());
                        }
                    }
                }
            }
        }
    }
    
    /// Initialize the service
    pub async fn initialize(&self) -> Result<()> {
        Ok(())
    }
    
    /// Shutdown the service
    pub async fn shutdown(&self) -> Result<()> {
        Ok(())
    }
}