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, Serialize, Deserialize, PartialEq, 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,
    /// AWS: IAM role ARN. GCP: service account email. Azure: subscription ID or 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>,
}

#[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,
        }
    }

    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,
        }
    }

    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,
        }
    }

    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;
    }
}

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

impl SessionStore {
    pub fn new() -> Result<Self> {
        let dir = dirs::data_local_dir()
            .unwrap_or_else(|| PathBuf::from("."))
            .join("audex")
            .join("sessions");
        std::fs::create_dir_all(&dir)?;
        Ok(Self { dir })
    }

    pub fn save(&self, session: &Session) -> Result<()> {
        let path = self.dir.join(format!("{}.json", session.id));
        let json = serde_json::to_string_pretty(session)?;
        std::fs::write(path, json)?;
        Ok(())
    }

    pub fn load(&self, id: &str) -> Result<Session> {
        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>> {
        let sessions = self.list()?;
        let mut allow_sorted: Vec<&str> = allow_str.split(',').map(|s| s.trim()).collect();
        allow_sorted.sort();

        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;
            }

            // Compare actions
            let mut session_actions: Vec<String> = session
                .policy
                .actions
                .iter()
                .map(|a| a.to_iam_action())
                .collect();
            session_actions.sort();
            let session_actions_str: Vec<&str> =
                session_actions.iter().map(|s| s.as_str()).collect();

            if session_actions_str == allow_sorted {
                return Ok(Some(session));
            }
        }
        Ok(None)
    }
}