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> { 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> { 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, Vec, Vec), Box> { // 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> { // 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> { 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> { 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"); } }