lakehouse/crates/gateway/src/access.rs

/// Access control: field-level sensitivity enforcement, column masking, query audit.
/// Evaluates policies at query time — not at storage level.

use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
use shared::types::Sensitivity;
use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::RwLock;

/// An agent's role determines what sensitivity levels they can see.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AgentRole {
    pub agent_name: String,
    pub role: String,                        // "recruiter", "admin", "analyst", "agent"
    pub allowed_sensitivity: Vec<Sensitivity>, // what they can see unmasked
}

/// A query audit entry.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct QueryAudit {
    pub id: String,
    pub agent: String,
    pub sql: String,
    pub datasets_accessed: Vec<String>,
    pub pii_fields_accessed: Vec<String>,
    pub timestamp: DateTime<Utc>,
    pub row_count: usize,
    pub allowed: bool,
    pub masked_fields: Vec<String>,
}

/// Access control manager.
#[derive(Clone)]
pub struct AccessControl {
    roles: Arc<RwLock<HashMap<String, AgentRole>>>,
    audit_log: Arc<RwLock<Vec<QueryAudit>>>,
    enabled: bool,
}

impl AccessControl {
    pub fn new(enabled: bool) -> Self {
        let ac = Self {
            roles: Arc::new(RwLock::new(HashMap::new())),
            audit_log: Arc::new(RwLock::new(Vec::new())),
            enabled,
        };
        ac
    }

    /// Register default roles.
    pub async fn register_defaults(&self) {
        let defaults = vec![
            AgentRole {
                agent_name: "admin".into(),
                role: "admin".into(),
                allowed_sensitivity: vec![
                    Sensitivity::Public, Sensitivity::Internal,
                    Sensitivity::Pii, Sensitivity::Phi, Sensitivity::Financial,
                ],
            },
            AgentRole {
                agent_name: "recruiter".into(),
                role: "recruiter".into(),
                allowed_sensitivity: vec![
                    Sensitivity::Public, Sensitivity::Internal, Sensitivity::Pii,
                ],
            },
            AgentRole {
                agent_name: "analyst".into(),
                role: "analyst".into(),
                allowed_sensitivity: vec![
                    Sensitivity::Public, Sensitivity::Internal, Sensitivity::Financial,
                ],
            },
            AgentRole {
                agent_name: "agent".into(),
                role: "agent".into(),
                allowed_sensitivity: vec![
                    Sensitivity::Public, Sensitivity::Internal,
                ],
            },
        ];

        let mut roles = self.roles.write().await;
        for role in defaults {
            roles.insert(role.agent_name.clone(), role);
        }
    }

    /// Register or update an agent role.
    pub async fn set_role(&self, role: AgentRole) {
        self.roles.write().await.insert(role.agent_name.clone(), role);
    }

    /// Get an agent's role.
    pub async fn get_role(&self, agent: &str) -> Option<AgentRole> {
        self.roles.read().await.get(agent).cloned()
    }

    /// List all roles.
    pub async fn list_roles(&self) -> Vec<AgentRole> {
        self.roles.read().await.values().cloned().collect()
    }

    /// Check if an agent can see a field with given sensitivity.
    pub async fn can_access(&self, agent: &str, sensitivity: &Sensitivity) -> bool {
        if !self.enabled { return true; }
        match self.roles.read().await.get(agent) {
            Some(role) => role.allowed_sensitivity.contains(sensitivity),
            None => false, // unknown agent = no access
        }
    }

    /// Determine which fields should be masked for an agent.
    pub async fn masked_fields(
        &self,
        agent: &str,
        columns: &[shared::types::ColumnMeta],
    ) -> Vec<String> {
        if !self.enabled { return vec![]; }

        let role = match self.roles.read().await.get(agent) {
            Some(r) => r.clone(),
            None => return columns.iter().filter(|c| c.sensitivity.is_some()).map(|c| c.name.clone()).collect(),
        };

        columns.iter()
            .filter(|col| {
                if let Some(ref sens) = col.sensitivity {
                    !role.allowed_sensitivity.contains(sens)
                } else {
                    false
                }
            })
            .map(|col| col.name.clone())
            .collect()
    }

    /// Log a query for audit.
    pub async fn log_query(&self, audit: QueryAudit) {
        self.audit_log.write().await.push(audit);
    }

    /// Get recent audit entries.
    pub async fn recent_audit(&self, limit: usize) -> Vec<QueryAudit> {
        let log = self.audit_log.read().await;
        let start = log.len().saturating_sub(limit);
        log[start..].iter().rev().cloned().collect()
    }

    pub fn is_enabled(&self) -> bool {
        self.enabled
    }
}