use std::io; use std::io::{BufRead, BufReader, Write}; use std::net::{SocketAddr, TcpStream}; use crate::cmd::{Command, Response}; use crate::{ClientError, Config, Host, NutError}; /// A blocking NUT client connection. pub enum Connection { Tcp(TcpConnection), } impl Connection { pub fn new(config: Config) -> crate::Result { match &config.host { Host::Tcp(socket_addr) => { Ok(Self::Tcp(TcpConnection::new(config.clone(), socket_addr)?)) } } } pub fn list_ups(&mut self) -> crate::Result> { match self { Self::Tcp(conn) => conn.list_ups(), } } } /// A blocking TCP NUT client connection. #[derive(Debug)] pub struct TcpConnection { config: Config, tcp_stream: TcpStream, } impl TcpConnection { fn new(config: Config, socket_addr: &SocketAddr) -> crate::Result { // Create the TCP connection let tcp_stream = TcpStream::connect_timeout(socket_addr, config.timeout)?; let mut connection = Self { config, tcp_stream }; // Attempt login using `config.auth` connection.login()?; Ok(connection) } fn login(&mut self) -> crate::Result<()> { if let Some(auth) = &self.config.auth { // Pass username and check for 'OK' Self::write_cmd(&mut self.tcp_stream, Command::SetUsername(&auth.username))?; Self::read_response(&mut self.tcp_stream)?.expect_ok()?; // Pass password and check for 'OK' if let Some(password) = &auth.password { Self::write_cmd(&mut self.tcp_stream, Command::SetPassword(&password))?; Self::read_response(&mut self.tcp_stream)?.expect_ok()?; } } Ok(()) } fn list_ups(&mut self) -> crate::Result> { Self::write_cmd(&mut self.tcp_stream, Command::List(&["UPS"]))?; let list = Self::read_list(&mut self.tcp_stream, &["UPS"])?; Ok(list .into_iter() .map(|mut row| (row.remove(0), row.remove(0))) .collect()) } fn write_cmd(stream: &mut TcpStream, line: Command) -> crate::Result<()> { let line = format!("{}\n", line); stream.write_all(line.as_bytes())?; stream.flush()?; Ok(()) } fn parse_line(reader: &mut BufReader<&mut TcpStream>) -> crate::Result> { let mut raw = String::new(); reader.read_line(&mut raw)?; raw = raw[..raw.len() - 1].to_string(); // Strip off \n // Parse args by splitting whitespace, minding quotes for args with multiple words let args = shell_words::split(&raw) .map_err(|e| NutError::Generic(format!("Parsing server response failed: {}", e)))?; Ok(args) } fn read_response(stream: &mut TcpStream) -> crate::Result { let mut reader = io::BufReader::new(stream); let args = Self::parse_line(&mut reader)?; Response::from_args(args) } fn read_list(stream: &mut TcpStream, query: &[&str]) -> crate::Result>> { let mut reader = io::BufReader::new(stream); let args = Self::parse_line(&mut reader)?; Response::from_args(args)?.expect_begin_list(query)?; let mut lines: Vec> = Vec::new(); loop { let mut args = Self::parse_line(&mut reader)?; let resp = Response::from_args(args.clone()); if let Ok(resp) = resp { resp.expect_end_list(query)?; break; } else { let err = resp.unwrap_err(); if let ClientError::Nut(err) = err { if let NutError::UnknownResponseType(_) = err { // Likely an item entry, let's check... if args.len() < query.len() || &args[0..query.len()] != query { return Err(ClientError::Nut(err)); } else { let args = args.drain(query.len()..).collect(); lines.push(args); continue; } } else { return Err(ClientError::Nut(err)); } } else { return Err(err); } } } Ok(lines) } }