botserver/src/security/integration.rs

use anyhow::{Context, Result};
use reqwest::{Certificate, Client, ClientBuilder, Identity};
use std::collections::HashMap;
use std::fs;
use std::path::{Path, PathBuf};
use std::sync::{Arc, OnceLock};
use std::time::Duration;
use tracing::{info, warn};

#[derive(Debug, Clone)]
pub struct ServiceUrls {
    pub original: String,
    pub secure: String,
    pub port: u16,
    pub tls_port: u16,
}

#[derive(Debug)]
pub struct TlsIntegration {
    services: HashMap<String, ServiceUrls>,

    ca_cert: Option<Certificate>,

    client_certs: HashMap<String, Identity>,

    tls_enabled: bool,

    https_only: bool,
}

impl TlsIntegration {
    pub fn new(tls_enabled: bool) -> Self {
        let mut services = HashMap::new();

        services.insert(
            "api".to_string(),
            ServiceUrls {
                original: "http://localhost:8080".to_string(),
                secure: "https://localhost:8443".to_string(),
                port: 8080,
                tls_port: 8443,
            },
        );

        services.insert(
            "llm".to_string(),
            ServiceUrls {
                original: "http://localhost:8081".to_string(),
                secure: "https://localhost:8444".to_string(),
                port: 8081,
                tls_port: 8444,
            },
        );

        services.insert(
            "embedding".to_string(),
            ServiceUrls {
                original: "http://localhost:8082".to_string(),
                secure: "https://localhost:8445".to_string(),
                port: 8082,
                tls_port: 8445,
            },
        );

        services.insert(
            "qdrant".to_string(),
            ServiceUrls {
                original: "http://localhost:6333".to_string(),
                secure: "https://localhost:6334".to_string(),
                port: 6333,
                tls_port: 6334,
            },
        );

        services.insert(
            "redis".to_string(),
            ServiceUrls {
                original: "redis://localhost:6379".to_string(),
                secure: "rediss://localhost:6380".to_string(),
                port: 6379,
                tls_port: 6380,
            },
        );

        services.insert(
            "postgres".to_string(),
            ServiceUrls {
                original: "postgres://localhost:5432".to_string(),
                secure: "postgres://localhost:5433?sslmode=require".to_string(),
                port: 5432,
                tls_port: 5433,
            },
        );

        services.insert(
            "minio".to_string(),
            ServiceUrls {
                original: "http://localhost:9000".to_string(),
                secure: "https://localhost:9001".to_string(),
                port: 9000,
                tls_port: 9001,
            },
        );

        services.insert(
            "directory".to_string(),
            ServiceUrls {
                original: "http://localhost:8080".to_string(),
                secure: "https://localhost:8446".to_string(),
                port: 8080,
                tls_port: 8446,
            },
        );

        Self {
            services,
            ca_cert: None,
            client_certs: HashMap::new(),
            tls_enabled,
            https_only: tls_enabled,
        }
    }

    pub fn load_ca_cert(&mut self, ca_path: &Path) -> Result<()> {
        if ca_path.exists() {
            let ca_cert_pem = fs::read(ca_path)
                .with_context(|| format!("Failed to read CA certificate from {}", ca_path.display()))?;

            let ca_cert =
                Certificate::from_pem(&ca_cert_pem).context("Failed to parse CA certificate")?;

            self.ca_cert = Some(ca_cert);
            info!("Loaded CA certificate from {}", ca_path.display());
        } else {
            warn!("CA certificate not found at {}", ca_path.display());
        }

        Ok(())
    }

    pub fn load_client_cert(
        &mut self,
        service: &str,
        cert_path: &Path,
        key_path: &Path,
    ) -> Result<()> {
        if cert_path.exists() && key_path.exists() {
            let cert = fs::read(cert_path)
                .with_context(|| format!("Failed to read client cert from {}", cert_path.display()))?;

            let key = fs::read(key_path)
                .with_context(|| format!("Failed to read client key from {}", key_path.display()))?;

            let identity = Identity::from_pem(&[&cert[..], &key[..]].concat())
                .context("Failed to create client identity")?;

            self.client_certs.insert(service.to_string(), identity);
            info!("Loaded client certificate for service: {}", service);
        } else {
            warn!("Client certificate not found for service: {}", service);
        }

        Ok(())
    }

    pub fn convert_url(&self, url: &str) -> String {
        if !self.tls_enabled {
            return url.to_string();
        }

        for urls in self.services.values() {
            if url.starts_with(&urls.original) {
                return url.replace(&urls.original, &urls.secure);
            }
        }

        if url.starts_with("http://") {
            url.replace("http://", "https://")
        } else if url.starts_with("redis://") {
            url.replace("redis://", "rediss://")
        } else {
            url.to_string()
        }
    }

    pub fn get_service_url(&self, service: &str) -> Option<String> {
        self.services.get(service).map(|urls| {
            if self.tls_enabled {
                urls.secure.clone()
            } else {
                urls.original.clone()
            }
        })
    }

    pub fn create_client(&self, service: &str) -> Result<Client> {
        let mut builder = ClientBuilder::new()
            .timeout(Duration::from_secs(30))
            .connect_timeout(Duration::from_secs(10));

        if self.tls_enabled {
            builder = builder.use_rustls_tls();

            if let Some(ca_cert) = &self.ca_cert {
                builder = builder.add_root_certificate(ca_cert.clone());
            }

            if let Some(identity) = self.client_certs.get(service) {
                builder = builder.identity(identity.clone());
            }

            if cfg!(debug_assertions) {
                builder = builder.danger_accept_invalid_certs(true);
            }

            if self.https_only {
                builder = builder.https_only(true);
            }
        }

        builder.build().context("Failed to build HTTP client")
    }

    pub fn create_generic_client(&self) -> Result<Client> {
        self.create_client("generic")
    }

    pub fn is_tls_enabled(&self) -> bool {
        self.tls_enabled
    }

    pub fn get_secure_port(&self, service: &str) -> Option<u16> {
        self.services.get(service).map(|urls| {
            if self.tls_enabled {
                urls.tls_port
            } else {
                urls.port
            }
        })
    }

    pub fn update_postgres_url(&self, url: &str) -> String {
        if !self.tls_enabled {
            return url.to_string();
        }

        if url.contains("localhost:5432") || url.contains("127.0.0.1:5432") {
            let base = url
                .replace("localhost:5432", "localhost:5433")
                .replace("127.0.0.1:5432", "127.0.0.1:5433");

            if base.contains("sslmode=") {
                base
            } else if base.contains('?') {
                format!("{}&sslmode=require", base)
            } else {
                format!("{}?sslmode=require", base)
            }
        } else {
            url.to_string()
        }
    }

    pub fn update_redis_url(&self, url: &str) -> String {
        if !self.tls_enabled {
            return url.to_string();
        }

        if url.starts_with("redis://") {
            url.replace("redis://", "rediss://")
                .replace(":6379", ":6380")
        } else {
            url.to_string()
        }
    }

    pub fn load_all_certs_from_dir(&mut self, cert_dir: &Path) -> Result<()> {
        let ca_path = cert_dir.join("ca.crt");
        if ca_path.exists() {
            self.load_ca_cert(&ca_path)?;
        }

        for service in &[
            "api",
            "llm",
            "embedding",
            "qdrant",
            "postgres",
            "redis",
            "minio",
        ] {
            let service_dir = cert_dir.join(service);
            if service_dir.exists() {
                let cert_path = service_dir.join("client.crt");
                let key_path = service_dir.join("client.key");

                if cert_path.exists() && key_path.exists() {
                    self.load_client_cert(service, &cert_path, &key_path)?;
                }
            }
        }

        Ok(())
    }
}

static TLS_INTEGRATION: OnceLock<Arc<TlsIntegration>> = OnceLock::new();

pub fn init_tls_integration(tls_enabled: bool, cert_dir: Option<PathBuf>) -> Result<()> {
    let _ = TLS_INTEGRATION.get_or_init(|| {
        let mut integration = TlsIntegration::new(tls_enabled);

        if tls_enabled {
            if let Some(dir) = cert_dir {
                if let Err(e) = integration.load_all_certs_from_dir(&dir) {
                    warn!("Failed to load some certificates: {}", e);
                }
            }
        }

        info!("TLS integration initialized (TLS: {})", tls_enabled);
        Arc::new(integration)
    });

    Ok(())
}

pub fn get_tls_integration() -> Option<Arc<TlsIntegration>> {
    TLS_INTEGRATION.get().cloned()
}

pub fn to_secure_url(url: &str) -> String {
    if let Some(integration) = get_tls_integration() {
        integration.convert_url(url)
    } else {
        url.to_string()
    }
}

pub fn create_https_client(service: &str) -> Result<Client> {
    if let Some(integration) = get_tls_integration() {
        integration.create_client(service)
    } else {
        Client::builder()
            .timeout(Duration::from_secs(30))
            .build()
            .context("Failed to build default HTTP client")
    }
}