tryaudex_core/

credentials.rs

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

use serde::{Deserialize, Serialize};

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

/// Temporary AWS credentials returned by STS.
#[derive(Clone, Serialize, Deserialize)]
pub struct TempCredentials {
    pub access_key_id: String,
    pub secret_access_key: String,
    pub session_token: String,
    pub expires_at: chrono::DateTime<chrono::Utc>,
}

/// Cached GCP access token for session reuse (R6-H7).
#[derive(Clone, Serialize, Deserialize)]
pub struct CachedGcpToken {
    pub access_token: String,
    pub expires_at: chrono::DateTime<chrono::Utc>,
    /// Whether this token was successfully downscoped via Credential Access
    /// Boundaries. Used so a CAB-enforced session is never reused for a
    /// request that expected enforcement when the original fell back to
    /// advisory-only (R6-M19).
    #[serde(default)]
    pub scoping_enforced: bool,
}

impl std::fmt::Debug for CachedGcpToken {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("CachedGcpToken")
            .field("access_token", &"[REDACTED]")
            .field("expires_at", &self.expires_at)
            .field("scoping_enforced", &self.scoping_enforced)
            .finish()
    }
}

/// Cached Azure access token for session reuse (R6-H7).
#[derive(Clone, Serialize, Deserialize)]
pub struct CachedAzureToken {
    pub access_token: String,
    pub expires_at: chrono::DateTime<chrono::Utc>,
    /// ARM scope the token was minted for. Reuse only matches on identical
    /// scope so rotated / reused tokens target the correct audience.
    #[serde(default)]
    pub scope: Option<String>,
}

impl std::fmt::Debug for CachedAzureToken {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("CachedAzureToken")
            .field("access_token", &"[REDACTED]")
            .field("expires_at", &self.expires_at)
            .field("scope", &self.scope)
            .finish()
    }
}

impl std::fmt::Debug for TempCredentials {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("TempCredentials")
            .field("access_key_id", &self.access_key_id)
            .field("secret_access_key", &"[REDACTED]")
            .field("session_token", &"[REDACTED]")
            .field("expires_at", &self.expires_at)
            .finish()
    }
}

impl TempCredentials {
    /// Returns env vars to inject into the subprocess.
    pub fn as_env_vars(&self) -> Vec<(&str, &str)> {
        vec![
            ("AWS_ACCESS_KEY_ID", &self.access_key_id),
            ("AWS_SECRET_ACCESS_KEY", &self.secret_access_key),
            ("AWS_SESSION_TOKEN", &self.session_token),
        ]
    }
}

/// Issues scoped, short-lived AWS credentials via STS AssumeRole.
pub struct CredentialIssuer {
    sts_client: aws_sdk_sts::Client,
}

impl CredentialIssuer {
    pub async fn new() -> Result<Self> {
        Self::with_region(None).await
    }

    pub async fn with_region(region: Option<&str>) -> Result<Self> {
        let mut loader = aws_config::defaults(aws_config::BehaviorVersion::latest());
        if let Some(region) = region {
            loader = loader.region(aws_config::Region::new(region.to_string()));
        }
        let config = loader.load().await;
        Ok(Self {
            sts_client: aws_sdk_sts::Client::new(&config),
        })
    }

    /// Assume a role with an inline policy that restricts permissions to the scoped policy.
    /// Optionally attaches a permissions boundary ARN as an additional ceiling.
    pub async fn issue(
        &self,
        session: &Session,
        policy: &ScopedPolicy,
        ttl: Duration,
    ) -> Result<TempCredentials> {
        self.issue_with_boundary(session, policy, ttl, None).await
    }

    /// Assume a role with an inline policy and optional permissions boundary.
    pub async fn issue_with_boundary(
        &self,
        session: &Session,
        policy: &ScopedPolicy,
        ttl: Duration,
        permissions_boundary: Option<&str>,
    ) -> Result<TempCredentials> {
        self.issue_full(session, policy, ttl, permissions_boundary, None)
            .await
    }

    /// Assume a role with inline policy, optional boundary, and optional network conditions.
    pub async fn issue_full(
        &self,
        session: &Session,
        policy: &ScopedPolicy,
        ttl: Duration,
        permissions_boundary: Option<&str>,
        network: Option<&crate::policy::NetworkPolicy>,
    ) -> Result<TempCredentials> {
        self.issue_full_with_tag_lock(
            session,
            policy,
            ttl,
            permissions_boundary,
            network,
            None,
            None,
        )
        .await
    }

    /// Assume a role with inline policy plus an optional tag-lock deny that
    /// prevents the caller from modifying/removing a specific tag key. Used
    /// by `--tag-session` so agents can't strip the tryaudex-session marker.
    #[allow(clippy::too_many_arguments)]
    pub async fn issue_full_with_tag_lock(
        &self,
        session: &Session,
        policy: &ScopedPolicy,
        ttl: Duration,
        permissions_boundary: Option<&str>,
        network: Option<&crate::policy::NetworkPolicy>,
        tag_lock_key: Option<&str>,
        external_id: Option<&str>,
    ) -> Result<TempCredentials> {
        if let Some(net) = network {
            net.validate()?;
        }
        let policy_json = match tag_lock_key {
            Some(key) if network.is_none() => policy.to_iam_policy_json_with_tag_lock(key)?,
            Some(key) => policy.to_iam_policy_json_with_network_and_tag_lock(network, Some(key))?,
            None => policy.to_iam_policy_json_with_network(network)?,
        };
        // AWS STS AssumeRole requires 900s <= DurationSeconds <= 43200s (15 minutes to 12 hours).
        // Clamp at both ends rather than letting STS reject the call with a cryptic message.
        let requested = ttl.as_secs();
        let ttl_secs = requested.clamp(900, 43200) as i32;
        if requested < 900 {
            tracing::warn!(
                requested_secs = requested,
                clamped_to_secs = ttl_secs,
                "TTL below AWS STS minimum (900s / 15m); clamping up"
            );
        } else if requested > 43200 {
            tracing::warn!(
                requested_secs = requested,
                clamped_to_secs = ttl_secs,
                "TTL above AWS STS maximum (43200s / 12h); clamping down"
            );
        }

        tracing::info!(
            session_id = %session.id,
            role_arn = %session.role_arn,
            ttl_secs = ttl_secs,
            permissions_boundary = ?permissions_boundary,
            "Assuming role with scoped policy"
        );

        let mut request = self
            .sts_client
            .assume_role()
            .role_arn(&session.role_arn)
            .role_session_name(format!("av-{}", session.id.get(..8).unwrap_or(&session.id)))
            .policy(&policy_json)
            .duration_seconds(ttl_secs);

        // Attach user-supplied session tags to the STS call so they flow
        // into CloudTrail and are available for attribute-based access control.
        for (key, value) in &session.tags {
            request = request.tags(
                aws_sdk_sts::types::Tag::builder()
                    .key(key)
                    .value(value)
                    .build()
                    .map_err(|e| AvError::Sts(format!("Invalid STS tag: {}", e)))?,
            );
        }

        // Attach permissions boundary as a managed policy ARN ceiling
        if let Some(boundary_arn) = permissions_boundary {
            request = request.policy_arns(
                aws_sdk_sts::types::PolicyDescriptorType::builder()
                    .arn(boundary_arn)
                    .build(),
            );
        }

        // Attach external ID for cross-account role assumption
        if let Some(ext_id) = external_id {
            request = request.external_id(ext_id);
        }

        let result = request.send().await.map_err(|e| {
            // AWS SDK's Display impl collapses ServiceError variants to "service error".
            // Extract the actual error code + message so the user can act on it.
            let detail = match e.as_service_error() {
                Some(svc) => {
                    let code = svc.meta().code().unwrap_or("Unknown");
                    let message = svc.meta().message().unwrap_or("(no message)");
                    format!("{code}: {message}")
                }
                None => e.to_string(),
            };
            AvError::Sts(detail)
        })?;

        let creds = result
            .credentials()
            .ok_or_else(|| AvError::Sts("No credentials returned by STS".to_string()))?;

        let exp = creds.expiration();
        let expires_at = chrono::DateTime::from_timestamp(exp.secs(), exp.subsec_nanos())
            .unwrap_or_else(|| {
                tracing::warn!(
                    secs = exp.secs(),
                    "STS returned unparseable expiration timestamp; falling back to Utc::now() + TTL"
                );
                chrono::Utc::now() + chrono::Duration::seconds(ttl_secs as i64)
            });

        Ok(TempCredentials {
            access_key_id: creds.access_key_id().to_string(),
            secret_access_key: creds.secret_access_key().to_string(),
            session_token: creds.session_token().to_string(),
            expires_at,
        })
    }
}

/// Caches credentials on disk, keyed by session ID.
pub struct CredentialCache {
    dir: PathBuf,
}

impl CredentialCache {
    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("cred_cache");
        std::fs::create_dir_all(&dir)?;
        #[cfg(unix)]
        {
            use std::os::unix::fs::PermissionsExt;
            std::fs::set_permissions(&dir, std::fs::Permissions::from_mode(0o700))?;
        }
        Ok(Self { dir })
    }

    /// Save credentials for a session (encrypted at rest).
    pub fn save(&self, session_id: &str, creds: &TempCredentials) -> Result<()> {
        crate::validate::session_id(session_id)?;
        let path = self.dir.join(format!("{}.json", session_id));
        crate::keystore::encrypt_to_file(&path, creds)
    }

    /// Load cached credentials for a session. Returns None if expired or missing.
    /// Handles both encrypted and legacy plaintext files for migration.
    pub fn load(&self, session_id: &str) -> Result<Option<TempCredentials>> {
        crate::validate::session_id(session_id)?;
        let path = self.dir.join(format!("{}.json", session_id));
        let creds: Option<TempCredentials> = crate::keystore::decrypt_from_file(&path)?;
        match creds {
            Some(c) if c.expires_at <= chrono::Utc::now() => {
                let _ = std::fs::remove_file(&path);
                Ok(None)
            }
            other => Ok(other),
        }
    }

    /// Remove cached credentials for a session.
    pub fn remove(&self, session_id: &str) -> Result<()> {
        crate::validate::session_id(session_id)?;
        let path = self.dir.join(format!("{}.json", session_id));
        let _ = std::fs::remove_file(path);
        // Also clean up provider-specific caches.
        let _ = std::fs::remove_file(self.dir.join(format!("{}.gcp.json", session_id)));
        let _ = std::fs::remove_file(self.dir.join(format!("{}.azure.json", session_id)));
        Ok(())
    }

    /// Save a GCP token for a session (encrypted at rest).
    pub fn save_gcp(&self, session_id: &str, token: &CachedGcpToken) -> Result<()> {
        crate::validate::session_id(session_id)?;
        let path = self.dir.join(format!("{}.gcp.json", session_id));
        crate::keystore::encrypt_to_file(&path, token)
    }

    /// Load a cached GCP token. Returns None if expired or missing.
    pub fn load_gcp(&self, session_id: &str) -> Result<Option<CachedGcpToken>> {
        crate::validate::session_id(session_id)?;
        let path = self.dir.join(format!("{}.gcp.json", session_id));
        let token: Option<CachedGcpToken> = crate::keystore::decrypt_from_file(&path)?;
        match token {
            Some(t) if t.expires_at <= chrono::Utc::now() => {
                let _ = std::fs::remove_file(&path);
                Ok(None)
            }
            other => Ok(other),
        }
    }

    /// Save an Azure token for a session (encrypted at rest).
    pub fn save_azure(&self, session_id: &str, token: &CachedAzureToken) -> Result<()> {
        crate::validate::session_id(session_id)?;
        let path = self.dir.join(format!("{}.azure.json", session_id));
        crate::keystore::encrypt_to_file(&path, token)
    }

    /// Load a cached Azure token. Returns None if expired or missing.
    pub fn load_azure(&self, session_id: &str) -> Result<Option<CachedAzureToken>> {
        crate::validate::session_id(session_id)?;
        let path = self.dir.join(format!("{}.azure.json", session_id));
        let token: Option<CachedAzureToken> = crate::keystore::decrypt_from_file(&path)?;
        match token {
            Some(t) if t.expires_at <= chrono::Utc::now() => {
                let _ = std::fs::remove_file(&path);
                Ok(None)
            }
            other => Ok(other),
        }
    }
}