tryaudex_core/

session.rs

use std::path::PathBuf;
use std::time::Duration;

use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
use uuid::Uuid;

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

/// Cloud provider for a session.
#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, Default)]
#[serde(rename_all = "lowercase")]
pub enum CloudProvider {
    #[default]
    Aws,
    Gcp,
    Azure,
}

impl std::fmt::Display for CloudProvider {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::Aws => write!(f, "aws"),
            Self::Gcp => write!(f, "gcp"),
            Self::Azure => write!(f, "azure"),
        }
    }
}

/// Represents an active agent credential session.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Session {
    pub id: String,
    pub created_at: DateTime<Utc>,
    pub expires_at: DateTime<Utc>,
    pub ttl_seconds: u64,
    pub budget: Option<f64>,
    pub policy: ScopedPolicy,
    pub status: SessionStatus,
    /// Provider-agnostic principal identifier despite the AWS-centric field name.
    ///
    /// - **AWS**: IAM role ARN (e.g. `arn:aws:iam::123456789012:role/MyRole`)
    /// - **GCP**: Service account email (e.g. `my-sa@project.iam.gserviceaccount.com`)
    /// - **Azure**: Subscription ID or principal (e.g. `00000000-0000-0000-0000-000000000000`)
    ///
    /// The field is named `role_arn` for historical reasons and backwards
    /// serialization compatibility. New code should use the [`Session::principal_id`]
    /// helper method instead of accessing this field directly.
    /// JSON deserialization also accepts the alias `"role_principal"` to ease
    /// migration in tooling that wants provider-neutral field names.
    #[serde(alias = "role_principal")]
    pub role_arn: String,
    pub command: Vec<String>,
    /// AWS access key for this session (populated after credential creation).
    pub access_key_id: Option<String>,
    /// Cloud provider (defaults to AWS for backwards compat).
    #[serde(default)]
    pub provider: CloudProvider,
    /// Agent identity (from AUDEX_AGENT_ID env var) — tracks which AI agent/model issued the request.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub agent_id: Option<String>,
    /// Azure principal ID (client ID or user object ID) to disambiguate sessions
    /// on the same subscription but different identities.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub principal_id: Option<String>,
    /// User-supplied key:value tags for audit filtering and STS session tags.
    #[serde(default, skip_serializing_if = "std::collections::HashMap::is_empty")]
    pub tags: std::collections::HashMap<String, String>,
}

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
#[serde(rename_all = "snake_case")]
pub enum SessionStatus {
    Active,
    Completed,
    Expired,
    Revoked,
    BudgetExceeded,
    Failed,
}

impl std::fmt::Display for SessionStatus {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::Active => write!(f, "active"),
            Self::Completed => write!(f, "completed"),
            Self::Expired => write!(f, "expired"),
            Self::Revoked => write!(f, "revoked"),
            Self::BudgetExceeded => write!(f, "budget_exceeded"),
            Self::Failed => write!(f, "failed"),
        }
    }
}

impl Session {
    pub fn new(
        ttl: Duration,
        budget: Option<f64>,
        policy: ScopedPolicy,
        role_arn: String,
        command: Vec<String>,
    ) -> Self {
        let now = Utc::now();
        let expires_at = now + chrono::Duration::seconds(ttl.as_secs() as i64);
        Self {
            id: Uuid::new_v4().to_string(),
            created_at: now,
            expires_at,
            ttl_seconds: ttl.as_secs(),
            budget,
            policy,
            status: SessionStatus::Active,
            role_arn,
            command,
            access_key_id: None,
            provider: CloudProvider::default(),
            agent_id: None,
            principal_id: None,
            tags: std::collections::HashMap::new(),
        }
    }

    pub fn new_gcp(
        ttl: Duration,
        budget: Option<f64>,
        policy: ScopedPolicy,
        service_account: String,
        command: Vec<String>,
    ) -> Self {
        let now = Utc::now();
        let expires_at = now + chrono::Duration::seconds(ttl.as_secs() as i64);
        Self {
            id: Uuid::new_v4().to_string(),
            created_at: now,
            expires_at,
            ttl_seconds: ttl.as_secs(),
            budget,
            policy,
            status: SessionStatus::Active,
            role_arn: service_account,
            command,
            access_key_id: None,
            provider: CloudProvider::Gcp,
            agent_id: None,
            principal_id: None,
            tags: std::collections::HashMap::new(),
        }
    }

    pub fn new_azure(
        ttl: Duration,
        budget: Option<f64>,
        policy: ScopedPolicy,
        subscription_id: String,
        command: Vec<String>,
    ) -> Self {
        let now = Utc::now();
        let expires_at = now + chrono::Duration::seconds(ttl.as_secs() as i64);
        Self {
            id: Uuid::new_v4().to_string(),
            created_at: now,
            expires_at,
            ttl_seconds: ttl.as_secs(),
            budget,
            policy,
            status: SessionStatus::Active,
            role_arn: subscription_id,
            command,
            access_key_id: None,
            provider: CloudProvider::Azure,
            agent_id: None,
            principal_id: None,
            tags: std::collections::HashMap::new(),
        }
    }

    /// Return the first 8 characters of the session ID for display, or the
    /// full ID if it is shorter than 8 characters. Avoids panics from direct
    /// `session.id[..8]` slicing on corrupt/hand-edited session files.
    pub fn short_id(&self) -> &str {
        self.id.get(..8).unwrap_or(&self.id)
    }

    pub fn is_expired(&self) -> bool {
        Utc::now() > self.expires_at
    }

    pub fn remaining_seconds(&self) -> i64 {
        (self.expires_at - Utc::now()).num_seconds().max(0)
    }

    pub fn complete(&mut self) {
        self.status = SessionStatus::Completed;
    }

    pub fn expire(&mut self) {
        self.status = SessionStatus::Expired;
    }

    pub fn revoke(&mut self) {
        self.status = SessionStatus::Revoked;
    }

    pub fn fail(&mut self) {
        self.status = SessionStatus::Failed;
    }

    /// Returns the provider-agnostic principal identifier for this session.
    ///
    /// This is a migration-path accessor for the `role_arn` field, which stores
    /// different things depending on the cloud provider (IAM role ARN for AWS,
    /// service account email for GCP, subscription/principal ID for Azure).
    /// When the optional `principal_id` field is populated (e.g. for Azure
    /// service-principal disambiguation), that value is returned instead.
    /// Prefer this method over accessing `role_arn` directly in new code.
    pub fn principal_id(&self) -> &str {
        self.principal_id.as_deref().unwrap_or(&self.role_arn)
    }
}

/// Manages session persistence on disk.
pub struct SessionStore {
    dir: PathBuf,
}

impl SessionStore {
    pub fn new() -> Result<Self> {
        let base = dirs::data_local_dir().ok_or_else(|| {
            AvError::InvalidPolicy(
                "Could not determine local data directory. Set XDG_DATA_HOME or HOME.".to_string(),
            )
        })?;
        let dir = base.join("audex").join("sessions");
        std::fs::create_dir_all(&dir)?;
        Ok(Self { dir })
    }

    /// Maximum number of session files retained on disk. When a new save
    /// would push the directory past this cap, the oldest files (by mtime)
    /// are evicted one-at-a-time to stay exactly at `SESSION_MAX` and
    /// prevent unbounded disk growth (R3-L12, R6-H22).
    const SESSION_MAX: usize = 1000;

    fn write_file(&self, session: &Session) -> Result<()> {
        let path = self.dir.join(format!("{}.json", session.id));
        let json = serde_json::to_string_pretty(session)?;
        #[cfg(unix)]
        {
            use std::os::unix::fs::OpenOptionsExt;
            let mut file = std::fs::OpenOptions::new()
                .write(true)
                .create(true)
                .truncate(true)
                .mode(0o600)
                .open(&path)?;
            std::io::Write::write_all(&mut file, json.as_bytes())?;
        }
        #[cfg(not(unix))]
        {
            std::fs::write(&path, json)?;
        }
        Ok(())
    }

    pub fn save(&self, session: &Session) -> Result<()> {
        // R6-H22: previously the prune check fired at `entries.len() >=
        // SESSION_MAX` and then blindly deleted back to `PRUNE_TARGET - 1`
        // (899), which meant every save above 1000 evicted 101 files —
        // even in the common case where we were only one over the limit.
        // Combined with `update()` aliasing `save()`, a single session
        // update could churn 101 unrelated sessions out of existence.
        //
        // Only prune when the new file would actually push us over the
        // cap, and only delete the excess down to SESSION_MAX so steady-
        // state writes evict exactly one file per new session.  Also
        // skip eviction when the session file already exists (update
        // path): updating an existing record must never delete other
        // sessions.
        let path = self.dir.join(format!("{}.json", session.id));
        let is_new = !path.exists();

        if is_new {
            let entries: Result<Vec<_>> = std::fs::read_dir(&self.dir)?
                .filter_map(|e| e.ok())
                .filter(|e| {
                    e.path()
                        .extension()
                        .and_then(|x| x.to_str())
                        .map(|x| x == "json")
                        .unwrap_or(false)
                })
                .map(|e| {
                    let mtime = e
                        .metadata()
                        .and_then(|m| m.modified())
                        .unwrap_or(std::time::SystemTime::UNIX_EPOCH);
                    Ok((mtime, e.path()))
                })
                .collect();
            let mut entries = entries?;
            // Evict just enough to fit the new record under the cap.
            if entries.len() + 1 > Self::SESSION_MAX {
                entries.sort_unstable_by_key(|(mtime, _)| *mtime);
                let to_delete = entries.len() + 1 - Self::SESSION_MAX;
                for (_, p) in entries.iter().take(to_delete) {
                    let _ = std::fs::remove_file(p);
                }
            }
        }

        self.write_file(session)
    }

    pub fn load(&self, id: &str) -> Result<Session> {
        crate::validate::session_id(id)?;
        let path = self.dir.join(format!("{}.json", id));
        if !path.exists() {
            return Err(AvError::SessionNotFound { id: id.to_string() });
        }
        let json = std::fs::read_to_string(path)?;
        let session: Session = serde_json::from_str(&json)?;
        Ok(session)
    }

    pub fn list(&self) -> Result<Vec<Session>> {
        let mut sessions = Vec::new();
        for entry in std::fs::read_dir(&self.dir)? {
            let entry = entry?;
            let path = entry.path();
            if path.extension().is_some_and(|ext| ext == "json") {
                let json = std::fs::read_to_string(&path)?;
                if let Ok(session) = serde_json::from_str::<Session>(&json) {
                    sessions.push(session);
                }
            }
        }
        sessions.sort_by(|a, b| b.created_at.cmp(&a.created_at));
        Ok(sessions)
    }

    pub fn update(&self, session: &Session) -> Result<()> {
        self.save(session)
    }

    /// Find an active session with matching allow actions, role ARN, and enough remaining TTL.
    /// Used for credential caching — reuse existing sessions instead of issuing new ones.
    pub fn find_reusable(
        &self,
        allow_str: &str,
        role_arn: &str,
        min_remaining_secs: i64,
    ) -> Result<Option<Session>> {
        self.find_reusable_for_provider(
            allow_str,
            role_arn,
            min_remaining_secs,
            CloudProvider::Aws,
            &[],
        )
    }

    /// Provider-aware session reuse lookup. Formats stored actions using the
    /// correct delimiter for the provider so GCP/Azure sessions can be matched.
    /// For Azure, `principal_id` disambiguates sessions on the same subscription
    /// but different identities (service principals, users).
    pub fn find_reusable_for_provider(
        &self,
        allow_str: &str,
        role_arn: &str,
        min_remaining_secs: i64,
        provider: CloudProvider,
        resources: &[String],
    ) -> Result<Option<Session>> {
        let sessions = self.list()?;
        // R6-H5: sort + dedup so semantically equivalent inputs
        // (reordered or repeated entries) still match the same session.
        let mut allow_sorted: Vec<&str> = allow_str.split(',').map(|s| s.trim()).collect();
        allow_sorted.sort();
        allow_sorted.dedup();

        let mut requested_resources: Vec<&str> = resources.iter().map(|s| s.as_str()).collect();
        requested_resources.sort();
        requested_resources.dedup();

        for session in sessions {
            if session.status != SessionStatus::Active {
                continue;
            }
            if session.is_expired() {
                continue;
            }
            if session.remaining_seconds() < min_remaining_secs {
                continue;
            }
            if session.role_arn != role_arn {
                continue;
            }
            // For Azure, also require matching principal to avoid reusing
            // credentials from a different identity on the same subscription.
            // Only enforced when AZURE_CLIENT_ID is set (service principal flow).
            // Users via `az login` don't have this env var, so we skip the
            // check and rely on subscription (role_arn) matching alone.
            if provider == CloudProvider::Azure {
                if let Some(ref stored_pid) = session.principal_id {
                    if let Ok(current_pid) = std::env::var("AZURE_CLIENT_ID") {
                        if current_pid != *stored_pid {
                            continue;
                        }
                    }
                }
                // R6-H11: match on tenant ID when available. A subscription
                // migrated between tenants (e.g. CSP transfer) keeps the same
                // subscription GUID but changes the token audience; reusing
                // the old token against the new tenant leaks directory data.
                let current_tenant = std::env::var("AZURE_TENANT_ID").ok();
                let stored_tenant = session.tags.get("azure:tenant_id").cloned();
                if stored_tenant != current_tenant {
                    continue;
                }
            }

            // Compare actions using provider-appropriate format
            let mut session_actions: Vec<String> = session
                .policy
                .actions
                .iter()
                .map(|a| match provider {
                    CloudProvider::Gcp => a.to_gcp_permission(),
                    CloudProvider::Azure => a.to_azure_permission(),
                    CloudProvider::Aws => a.to_iam_action(),
                })
                .collect();
            session_actions.sort();
            session_actions.dedup();
            let session_actions_str: Vec<&str> =
                session_actions.iter().map(|s| s.as_str()).collect();

            if session_actions_str != allow_sorted {
                continue;
            }

            // Also require matching resource restrictions so a session scoped to
            // a specific ARN is never handed out for a request that expects a
            // different (or wider) set of resources.
            let mut session_resources: Vec<&str> = session
                .policy
                .resources
                .iter()
                .map(|s| s.as_str())
                .collect();
            session_resources.sort();
            session_resources.dedup();
            if session_resources != requested_resources {
                continue;
            }

            return Ok(Some(session));
        }
        Ok(None)
    }
}