bascule-oss/crates/bascule-core/src/config.rs

//! Configuration — loaded from TOML file.

use serde::Deserialize;

#[derive(Debug, Deserialize)]
pub struct BasculeConfig {
    /// Address to listen on (default: 0.0.0.0:2222)
    #[serde(default = "default_listen")]
    pub listen_addr: String,

    /// Path to host key (generated if not present)
    pub host_key_path: Option<String>,

    /// Command to spawn for shell sessions.
    /// If not set, uses the user's login shell.
    pub shell_command: Option<String>,

    /// Arguments for shell_command.
    #[serde(default)]
    pub shell_args: Vec<String>,

    /// Authentication configuration.
    #[serde(default)]
    pub auth: AuthConfig,

    /// Session banner (shown after auth).
    pub banner: Option<String>,

    /// Maximum concurrent sessions (0 = unlimited).
    #[serde(default)]
    pub max_sessions: usize,

    /// Remote proxy configuration.
    /// When set, sessions are forwarded to a target SSH host
    /// instead of spawning a local shell.
    pub proxy: Option<ProxyConfig>,

    /// Container backend configuration.
    /// When set, sessions spawn an ephemeral container per connection.
    /// Priority: proxy > container > local PTY.
    pub container: Option<ContainerConfig>,

    /// K8s backend configuration.
    /// When running in-cluster, exec into a shell sidecar instead of local PTY.
    pub k8s: Option<K8sConfig>,

    /// Telemetry (OTel tracing).
    #[serde(default)]
    pub telemetry: TelemetryConfig,

    /// Prometheus metrics.
    #[serde(default)]
    pub metrics: MetricsConfig,
}

#[derive(Debug, Deserialize, Clone)]
pub struct ProxyConfig {
    /// Target SSH host to forward sessions to.
    pub target_host: String,
    /// Target SSH port (default: 22).
    #[serde(default = "default_ssh_port")]
    pub target_port: u16,
    /// Username on the target host.
    /// If not set, uses the authenticated principal.
    pub target_user: Option<String>,
    /// Path to private key for target host authentication.
    /// If not set, uses agent forwarding or password from the client.
    pub target_key_path: Option<String>,
    /// Accept any host key from target (dev only — disable in production).
    #[serde(default)]
    pub accept_target_host_key: bool,
}

#[derive(Debug, Deserialize)]
pub struct AuthConfig {
    /// Auth mode: "accept-all" (dev), "authorized-keys"
    #[serde(default = "default_auth_mode")]
    pub mode: String,

    /// Path to authorized_keys file (for authorized-keys mode).
    pub authorized_keys_path: Option<String>,

    /// AI agent authentication via Microsoft Entra Agent ID.
    pub agent_id: Option<AgentIdConfig>,

    /// SPIFFE/SPIRE workload identity authentication.
    pub spiffe: Option<SpiffeConfig>,
}

impl Default for AuthConfig {
    fn default() -> Self {
        Self {
            mode: default_auth_mode(),
            authorized_keys_path: None,
            agent_id: None,
            spiffe: None,
        }
    }
}

#[derive(Debug, Deserialize)]
pub struct AgentIdConfig {
    /// Entra tenant ID.
    pub tenant_id: String,
    /// Expected token audiences.
    #[serde(default)]
    pub audiences: Vec<String>,
    /// Accept agents from any tenant (multi-tenant mode).
    #[serde(default)]
    pub multi_tenant: bool,
}

impl Default for BasculeConfig {
    fn default() -> Self {
        Self {
            listen_addr: default_listen(),
            host_key_path: None,
            shell_command: None,
            shell_args: vec![],
            auth: AuthConfig::default(),
            banner: None,
            max_sessions: 0,
            proxy: None,
            container: None,
            k8s: None,
            telemetry: TelemetryConfig::default(),
            metrics: MetricsConfig::default(),
        }
    }
}

impl BasculeConfig {
    pub fn from_toml(toml_str: &str) -> anyhow::Result<Self> {
        Ok(toml::from_str(toml_str)?)
    }

    pub fn from_file(path: &str) -> anyhow::Result<Self> {
        let content = std::fs::read_to_string(path)?;
        Self::from_toml(&content)
    }
}

fn default_listen() -> String {
    "0.0.0.0:2222".to_string()
}

fn default_auth_mode() -> String {
    "accept-all".to_string()
}

fn default_ssh_port() -> u16 {
    22
}

/// Container backend configuration.
#[derive(Debug, Deserialize, Clone)]
pub struct ContainerConfig {
    /// Container runtime: "docker", "podman", "nerdctl", "auto".
    #[serde(default = "default_runtime")]
    pub runtime: String,
    /// Container image to use.
    pub image: String,
    /// Image pull policy: "always", "if-not-present", "never".
    #[serde(default = "default_pull_policy")]
    pub pull_policy: String,
    /// Volume mounts.
    #[serde(default)]
    pub mounts: Vec<MountConfig>,
    /// Extra environment variables.
    #[serde(default)]
    pub env: std::collections::HashMap<String, String>,
    /// Memory limit (e.g. "512m", "1g").
    pub memory_limit: Option<String>,
    /// CPU limit (e.g. "1.0", "0.5").
    pub cpu_limit: Option<String>,
    /// Shell to run inside the container.
    pub shell: Option<String>,
    /// User to run as inside the container.
    pub user: Option<String>,
    /// Destroy container on session end (default: true).
    #[serde(default = "default_true")]
    pub ephemeral: bool,
    /// Drop all capabilities, add back minimal set (default: true).
    #[serde(default = "default_true")]
    pub hardened: bool,
    /// Read-only root filesystem.
    #[serde(default)]
    pub read_only_rootfs: bool,
    /// Network mode (e.g. "none", "bridge", "host").
    pub network: Option<String>,
}

#[derive(Debug, Deserialize, Clone)]
pub struct MountConfig {
    pub source: String,
    pub target: String,
    #[serde(default)]
    pub readonly: bool,
}

/// SPIFFE/SPIRE authentication configuration.
#[derive(Debug, Deserialize, Clone)]
pub struct SpiffeConfig {
    /// SPIFFE trust domain.
    pub trust_domain: String,
    /// Path to trust bundle PEM file.
    pub trust_bundle_path: Option<String>,
    /// SPIRE Workload API socket path.
    pub workload_api_socket: Option<String>,
}

/// K8s backend configuration.
#[derive(Debug, Deserialize, Clone)]
pub struct K8sConfig {
    /// Enable K8s backend.
    #[serde(default)]
    pub enabled: bool,
    /// Namespace (auto-detected from downward API if not set).
    pub namespace: Option<String>,
    /// Pod name (auto-detected from downward API if not set).
    pub pod_name: Option<String>,
    /// Shell container name in the Pod.
    #[serde(default = "default_shell_container")]
    pub shell_container: String,
    /// Shell command inside the container.
    #[serde(default = "default_k8s_shell")]
    pub shell: String,
}

fn default_shell_container() -> String { "shell".to_string() }
fn default_k8s_shell() -> String { "/bin/bash".to_string() }

/// Telemetry configuration (OTel + metrics).
#[derive(Debug, Deserialize, Clone, Default)]
pub struct TelemetryConfig {
    /// OTLP endpoint for trace export.
    pub otlp_endpoint: Option<String>,
    /// Service name for OTel spans.
    #[serde(default = "default_service_name")]
    pub service_name: String,
}

/// Prometheus metrics endpoint configuration.
#[derive(Debug, Deserialize, Clone, Default)]
pub struct MetricsConfig {
    /// Enable metrics endpoint.
    #[serde(default)]
    pub enabled: bool,
    /// Port for /metrics endpoint.
    #[serde(default = "default_metrics_port")]
    pub port: u16,
}

impl ContainerConfig {
    /// Validate config values to prevent CLI argument injection.
    /// Call at startup before accepting any connections.
    pub fn validate(&self) -> anyhow::Result<()> {
        // Memory limit: digits + optional k/m/g suffix
        if let Some(ref mem) = self.memory_limit {
            if !mem
                .chars()
                .all(|c| c.is_ascii_digit() || "kmgbKMGB".contains(c))
            {
                anyhow::bail!("Invalid memory_limit: '{}'. Expected format: 512m, 1g", mem);
            }
        }

        // CPU limit: valid float
        if let Some(ref cpu) = self.cpu_limit {
            if cpu.parse::<f64>().is_err() {
                anyhow::bail!("Invalid cpu_limit: '{}'. Expected format: 1.0, 0.5", cpu);
            }
        }

        // Image: no shell metacharacters, no flags
        if self.image.starts_with('-')
            || self.image.contains(';')
            || self.image.contains('|')
            || self.image.contains('&')
            || self.image.contains('$')
            || self.image.contains('`')
        {
            anyhow::bail!(
                "Invalid image name: '{}'. Contains shell metacharacters.",
                self.image
            );
        }

        // Network: alphanumeric + hyphens/underscores/colons only
        if let Some(ref net) = self.network {
            if !net
                .chars()
                .all(|c| c.is_alphanumeric() || c == '-' || c == '_' || c == ':')
            {
                anyhow::bail!(
                    "Invalid network: '{}'. Alphanumeric, hyphens, colons only.",
                    net
                );
            }
        }

        // User: alphanumeric + common user chars
        if let Some(ref user) = self.user {
            if !user
                .chars()
                .all(|c| c.is_alphanumeric() || c == '-' || c == '_' || c == ':')
            {
                anyhow::bail!(
                    "Invalid user: '{}'. Alphanumeric, hyphens, colons only.",
                    user
                );
            }
        }

        // Mount paths: no shell metacharacters
        for mount in &self.mounts {
            for path in [&mount.source, &mount.target] {
                if path.contains(';')
                    || path.contains('|')
                    || path.contains('&')
                    || path.contains('$')
                    || path.contains('`')
                {
                    anyhow::bail!(
                        "Invalid mount path: '{}'. Contains shell metacharacters.",
                        path
                    );
                }
            }
        }

        // Env keys: alphanumeric + underscores
        for key in self.env.keys() {
            if !key.chars().all(|c| c.is_alphanumeric() || c == '_') {
                anyhow::bail!(
                    "Invalid env var name: '{}'. Alphanumeric + underscores only.",
                    key
                );
            }
        }

        Ok(())
    }
}

fn default_service_name() -> String {
    "bascule".to_string()
}
fn default_metrics_port() -> u16 {
    9090
}

fn default_runtime() -> String {
    "auto".to_string()
}
fn default_pull_policy() -> String {
    "if-not-present".to_string()
}
fn default_true() -> bool {
    true
}

#[cfg(test)]
mod tests {
    use super::*;

    fn valid_container_config() -> ContainerConfig {
        ContainerConfig {
            runtime: "docker".to_string(),
            image: "bascule-shell:minimal".to_string(),
            pull_policy: "if-not-present".to_string(),
            mounts: vec![],
            env: std::collections::HashMap::new(),
            memory_limit: Some("512m".to_string()),
            cpu_limit: Some("1.0".to_string()),
            shell: None,
            user: Some("operator".to_string()),
            ephemeral: true,
            hardened: true,
            read_only_rootfs: false,
            network: Some("none".to_string()),
        }
    }

    #[test]
    fn test_valid_container_config() {
        assert!(valid_container_config().validate().is_ok());
    }

    #[test]
    fn test_invalid_memory_limit() {
        let mut cfg = valid_container_config();
        cfg.memory_limit = Some("--privileged".to_string());
        assert!(cfg.validate().is_err());
    }

    #[test]
    fn test_invalid_cpu_limit() {
        let mut cfg = valid_container_config();
        cfg.cpu_limit = Some("--privileged".to_string());
        assert!(cfg.validate().is_err());
    }

    #[test]
    fn test_invalid_image_flag() {
        let mut cfg = valid_container_config();
        cfg.image = "--privileged".to_string();
        assert!(cfg.validate().is_err());
    }

    #[test]
    fn test_invalid_image_injection() {
        let mut cfg = valid_container_config();
        cfg.image = "ubuntu; rm -rf /".to_string();
        assert!(cfg.validate().is_err());
    }

    #[test]
    fn test_invalid_network() {
        let mut cfg = valid_container_config();
        cfg.network = Some("host; malicious".to_string());
        assert!(cfg.validate().is_err());
    }

    #[test]
    fn test_invalid_env_key() {
        let mut cfg = valid_container_config();
        cfg.env.insert("VALID_KEY".to_string(), "ok".to_string());
        assert!(cfg.validate().is_ok());
        cfg.env.insert("BAD;KEY".to_string(), "bad".to_string());
        assert!(cfg.validate().is_err());
    }

    #[test]
    fn test_invalid_mount_path() {
        let mut cfg = valid_container_config();
        cfg.mounts.push(MountConfig {
            source: "/host/path".to_string(),
            target: "/container; rm -rf /".to_string(),
            readonly: true,
        });
        assert!(cfg.validate().is_err());
    }

    #[test]
    fn test_config_from_toml() {
        let toml = r#"
listen_addr = "127.0.0.1:2222"
[auth]
mode = "authorized-keys"
authorized_keys_path = "/etc/bascule/authorized_keys"
"#;
        let config = BasculeConfig::from_toml(toml).unwrap();
        assert_eq!(config.listen_addr, "127.0.0.1:2222");
        assert_eq!(config.auth.mode, "authorized-keys");
        assert_eq!(
            config.auth.authorized_keys_path.unwrap(),
            "/etc/bascule/authorized_keys"
        );
    }

    #[test]
    fn test_config_defaults() {
        let config = BasculeConfig::default();
        assert_eq!(config.listen_addr, "0.0.0.0:2222");
        assert_eq!(config.auth.mode, "accept-all");
        assert_eq!(config.max_sessions, 0);
    }
}