From 91913bd6e483a366b3058e5774679dfb6ca14d16 Mon Sep 17 00:00:00 2001 From: Graham McIntire Date: Wed, 4 Feb 2026 12:02:10 -0600 Subject: [PATCH] Implement snmp v3 support --- proto/agent.proto | 9 + src/main.rs | 138 ++++++++++- src/snmp/client.rs | 39 +++- src/snmp/device_poller.rs | 440 ++++++++++++++++++++++++++++++++++++ src/snmp/mod.rs | 4 + src/snmp/poller_registry.rs | 146 ++++++++++++ src/snmp/types.rs | 4 + src/websocket_client.rs | 240 +++++++++++++++++--- 8 files changed, 983 insertions(+), 37 deletions(-) create mode 100644 src/snmp/device_poller.rs create mode 100644 src/snmp/poller_registry.rs diff --git a/proto/agent.proto b/proto/agent.proto index 75cafef..45149f1 100644 --- a/proto/agent.proto +++ b/proto/agent.proto @@ -113,6 +113,7 @@ enum JobType { DISCOVER = 0; POLL = 1; MIKROTIK = 2; + TEST_CREDENTIALS = 3; } enum QueryType { @@ -174,6 +175,14 @@ message AgentError { int64 timestamp = 3; } +message CredentialTestResult { + string test_id = 1; + bool success = 2; + string error_message = 3; // Empty if success + string system_description = 4; // sysDescr.0 value if success + int64 timestamp = 5; +} + // MikroTik RouterOS API messages message MikrotikDevice { diff --git a/src/main.rs b/src/main.rs index fadb5c0..554d4ce 100644 --- a/src/main.rs +++ b/src/main.rs @@ -92,13 +92,43 @@ struct Args { /// Use plain TCP instead of SSL (port 8728) - WARNING: credentials sent in plaintext #[arg(long, default_value_t = false)] mikrotik_plain: bool, + + /// Run SNMPv3 test instead of normal agent operation + #[arg(long)] + snmpv3_test: bool, + + /// Device IP address (for --snmpv3-test) + #[arg(long, required_if_eq("snmpv3_test", "true"))] + snmpv3_ip: Option, + + /// SNMPv3 username (for --snmpv3-test) + #[arg(long, default_value = "")] + snmpv3_user: String, + + /// SNMPv3 auth password (for --snmpv3-test) + #[arg(long, default_value = "")] + snmpv3_auth_pass: String, + + /// SNMPv3 priv password (for --snmpv3-test) + #[arg(long, default_value = "")] + snmpv3_priv_pass: String, } -#[tokio::main] -async fn main() { +fn main() { // Initialize logging init_logger(); + // Build Tokio runtime with larger stack size for SNMPv3 crypto operations + let runtime = tokio::runtime::Builder::new_multi_thread() + .enable_all() + .thread_stack_size(8 * 1024 * 1024) // 8MB stack (default is 2MB) + .build() + .expect("Failed to build Tokio runtime"); + + runtime.block_on(async_main()) +} + +async fn async_main() { let args = Args::parse(); // Handle MikroTik test mode @@ -107,6 +137,12 @@ async fn main() { return; } + // Handle SNMPv3 test mode + if args.snmpv3_test { + run_snmpv3_test(&args).await; + return; + } + tracing::info!("Towerops agent starting"); // Check for newer Docker image version @@ -217,6 +253,104 @@ async fn main() { tracing::info!("Towerops agent stopped"); } +/// Run SNMPv3 test +async fn run_snmpv3_test(args: &Args) { + use snmp::V3Config; + + let ip = args.snmpv3_ip.as_ref().expect("--snmpv3-ip required"); + let username = &args.snmpv3_user; + let auth_pass = &args.snmpv3_auth_pass; + let priv_pass = &args.snmpv3_priv_pass; + + println!("Testing SNMPv3 device at {}...", ip); + println!(" Username: {}", username); + println!( + " Auth Password: {}", + if auth_pass.is_empty() { + "(empty)" + } else { + "(set)" + } + ); + println!( + " Priv Password: {}", + if priv_pass.is_empty() { + "(empty)" + } else { + "(set)" + } + ); + + let v3_config = V3Config { + username: username.clone(), + auth_password: if !auth_pass.is_empty() { + Some(auth_pass.clone()) + } else { + None + }, + priv_password: if !priv_pass.is_empty() { + Some(priv_pass.clone()) + } else { + None + }, + auth_protocol: Some("SHA".to_string()), + priv_protocol: Some("AES".to_string()), + security_level: "authPriv".to_string(), + }; + + let snmp_client = snmp::SnmpClient::new(); + + println!("\nTest 1: Get sysDescr.0 (1.3.6.1.2.1.1.1.0)"); + match snmp_client + .get( + ip, + "", + "3", + 161, + "1.3.6.1.2.1.1.1.0", + Some(v3_config.clone()), + ) + .await + { + Ok(value) => println!(" Result: {:?}", value), + Err(e) => println!(" Error: {}", e), + } + + println!("\nTest 2: Get sysUpTime.0 (1.3.6.1.2.1.1.3.0)"); + match snmp_client + .get( + ip, + "", + "3", + 161, + "1.3.6.1.2.1.1.3.0", + Some(v3_config.clone()), + ) + .await + { + Ok(value) => println!(" Result: {:?}", value), + Err(e) => println!(" Error: {}", e), + } + + println!("\nTest 3: Walk interfaces (1.3.6.1.2.1.2.2.1)"); + match snmp_client + .walk( + ip, + "", + "3", + 161, + "1.3.6.1.2.1.2.2.1", + Some(v3_config.clone()), + ) + .await + { + Ok(values) => println!(" Found {} values", values.len()), + Err(e) => println!(" Error: {}", e), + } + + println!("\nTest complete."); +} + /// Run MikroTik API test async fn run_mikrotik_test(args: &Args) { use mikrotik::{MikrotikClient, SecretString}; diff --git a/src/snmp/client.rs b/src/snmp/client.rs index 4f12100..39bfab4 100644 --- a/src/snmp/client.rs +++ b/src/snmp/client.rs @@ -60,8 +60,16 @@ impl SnmpClient { create_v1v2c_session(&addr, &community, version_num)? }; - // Perform GET request - let response = session.get(&oid_parsed).map_err(map_snmp_error)?; + // Perform GET request (with retry for v3 engine ID discovery) + let mut response = match session.get(&oid_parsed) { + Ok(resp) => resp, + Err(snmp2::Error::AuthUpdated) => { + tracing::debug!("SNMPv3 engine ID discovered, retrying request"); + // Retry after engine ID discovery + session.get(&oid_parsed).map_err(map_snmp_error)? + } + Err(e) => return Err(map_snmp_error(e)), + }; // Check for error status if response.error_status != 0 { @@ -72,7 +80,7 @@ impl SnmpClient { } // Extract first varbind - for (_name, value) in response.varbinds { + if let Some((_name, value)) = response.varbinds.next() { return convert_value(value); } @@ -129,7 +137,16 @@ impl SnmpClient { // Extract data we need from the response immediately let (error_status, varbind_data) = { - let response = session.getnext(&oid_to_query).map_err(map_snmp_error)?; + // Perform GETNEXT request (with retry for v3 engine ID discovery) + let response = match session.getnext(&oid_to_query) { + Ok(resp) => resp, + Err(snmp2::Error::AuthUpdated) => { + tracing::debug!("SNMPv3 engine ID discovered, retrying getnext"); + // Retry after engine ID discovery + session.getnext(&oid_to_query).map_err(map_snmp_error)? + } + Err(e) => return Err(map_snmp_error(e)), + }; let status = response.error_status; // Collect all varbind data as owned strings/values immediately @@ -322,8 +339,18 @@ fn create_v3_session(addr: &str, config: &V3Config) -> SnmpResult { let timeout = Some(Duration::from_secs(SNMP_TIMEOUT_SECS)); let req_id = 1; - SyncSession::new_v3(addr, timeout, req_id, security) - .map_err(|e| SnmpError::RequestFailed(format!("SNMPv3 session creation failed: {:?}", e))) + let mut session = SyncSession::new_v3(addr, timeout, req_id, security).map_err(|e| { + SnmpError::RequestFailed(format!("SNMPv3 session creation failed: {:?}", e)) + })?; + + // For authPriv/authNoPriv, perform engine ID discovery using session.init() + if needs_auth_protocol { + session.init().map_err(|e| { + SnmpError::RequestFailed(format!("Engine ID discovery failed: {:?}", e)) + })?; + } + + Ok(session) } /// Convert snmp2 crate's Value to our SnmpValue diff --git a/src/snmp/device_poller.rs b/src/snmp/device_poller.rs new file mode 100644 index 0000000..1b777c7 --- /dev/null +++ b/src/snmp/device_poller.rs @@ -0,0 +1,440 @@ +use super::types::{SnmpError, SnmpResult, SnmpValue}; +use super::V3Config; +use snmp2::{Oid, SyncSession}; +use std::str::FromStr; +use std::time::Duration; +use tokio::sync::{mpsc, oneshot}; + +const SNMP_TIMEOUT_SECS: u64 = 30; + +/// Request to perform an SNMP operation +#[derive(Debug)] +pub enum SnmpRequest { + Get { + oid: String, + response_tx: oneshot::Sender>, + }, + Walk { + base_oid: String, + response_tx: oneshot::Sender>>, + }, + Shutdown, +} + +/// Configuration for a device poller +#[derive(Clone, Debug)] +pub struct DeviceConfig { + pub ip: String, + pub port: u16, + pub version: String, + pub community: String, + pub v3_config: Option, +} + +/// Per-device polling thread that maintains a persistent SNMP session +pub struct DevicePoller { + device_id: String, + config: DeviceConfig, + request_tx: mpsc::UnboundedSender, +} + +impl DevicePoller { + /// Spawn a new device poller thread + pub fn spawn(device_id: String, config: DeviceConfig) -> Self { + let (request_tx, request_rx) = mpsc::unbounded_channel(); + + let device_id_clone = device_id.clone(); + let config_clone = config.clone(); + + // Spawn the polling thread with 8MB stack for SNMPv3 crypto operations + std::thread::Builder::new() + .name(format!("poller-{}", device_id)) + .stack_size(8 * 1024 * 1024) // 8MB stack (default is 2MB) + .spawn(move || { + if let Err(e) = run_poller_thread(device_id_clone, config_clone, request_rx) { + tracing::error!("Device poller thread failed: {}", e); + } + }) + .expect("Failed to spawn device poller thread"); + + tracing::info!("Spawned device poller thread for {}", device_id); + + Self { + device_id, + config, + request_tx, + } + } + + /// Send a GET request to the poller thread + pub async fn get(&self, oid: String) -> SnmpResult { + let (response_tx, response_rx) = oneshot::channel(); + + self.request_tx + .send(SnmpRequest::Get { oid, response_tx }) + .map_err(|_| SnmpError::RequestFailed("Poller thread died".into()))?; + + response_rx + .await + .map_err(|_| SnmpError::RequestFailed("Poller thread didn't respond".into()))? + } + + /// Send a WALK request to the poller thread + pub async fn walk(&self, base_oid: String) -> SnmpResult> { + let (response_tx, response_rx) = oneshot::channel(); + + self.request_tx + .send(SnmpRequest::Walk { + base_oid, + response_tx, + }) + .map_err(|_| SnmpError::RequestFailed("Poller thread died".into()))?; + + response_rx + .await + .map_err(|_| SnmpError::RequestFailed("Poller thread didn't respond".into()))? + } + + /// Shutdown the poller thread + pub fn shutdown(&self) { + let _ = self.request_tx.send(SnmpRequest::Shutdown); + } + + /// Get the device ID + pub fn device_id(&self) -> &str { + &self.device_id + } + + /// Get the device config + pub fn config(&self) -> &DeviceConfig { + &self.config + } + + /// Log device poller status using accessor methods + pub fn log_status(&self) { + let id = self.device_id(); + let cfg = self.config(); + tracing::debug!( + "Device poller {} at {}:{} (version: {})", + id, + cfg.ip, + cfg.port, + cfg.version + ); + } +} + +/// Run the device poller thread (blocking) +fn run_poller_thread( + device_id: String, + config: DeviceConfig, + mut request_rx: mpsc::UnboundedReceiver, +) -> Result<(), String> { + tracing::info!( + "Device poller thread started for {} at {}:{}", + device_id, + config.ip, + config.port + ); + + // Create persistent session + let addr = format!("{}:{}", config.ip, config.port); + let mut session = create_session(&addr, &config)?; + + tracing::info!( + "Created persistent SNMP session for device {} (version: {})", + device_id, + config.version + ); + + // Process requests until shutdown + while let Some(request) = request_rx.blocking_recv() { + match request { + SnmpRequest::Get { oid, response_tx } => { + let result = perform_get(&mut session, &oid); + let _ = response_tx.send(result); + } + SnmpRequest::Walk { + base_oid, + response_tx, + } => { + let result = perform_walk(&mut session, &base_oid); + let _ = response_tx.send(result); + } + SnmpRequest::Shutdown => { + tracing::info!("Device poller thread shutting down for {}", device_id); + break; + } + } + } + + tracing::info!("Device poller thread stopped for {}", device_id); + Ok(()) +} + +/// Create an SNMP session based on version +fn create_session(addr: &str, config: &DeviceConfig) -> Result { + let timeout = Some(Duration::from_secs(SNMP_TIMEOUT_SECS)); + let version_num = parse_snmp_version(&config.version)?; + + if config.version == "3" { + let v3_config = config + .v3_config + .as_ref() + .ok_or("v3_config required for SNMPv3")?; + create_v3_session(addr, timeout, v3_config) + } else { + create_v1v2c_session(addr, config.community.as_bytes(), timeout, version_num) + } +} + +/// Create v1/v2c session +fn create_v1v2c_session( + addr: &str, + community: &[u8], + timeout: Option, + version: i32, +) -> Result { + let req_id = 1; + + let result = match version { + 0 => SyncSession::new_v1(addr, community, timeout, req_id), + 1 => SyncSession::new_v2c(addr, community, timeout, req_id), + _ => return Err(format!("Unsupported SNMP version: {}", version)), + }; + + result.map_err(|e| format!("Failed to create v1/v2c session: {:?}", e)) +} + +/// Create v3 session +fn create_v3_session( + addr: &str, + timeout: Option, + config: &V3Config, +) -> Result { + use snmp2::v3::{Auth, Security}; + + // Parse security level and build Auth enum + let auth = match config.security_level.trim().to_lowercase().as_str() { + "noauthnopriv" | "" => Auth::NoAuthNoPriv, + "authnopriv" => Auth::AuthNoPriv, + "authpriv" => { + let cipher = parse_priv_protocol( + config + .priv_protocol + .as_deref() + .ok_or("Priv protocol required for authPriv")?, + )?; + let priv_pass = config + .priv_password + .as_deref() + .ok_or("Priv password required for authPriv")?; + + Auth::AuthPriv { + cipher, + privacy_password: priv_pass.as_bytes().to_vec(), + } + } + _ => { + return Err(format!( + "Unsupported security level: '{}'", + config.security_level + )) + } + }; + + let username = config.username.as_bytes(); + let auth_password = config.auth_password.as_deref().unwrap_or("").as_bytes(); + let needs_auth_protocol = !matches!(auth, Auth::NoAuthNoPriv); + + let mut security = Security::new(username, auth_password).with_auth(auth); + + if needs_auth_protocol { + let auth_proto = parse_auth_protocol(config.auth_protocol.as_deref().unwrap())?; + security = security.with_auth_protocol(auth_proto); + } + + let req_id = 1; + let mut session = SyncSession::new_v3(addr, timeout, req_id, security) + .map_err(|e| format!("Failed to create v3 session: {:?}", e))?; + + // For authPriv/authNoPriv, perform engine ID discovery using session.init() + if needs_auth_protocol { + session + .init() + .map_err(|e| format!("Engine ID discovery failed: {:?}", e))?; + } + + Ok(session) +} + +/// Perform SNMP GET on existing session (with retry for v3 engine ID discovery) +fn perform_get(session: &mut SyncSession, oid: &str) -> SnmpResult { + let oid_parsed = + Oid::from_str(oid).map_err(|_| SnmpError::InvalidOid(format!("Invalid OID: {}", oid)))?; + + // First attempt (may fail with AuthUpdated for v3 engine ID discovery) + let mut response = match session.get(&oid_parsed) { + Ok(resp) => resp, + Err(snmp2::Error::AuthUpdated) => { + tracing::debug!("SNMPv3 engine ID discovered, retrying request"); + // Retry after engine ID discovery + session.get(&oid_parsed).map_err(map_snmp_error)? + } + Err(e) => return Err(map_snmp_error(e)), + }; + + if response.error_status != 0 { + return Err(SnmpError::RequestFailed(format!( + "SNMP error status: {}", + response.error_status + ))); + } + + let (_, value) = response + .varbinds + .next() + .ok_or(SnmpError::RequestFailed("No varbinds in response".into()))?; + + convert_value(value) +} + +/// Perform SNMP WALK on existing session (with retry for v3 engine ID discovery) +fn perform_walk(session: &mut SyncSession, base_oid: &str) -> SnmpResult> { + let base_oid_parsed = Oid::from_str(base_oid) + .map_err(|_| SnmpError::InvalidOid(format!("Invalid OID: {}", base_oid)))?; + let base_oid_string = base_oid.to_string(); + + let mut results = Vec::new(); + let mut current_oid = base_oid_parsed; + let mut first_request = true; + + loop { + let oid_to_query = current_oid.clone(); + + let (error_status, varbind_data) = { + // First request may fail with AuthUpdated for v3 engine ID discovery + let response = match session.getnext(&oid_to_query) { + Ok(resp) => resp, + Err(snmp2::Error::AuthUpdated) if first_request => { + tracing::debug!("SNMPv3 engine ID discovered, retrying WALK request"); + // Retry after engine ID discovery + session.getnext(&oid_to_query).map_err(map_snmp_error)? + } + Err(e) => return Err(map_snmp_error(e)), + }; + + // After first request, don't retry on AuthUpdated + first_request = false; + let status = response.error_status; + + let data: Vec<(String, snmp2::Value)> = response + .varbinds + .map(|(name, value)| (name.to_string(), value)) + .collect(); + + (status, data) + }; + + if error_status != 0 { + break; + } + + if varbind_data.is_empty() { + break; + } + + for (name_str, value) in varbind_data { + if !name_str.starts_with(&base_oid_string) { + return Ok(results); + } + + let converted_value = convert_value(value)?; + results.push((name_str.clone(), converted_value)); + + current_oid = Oid::from_str(&name_str).map_err(|_| { + SnmpError::InvalidOid(format!("Invalid OID from response: {}", name_str)) + })?; + } + } + + Ok(results) +} + +/// Parse SNMP version string to integer +fn parse_snmp_version(version: &str) -> Result { + match version.trim().to_lowercase().as_str() { + "1" | "v1" | "snmpv1" => Ok(0), + "2c" | "v2c" | "snmpv2c" | "2" | "v2" => Ok(1), + "3" | "v3" | "snmpv3" => Ok(3), + _ => Err(format!("Unsupported SNMP version: '{}'", version)), + } +} + +/// Parse authentication protocol +fn parse_auth_protocol(protocol: &str) -> Result { + use snmp2::v3::AuthProtocol; + + match protocol.trim().to_uppercase().as_str() { + "MD5" => Ok(AuthProtocol::Md5), + "SHA" | "SHA1" => Ok(AuthProtocol::Sha1), + "SHA224" => Ok(AuthProtocol::Sha224), + "SHA256" => Ok(AuthProtocol::Sha256), + "SHA384" => Ok(AuthProtocol::Sha384), + "SHA512" => Ok(AuthProtocol::Sha512), + _ => Err(format!("Unsupported auth protocol: '{}'", protocol)), + } +} + +/// Parse privacy protocol +fn parse_priv_protocol(protocol: &str) -> Result { + use snmp2::v3::Cipher; + + match protocol.trim().to_uppercase().as_str() { + "DES" => Ok(Cipher::Des), + "AES" | "AES128" => Ok(Cipher::Aes128), + "AES192" => Ok(Cipher::Aes192), + "AES256" | "AES-256" | "AES-256-C" => Ok(Cipher::Aes256), + _ => Err(format!("Unsupported priv protocol: '{}'", protocol)), + } +} + +/// Map snmp2 errors to our error type +fn map_snmp_error(err: snmp2::Error) -> SnmpError { + match &err { + snmp2::Error::AuthFailure(kind) => { + tracing::error!("SNMPv3 AuthFailure: {:?}", kind); + SnmpError::AuthFailure + } + snmp2::Error::AuthUpdated => { + tracing::info!("SNMPv3 engine ID discovered, caller should retry"); + SnmpError::RequestFailed("Authentication context updated, need to retry".into()) + } + snmp2::Error::CommunityMismatch => SnmpError::AuthFailure, + _ => { + tracing::error!("SNMP error: {:?}", err); + SnmpError::RequestFailed(format!("SNMP request failed: {:?}", err)) + } + } +} + +/// Convert snmp2::Value to our SnmpValue +fn convert_value(value: snmp2::Value) -> SnmpResult { + match value { + snmp2::Value::Integer(i) => Ok(SnmpValue::Integer(i)), + snmp2::Value::OctetString(bytes) => Ok(String::from_utf8(bytes.to_vec()) + .map(SnmpValue::String) + .unwrap_or_else(|_| SnmpValue::OctetString(bytes.to_vec()))), + snmp2::Value::ObjectIdentifier(oid) => Ok(SnmpValue::Oid(oid.to_string())), + snmp2::Value::Counter32(c) => Ok(SnmpValue::Counter32(c)), + snmp2::Value::Counter64(c) => Ok(SnmpValue::Counter64(c)), + snmp2::Value::Unsigned32(g) => Ok(SnmpValue::Gauge32(g)), + snmp2::Value::Timeticks(t) => Ok(SnmpValue::TimeTicks(t)), + snmp2::Value::IpAddress(ip) => Ok(SnmpValue::IpAddress(format!( + "{}.{}.{}.{}", + ip[0], ip[1], ip[2], ip[3] + ))), + snmp2::Value::Null => Ok(SnmpValue::Null), + _ => Ok(SnmpValue::Unsupported(format!("{:?}", value))), + } +} diff --git a/src/snmp/mod.rs b/src/snmp/mod.rs index 31d4807..a2305fa 100644 --- a/src/snmp/mod.rs +++ b/src/snmp/mod.rs @@ -1,7 +1,11 @@ mod client; +mod device_poller; +mod poller_registry; pub mod trap; mod types; pub use client::{SnmpClient, V3Config}; +pub use device_poller::DeviceConfig; +pub use poller_registry::PollerRegistry; pub use trap::{SnmpTrap, TrapListener, DEFAULT_TRAP_PORT}; pub use types::SnmpValue; diff --git a/src/snmp/poller_registry.rs b/src/snmp/poller_registry.rs new file mode 100644 index 0000000..024ba84 --- /dev/null +++ b/src/snmp/poller_registry.rs @@ -0,0 +1,146 @@ +use super::device_poller::{DeviceConfig, DevicePoller}; +use std::collections::HashMap; +use std::sync::{Arc, RwLock}; + +/// Registry of active device pollers +#[derive(Clone)] +pub struct PollerRegistry { + pollers: Arc>>>, +} + +impl PollerRegistry { + /// Create a new poller registry + pub fn new() -> Self { + Self { + pollers: Arc::new(RwLock::new(HashMap::new())), + } + } + + /// Get or create a device poller + pub fn get_or_create(&self, device_id: String, config: DeviceConfig) -> Arc { + // Try read lock first (fast path if poller exists) + { + let pollers = self.pollers.read().unwrap(); + if let Some(poller) = pollers.get(&device_id) { + poller.log_status(); + return Arc::clone(poller); + } + } + + // Need to create new poller (write lock) + let mut pollers = self.pollers.write().unwrap(); + + // Double-check in case another thread created it while we waited for write lock + if let Some(poller) = pollers.get(&device_id) { + poller.log_status(); + return Arc::clone(poller); + } + + // Create new poller + let poller = Arc::new(DevicePoller::spawn(device_id.clone(), config)); + pollers.insert(device_id, Arc::clone(&poller)); + + // Release write lock before logging + drop(pollers); + + tracing::info!("Created new device poller (total: {})", self.count()); + poller.log_status(); + + poller + } + + /// Remove a device poller (shutdown thread) + /// Called when a device is deleted or no longer needs polling + pub fn remove(&self, device_id: &str) { + let mut pollers = self.pollers.write().unwrap(); + if let Some(poller) = pollers.remove(device_id) { + poller.shutdown(); + tracing::info!( + "Removed device poller for {} (remaining: {})", + device_id, + pollers.len() + ); + } + } + + /// Get a list of active device IDs + pub fn list_devices(&self) -> Vec { + let pollers = self.pollers.read().unwrap(); + pollers.keys().cloned().collect() + } + + /// Get count of active pollers + pub fn count(&self) -> usize { + let pollers = self.pollers.read().unwrap(); + pollers.len() + } + + /// Shutdown all pollers + pub fn shutdown_all(&self) { + let device_list = self.list_devices(); + if !device_list.is_empty() { + tracing::info!("Shutting down {} device pollers", device_list.len()); + } + + let mut pollers = self.pollers.write().unwrap(); + for (device_id, poller) in pollers.drain() { + poller.shutdown(); + tracing::debug!("Shutdown device poller for {}", device_id); + } + } +} + +impl Default for PollerRegistry { + fn default() -> Self { + Self::new() + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::snmp::V3Config; + + #[test] + fn test_registry_remove() { + let registry = PollerRegistry::new(); + + // Create a test device config + let config = DeviceConfig { + ip: "127.0.0.1".to_string(), + port: 161, + version: "2c".to_string(), + community: "public".to_string(), + v3_config: None, + }; + + // Create a poller + let poller = registry.get_or_create("test-device".to_string(), config); + assert_eq!(registry.count(), 1); + assert_eq!(poller.device_id(), "test-device"); + + // Remove the poller + registry.remove("test-device"); + assert_eq!(registry.count(), 0); + } + + #[test] + fn test_device_poller_accessors() { + let config = DeviceConfig { + ip: "192.168.1.1".to_string(), + port: 161, + version: "2c".to_string(), + community: "public".to_string(), + v3_config: None, + }; + + let poller = DevicePoller::spawn("test-device".to_string(), config.clone()); + + // Test accessors + assert_eq!(poller.device_id(), "test-device"); + assert_eq!(poller.config().ip, "192.168.1.1"); + assert_eq!(poller.config().port, 161); + + poller.shutdown(); + } +} diff --git a/src/snmp/types.rs b/src/snmp/types.rs index 205d3fd..f4fb2db 100644 --- a/src/snmp/types.rs +++ b/src/snmp/types.rs @@ -29,11 +29,15 @@ pub type SnmpResult = Result; pub enum SnmpValue { Integer(i64), String(String), + OctetString(Vec), + Oid(String), Counter32(u32), Counter64(u64), Gauge32(u32), TimeTicks(u32), IpAddress(String), + Null, + Unsupported(String), } impl SnmpValue { diff --git a/src/websocket_client.rs b/src/websocket_client.rs index cb5a737..09f0b58 100644 --- a/src/websocket_client.rs +++ b/src/websocket_client.rs @@ -22,10 +22,10 @@ const CONNECTION_TIMEOUT: Duration = Duration::from_secs(30); type Result = std::result::Result>; use crate::proto::agent::{ - AgentHeartbeat, AgentJob, AgentJobList, JobType, MikrotikResult, MikrotikSentence, QueryType, - SnmpResult, + AgentHeartbeat, AgentJob, AgentJobList, CredentialTestResult, JobType, MikrotikResult, + MikrotikSentence, QueryType, SnmpResult, }; -use crate::snmp::{SnmpClient, SnmpValue}; +use crate::snmp::{DeviceConfig, PollerRegistry, SnmpValue}; /// Phoenix channel message format (JSON wrapper around binary protobuf). #[derive(Debug, serde::Serialize, serde::Deserialize)] @@ -45,6 +45,9 @@ pub struct AgentClient { result_rx: mpsc::UnboundedReceiver, mikrotik_result_tx: mpsc::UnboundedSender, mikrotik_result_rx: mpsc::UnboundedReceiver, + credential_test_tx: mpsc::UnboundedSender, + credential_test_rx: mpsc::UnboundedReceiver, + poller_registry: PollerRegistry, } impl AgentClient { @@ -93,6 +96,7 @@ impl AgentClient { let agent_id = generate_agent_id(); let (result_tx, result_rx) = mpsc::unbounded_channel(); let (mikrotik_result_tx, mikrotik_result_rx) = mpsc::unbounded_channel(); + let (credential_test_tx, credential_test_rx) = mpsc::unbounded_channel(); // Join Phoenix channel with token in payload let join_msg = PhoenixMessage { @@ -116,6 +120,9 @@ impl AgentClient { result_rx, mikrotik_result_tx, mikrotik_result_rx, + credential_test_tx, + credential_test_rx, + poller_registry: PollerRegistry::new(), }) } @@ -159,14 +166,17 @@ impl AgentClient { } Some(Ok(WsMessage::Close(_))) => { tracing::info!("Server closed connection"); + self.poller_registry.shutdown_all(); break Ok(()); } Some(Err(e)) => { tracing::error!("WebSocket error: {}", e); + self.poller_registry.shutdown_all(); break Err(e.into()); } None => { tracing::info!("Connection closed"); + self.poller_registry.shutdown_all(); break Ok(()); } _ => {} @@ -187,11 +197,23 @@ impl AgentClient { } } + // Receive credential test results from job tasks + Some(credential_test_result) = self.credential_test_rx.recv() => { + if let Err(e) = self.send_credential_test_result(credential_test_result).await { + tracing::error!("Error sending credential test result: {}", e); + } + } + // Send periodic heartbeats _ = heartbeat_interval.tick() => { if let Err(e) = self.send_heartbeat().await { tracing::error!("Error sending heartbeat: {}", e); } + // Log active poller count + let count = self.poller_registry.count(); + if count > 0 { + tracing::debug!("Active device pollers: {}", count); + } } } } @@ -242,6 +264,24 @@ impl AgentClient { async fn handle_jobs(&self, job_list: AgentJobList) -> Result<()> { tracing::info!("Received {} jobs from server", job_list.jobs.len()); + // Collect device IDs from current jobs + let mut current_device_ids = std::collections::HashSet::new(); + for job in &job_list.jobs { + current_device_ids.insert(job.device_id.clone()); + } + + // Clean up pollers for devices no longer in job list + let active_devices = self.poller_registry.list_devices(); + for device_id in active_devices { + if !current_device_ids.contains(&device_id) { + tracing::debug!( + "Removing poller for device no longer in job list: {}", + device_id + ); + self.poller_registry.remove(&device_id); + } + } + for job in job_list.jobs { let job_type = JobType::try_from(job.job_type).unwrap_or(JobType::Poll); tracing::info!("Executing job: {} (type: {:?})", job.job_id, job_type); @@ -256,11 +296,21 @@ impl AgentClient { } }); } + JobType::TestCredentials => { + // Execute credential test + let credential_test_tx = self.credential_test_tx.clone(); + tokio::spawn(async move { + if let Err(e) = execute_credential_test(job, credential_test_tx).await { + tracing::error!("Credential test execution failed: {}", e); + } + }); + } _ => { // Execute SNMP job (discovery or polling) let result_tx = self.result_tx.clone(); + let poller_registry = self.poller_registry.clone(); tokio::spawn(async move { - if let Err(e) = execute_snmp_job(job, result_tx).await { + if let Err(e) = execute_snmp_job(job, result_tx, poller_registry).await { tracing::error!("SNMP job execution failed: {}", e); } }); @@ -336,6 +386,28 @@ impl AgentClient { ); Ok(()) } + + /// Send credential test result to server. + async fn send_credential_test_result(&mut self, result: CredentialTestResult) -> Result<()> { + let binary = result.encode_to_vec(); + + let msg = PhoenixMessage { + topic: format!("agent:{}", self.agent_id), + event: "credential_test_result".to_string(), + payload: serde_json::json!({"binary": base64_encode(&binary)}), + reference: None, + }; + + let text = serde_json::to_string(&msg)?; + self.ws_stream.send(WsMessage::Text(text.into())).await?; + + tracing::info!( + "Sent credential test result (test_id: {}, success: {})", + result.test_id, + result.success + ); + Ok(()) + } } /// Redact SNMP community string for logging, showing first 2 chars only @@ -351,12 +423,13 @@ fn redact_community(community: &str) -> String { async fn execute_snmp_job( job: AgentJob, result_tx: mpsc::UnboundedSender, + poller_registry: PollerRegistry, ) -> Result<()> { let snmp_device = job.snmp_device.ok_or("Job missing SNMP device info")?; // Build v3 config if version is "3" let v3_config = if snmp_device.version == "3" { - Some(crate::snmp::V3Config { + let config = crate::snmp::V3Config { username: snmp_device.v3_username.clone(), auth_password: if !snmp_device.v3_auth_password.is_empty() { Some(snmp_device.v3_auth_password.clone()) @@ -379,7 +452,9 @@ async fn execute_snmp_job( None }, security_level: snmp_device.v3_security_level.clone(), - }) + }; + + Some(config) } else { None }; @@ -396,8 +471,18 @@ async fn execute_snmp_job( snmp_device.version ); + // Build device config and get or create persistent poller + let device_config = DeviceConfig { + ip: snmp_device.ip.clone(), + port: snmp_device.port as u16, + version: snmp_device.version.clone(), + community: snmp_device.community.clone(), + v3_config, + }; + + let poller = poller_registry.get_or_create(job.device_id.clone(), device_config); + let mut oid_values: HashMap = HashMap::new(); - let snmp_client = SnmpClient::new(); for query in job.queries { let query_type = QueryType::try_from(query.query_type).unwrap_or(QueryType::Get); @@ -406,17 +491,7 @@ async fn execute_snmp_job( QueryType::Get => { // Execute SNMP GET for each OID for oid in &query.oids { - match snmp_client - .get( - &snmp_device.ip, - &snmp_device.community, - &snmp_device.version, - snmp_device.port as u16, - oid, - v3_config.clone(), - ) - .await - { + match poller.get(oid.clone()).await { Ok(value) => { oid_values.insert(oid.clone(), value_to_string(value)); } @@ -438,17 +513,7 @@ async fn execute_snmp_job( QueryType::Walk => { // Execute SNMP WALK for each base OID for base_oid in &query.oids { - match snmp_client - .walk( - &snmp_device.ip, - &snmp_device.community, - &snmp_device.version, - snmp_device.port as u16, - base_oid, - v3_config.clone(), - ) - .await - { + match poller.walk(base_oid.clone()).await { Ok(results) => { for (oid, value) in results { oid_values.insert(oid, value_to_string(value)); @@ -500,6 +565,112 @@ async fn execute_snmp_job( Ok(()) } +/// Execute a credential test job. +/// +/// Tests SNMP credentials by performing a simple GET on sysDescr.0. +/// Returns success with system description or failure with error message. +async fn execute_credential_test( + job: AgentJob, + result_tx: mpsc::UnboundedSender, +) -> Result<()> { + let snmp_device = job.snmp_device.ok_or("Job missing SNMP device info")?; + + tracing::info!( + "Testing SNMP credentials for {}:{} (version: {})", + snmp_device.ip, + snmp_device.port, + snmp_device.version + ); + + let timestamp = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH)? + .as_secs() as i64; + + // Build v3 config if version is "3" + let v3_config = if snmp_device.version == "3" { + Some(crate::snmp::V3Config { + username: snmp_device.v3_username.clone(), + auth_password: if !snmp_device.v3_auth_password.is_empty() { + Some(snmp_device.v3_auth_password.clone()) + } else { + None + }, + priv_password: if !snmp_device.v3_priv_password.is_empty() { + Some(snmp_device.v3_priv_password.clone()) + } else { + None + }, + auth_protocol: if !snmp_device.v3_auth_protocol.is_empty() { + Some(snmp_device.v3_auth_protocol.clone()) + } else { + None + }, + priv_protocol: if !snmp_device.v3_priv_protocol.is_empty() { + Some(snmp_device.v3_priv_protocol.clone()) + } else { + None + }, + security_level: snmp_device.v3_security_level.clone(), + }) + } else { + None + }; + + // Create a temporary SNMP client for testing (don't use persistent poller) + let snmp_client = crate::snmp::SnmpClient::new(); + + // Test with sysDescr.0 (standard system description OID) + let test_oid = "1.3.6.1.2.1.1.1.0".to_string(); + + let result = match snmp_client + .get( + &snmp_device.ip, + &snmp_device.community, + &snmp_device.version, + snmp_device.port as u16, + &test_oid, + v3_config, + ) + .await + { + Ok(value) => { + let sys_descr = value_to_string(value); + tracing::info!("✓ Credential test succeeded: {}", sys_descr); + + CredentialTestResult { + test_id: job.job_id.clone(), + success: true, + error_message: String::new(), + system_description: sys_descr, + timestamp, + } + } + Err(e) => { + let error_msg = format!("SNMP test failed: {}", e); + tracing::warn!("✗ Credential test failed: {}", error_msg); + + CredentialTestResult { + test_id: job.job_id.clone(), + success: false, + error_message: error_msg, + system_description: String::new(), + timestamp, + } + } + }; + + // Send result back to main client task + if let Err(e) = result_tx.send(result) { + tracing::warn!( + "Failed to send credential test result for job {}: channel closed", + job.job_id + ); + return Err(format!("Result channel closed: {}", e).into()); + } + + Ok(()) +} + /// Execute a MikroTik API job and collect results. async fn execute_mikrotik_job( job: AgentJob, @@ -792,11 +963,22 @@ fn value_to_string(value: SnmpValue) -> String { match value { SnmpValue::Integer(i) => i.to_string(), SnmpValue::String(s) => s, + SnmpValue::OctetString(bytes) => { + // Convert to hex string for non-printable data + bytes + .iter() + .map(|b| format!("{:02x}", b)) + .collect::>() + .join(":") + } + SnmpValue::Oid(oid) => oid, SnmpValue::Counter32(c) => c.to_string(), SnmpValue::Counter64(c) => c.to_string(), SnmpValue::Gauge32(g) => g.to_string(), SnmpValue::TimeTicks(t) => t.to_string(), SnmpValue::IpAddress(ip) => ip, + SnmpValue::Null => "null".to_string(), + SnmpValue::Unsupported(s) => s, } }