tryaudex_core/

audit.rs

use std::path::PathBuf;

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

use crate::error::Result;
use crate::session::{CloudProvider, Session};

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AuditEntry {
    pub timestamp: DateTime<Utc>,
    pub session_id: String,
    /// Cloud provider for this entry (defaults to "aws" for backwards compat).
    #[serde(default = "default_provider")]
    pub provider: String,
    pub event: AuditEvent,
}

fn default_provider() -> String {
    "aws".to_string()
}

#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "snake_case")]
pub enum AuditEvent {
    SessionCreated {
        /// "Role ARN" for AWS, "Service Account" for GCP, "Subscription" for Azure
        role_arn: String,
        ttl_seconds: u64,
        budget: Option<f64>,
        allowed_actions: Vec<String>,
        command: Vec<String>,
        /// AI agent/model identity (from AUDEX_AGENT_ID env var).
        #[serde(default, skip_serializing_if = "Option::is_none")]
        agent_id: Option<String>,
    },
    CredentialsIssued {
        access_key_id: String,
        expires_at: DateTime<Utc>,
    },
    SessionEnded {
        status: String,
        duration_seconds: i64,
        exit_code: Option<i32>,
    },
    BudgetWarning {
        current_spend: f64,
        limit: f64,
    },
    BudgetExceeded {
        final_spend: f64,
        limit: f64,
    },
}

/// Append-only audit logger that writes to a local JSONL file.
pub struct AuditLog {
    path: PathBuf,
}

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

    /// Construct an AuditLog at a caller-provided path. Used by tests so they
    /// don't pollute the real audit log at `~/.local/share/audex/audit/audit.jsonl`.
    pub fn with_path(path: PathBuf) -> Result<Self> {
        if let Some(parent) = path.parent() {
            std::fs::create_dir_all(parent)?;
        }
        Ok(Self { path })
    }

    pub fn log(&self, session_id: &str, event: AuditEvent) -> Result<()> {
        self.log_with_provider(session_id, "aws", event)
    }

    pub fn log_with_provider(
        &self,
        session_id: &str,
        provider: &str,
        event: AuditEvent,
    ) -> Result<()> {
        let entry = AuditEntry {
            timestamp: Utc::now(),
            session_id: session_id.to_string(),
            provider: provider.to_string(),
            event,
        };
        let json_line = serde_json::to_string(&entry)?;

        // Redact any credential material before persisting
        let json_line = crate::leakdetect::redact_secrets(&json_line);

        // Take an exclusive OS-level lock on the audit log for the entire
        // read-hash-then-append sequence. Without this, concurrent writers
        // (threads OR separate tryaudex processes) would each read the same
        // `previous_hash`, compute HMACs against it, and append — breaking
        // the chain for every line after the first one to hit disk.
        use fs4::fs_std::FileExt;
        use std::io::Write;
        let mut file = std::fs::OpenOptions::new()
            .create(true)
            .append(true)
            .read(true)
            .open(&self.path)?;
        file.lock_exclusive()?;

        // Sign with chain HMAC for tamper detection — now inside the lock.
        let previous_hash = crate::integrity::last_chain_hash(&self.path);
        let signed_line = crate::integrity::sign_line(&json_line, &previous_hash);

        let mut output = signed_line;
        output.push('\n');

        let write_result = file.write_all(output.as_bytes());
        // Always release the lock, even if the write failed.
        let _ = FileExt::unlock(&file);
        write_result?;

        Ok(())
    }

    pub fn log_session_created(&self, session: &Session) -> Result<()> {
        let allowed_actions: Vec<String> = match session.provider {
            CloudProvider::Gcp => session
                .policy
                .actions
                .iter()
                .map(|a| a.to_gcp_permission())
                .collect(),
            CloudProvider::Azure => session
                .policy
                .actions
                .iter()
                .map(|a| a.to_azure_permission())
                .collect(),
            CloudProvider::Aws => session
                .policy
                .actions
                .iter()
                .map(|a| a.to_iam_action())
                .collect(),
        };

        self.log_with_provider(
            &session.id,
            &session.provider.to_string(),
            AuditEvent::SessionCreated {
                role_arn: session.role_arn.clone(),
                ttl_seconds: session.ttl_seconds,
                budget: session.budget,
                allowed_actions,
                command: session.command.clone(),
                agent_id: session.agent_id.clone(),
            },
        )
    }

    pub fn log_session_ended(&self, session: &Session, exit_code: Option<i32>) -> Result<()> {
        let duration = (Utc::now() - session.created_at).num_seconds();
        self.log_with_provider(
            &session.id,
            &session.provider.to_string(),
            AuditEvent::SessionEnded {
                status: session.status.to_string(),
                duration_seconds: duration,
                exit_code,
            },
        )
    }

    /// Read all audit entries, optionally filtered by session ID.
    pub fn read(&self, session_id: Option<&str>) -> Result<Vec<AuditEntry>> {
        self.read_filtered(session_id, None)
    }

    /// Read audit entries with optional session ID and provider filters.
    pub fn read_filtered(
        &self,
        session_id: Option<&str>,
        provider: Option<&str>,
    ) -> Result<Vec<AuditEntry>> {
        if !self.path.exists() {
            return Ok(Vec::new());
        }
        let content = std::fs::read_to_string(&self.path)?;
        let entries: Vec<AuditEntry> = content
            .lines()
            .filter(|l| !l.trim().is_empty())
            .filter_map(|l| {
                // Strip HMAC signature if present before parsing JSON
                let (json_content, _) = crate::integrity::parse_line(l);
                serde_json::from_str(json_content).ok()
            })
            .filter(|e: &AuditEntry| session_id.is_none_or(|id| e.session_id == id))
            .filter(|e: &AuditEntry| provider.is_none_or(|p| e.provider == p))
            .collect();
        Ok(entries)
    }

    /// Verify the integrity of the audit log using chain HMACs.
    pub fn verify(&self) -> Result<crate::integrity::VerificationResult> {
        crate::integrity::verify_audit_log(&self.path)
    }

    pub fn path(&self) -> &PathBuf {
        &self.path
    }
}