NxMesh/apps/nxmesh-agent/src/connector/master/ssh.rs

use std::{fs::File, io::Read};

use tokio::fs::read;

use nxmesh_proto::agent_service_client::AgentServiceClient;
use tonic::transport::{Certificate, ClientTlsConfig, Identity};
use tracing::warn;

use crate::config::settings::{MAuthSettings, TLSSettings};

use super::{AgentClient, MasterConnectorTrait};

pub struct SshMasterConnector {
    client: AgentClient,
}

impl SshMasterConnector {
    pub async fn new(
        settings: crate::config::settings::GrpcSettings,
    ) -> Result<Self, Box<dyn std::error::Error + Send + Sync>> {
        let tls_config = Self::generate_tls_config(&settings.m_auth).await?;
        // Create a gRPC channel
        let endpoint = tonic::transport::Channel::from_shared(settings.connection_string.clone())
            .map_err(|e| format!("Failed to create gRPC endpoint: {}", e))?
            .tls_config(tls_config)
            .map_err(|e| {
                format!(
                    "Failed to set TLS config: {}. Ensure TLS settings and certificates are correct.",
                    e
                )
            })?
            .connect_timeout(std::time::Duration::from_secs(5))
            .timeout(std::time::Duration::from_secs(10))
            .connect_lazy();

        // Create the gRPC client
        let client = AgentServiceClient::new(endpoint);
        Ok(Self { client })
    }

    async fn generate_tls_config(
        settings: &MAuthSettings,
    ) -> Result<ClientTlsConfig, Box<dyn std::error::Error + Send + Sync>> {
        let tls_config = match &settings {
            MAuthSettings::Tls(tls_settings) => {
                let (ca, cert, key) = match tls_settings {
                    TLSSettings::RawPath {
                        ca_path,
                        cert_path,
                        key_path,
                    } => {
                        // Read the certificate and key from the specified file paths
                        let ca = read(ca_path).await?;
                        let cert = read(cert_path).await?;
                        let key = read(key_path).await?;
                        (ca, cert, key)
                    }
                    TLSSettings::ZipPath { cert_zip_path } => {
                        // Extract the certificate and key from the zip file
                        Self::extract_certificate(cert_zip_path).await?
                    }
                };

                // TODO: allow skipping SANs validation if specified in the settings
                ClientTlsConfig::new()
                    .ca_certificate(Certificate::from_pem(&ca))
                    .identity(Identity::from_pem(&cert, &key))
            }
            #[allow(unreachable_patterns)]
            _ => {
                return Err("TLS settings are required for SSH connection".into());
            }
        };

        Ok(tls_config)
    }

    async fn extract_certificate(
        cert_zip_path: &str,
    ) -> Result<(Vec<u8>, Vec<u8>, Vec<u8>), Box<dyn std::error::Error + Send + Sync>> {
        // unzip the file and extract the cert, ca and key
        let file = File::open(cert_zip_path)?;
        let mut archive = zip::ZipArchive::new(file)?;
        let mut cert = Vec::new();
        let mut key = Vec::new();
        let mut ca = Vec::new();

        for i in 0..archive.len() {
            let mut file = archive.by_index(i)?;
            let outpath = match file.enclosed_name() {
                Some(path) => path.to_owned(),
                None => continue,
            };
            let file_name = outpath
                .file_name()
                .and_then(|n| n.to_str())
                .unwrap_or_default();
            if file_name != "cert.pem" && file_name != "key.pem" && file_name != "ca.pem" {
                warn!("Unexpected file in certificate zip: {}", file_name);
                continue;
            }
            if file_name == "cert.pem" {
                file.read_to_end(&mut cert)?;
            } else if file_name == "key.pem" {
                file.read_to_end(&mut key)?;
            } else if file_name == "ca.pem" {
                file.read_to_end(&mut ca)?;
            }
        }

        if cert.is_empty() || key.is_empty() || ca.is_empty() {
            return Err("Certificate zip must contain cert.pem, key.pem and ca.pem".into());
        }

        Ok((ca, cert, key))
    }
}

#[async_trait::async_trait]
impl MasterConnectorTrait for SshMasterConnector {
    async fn connect(
        &mut self,
        _settings: &crate::config::settings::Settings,
    ) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
        // ensure connection if required
        Ok(())
    }

    fn get_client(&self) -> AgentClient {
        self.client.clone()
    }
}

#[cfg(test)]
#[allow(clippy::expect_used)]
mod tests {
    use std::{
        fs::{self, File},
        io::Write,
        path::Path,
    };

    use tempfile::TempDir;

    use crate::config::settings::{MAuthSettings, TLSSettings};

    use super::SshMasterConnector;

    const CERT_PEM: &[u8] = b"-----BEGIN CERTIFICATE-----\nAQ==\n-----END CERTIFICATE-----\n";
    const KEY_PEM: &[u8] = b"-----BEGIN PRIVATE KEY-----\nAQ==\n-----END PRIVATE KEY-----\n";
    const CA_PEM: &[u8] = b"-----BEGIN CERTIFICATE-----\nAQ==\n-----END CERTIFICATE-----\n";

    fn create_zip_with_entries(
        dir: &TempDir,
        file_name: &str,
        entries: &[(&str, &[u8])],
    ) -> Result<String, Box<dyn std::error::Error + Send + Sync>> {
        let zip_path = dir.path().join(file_name);
        let file = File::create(&zip_path)?;
        let mut zip = zip::ZipWriter::new(file);
        let options = zip::write::SimpleFileOptions::default()
            .compression_method(zip::CompressionMethod::Deflated)
            .unix_permissions(0o600);

        for (entry_name, contents) in entries {
            zip.start_file(entry_name, options)?;
            zip.write_all(contents)?;
        }

        zip.finish()?;
        Ok(zip_path.to_string_lossy().to_string())
    }

    fn write_file(
        path: &Path,
        contents: &[u8],
    ) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
        fs::write(path, contents)?;
        Ok(())
    }

    #[tokio::test]
    async fn extract_certificate_succeeds_with_expected_files() {
        let temp_dir = TempDir::new().expect("failed to create temp dir");
        let zip_path = create_zip_with_entries(
            &temp_dir,
            "certs.zip",
            &[
                ("cert.pem", CERT_PEM),
                ("key.pem", KEY_PEM),
                ("ca.pem", CA_PEM),
                ("ignored.txt", b"ignored"),
            ],
        )
        .expect("failed to create zip");

        let (ca, cert, key) = SshMasterConnector::extract_certificate(&zip_path)
            .await
            .expect("expected cert extraction to succeed");

        assert_eq!(ca, CA_PEM);
        assert_eq!(cert, CERT_PEM);
        assert_eq!(key, KEY_PEM);
    }

    #[tokio::test]
    async fn extract_certificate_fails_when_required_files_are_missing() {
        let temp_dir = TempDir::new().expect("failed to create temp dir");
        let zip_path = create_zip_with_entries(
            &temp_dir,
            "missing-key.zip",
            &[("cert.pem", CERT_PEM), ("ca.pem", CA_PEM)],
        )
        .expect("failed to create zip");

        let err = SshMasterConnector::extract_certificate(&zip_path)
            .await
            .expect_err("expected extraction to fail when key.pem is missing");

        assert!(
            err.to_string()
                .contains("Certificate zip must contain cert.pem, key.pem and ca.pem")
        );
    }

    #[tokio::test]
    async fn generate_tls_config_succeeds_for_raw_paths() {
        let temp_dir = TempDir::new().expect("failed to create temp dir");
        let cert_path = temp_dir.path().join("cert.pem");
        let key_path = temp_dir.path().join("key.pem");
        let ca_path = temp_dir.path().join("ca.pem");

        write_file(&cert_path, CERT_PEM).expect("failed to write cert.pem");
        write_file(&key_path, KEY_PEM).expect("failed to write key.pem");
        write_file(&ca_path, CA_PEM).expect("failed to write ca.pem");

        let settings = MAuthSettings::Tls(TLSSettings::RawPath {
            ca_path: ca_path.to_string_lossy().to_string(),
            cert_path: cert_path.to_string_lossy().to_string(),
            key_path: key_path.to_string_lossy().to_string(),
        });

        let result = SshMasterConnector::generate_tls_config(&settings).await;
        assert!(result.is_ok(), "expected raw path TLS config to succeed");
    }

    #[tokio::test]
    async fn generate_tls_config_succeeds_for_zip_path() {
        let temp_dir = TempDir::new().expect("failed to create temp dir");
        let zip_path = create_zip_with_entries(
            &temp_dir,
            "certs.zip",
            &[
                ("cert.pem", CERT_PEM),
                ("key.pem", KEY_PEM),
                ("ca.pem", CA_PEM),
            ],
        )
        .expect("failed to create zip");

        let settings = MAuthSettings::Tls(TLSSettings::ZipPath {
            cert_zip_path: zip_path,
        });

        let result = SshMasterConnector::generate_tls_config(&settings).await;
        assert!(result.is_ok(), "expected zip path TLS config to succeed");
    }

    #[tokio::test]
    async fn generate_tls_config_fails_for_missing_raw_files() {
        let settings = MAuthSettings::Tls(TLSSettings::RawPath {
            ca_path: "/tmp/non-existent-ca.pem".to_string(),
            cert_path: "/tmp/non-existent-cert.pem".to_string(),
            key_path: "/tmp/non-existent-key.pem".to_string(),
        });

        let result = SshMasterConnector::generate_tls_config(&settings).await;
        assert!(result.is_err(), "expected raw path TLS config to fail");
    }
}