From 7ba49e4e4915e7155d03f6fe82880b5754d554b0 Mon Sep 17 00:00:00 2001 From: cool-mist Date: Wed, 8 Jul 2026 14:22:40 +0530 Subject: [PATCH] Add config file support --- src/config.rs | 247 ++++++++++++++++++++++++++++++++++++++++++++++++++ src/error.rs | 38 ++++++++ src/main.rs | 77 ++++++++++++---- 3 files changed, 342 insertions(+), 20 deletions(-) create mode 100644 src/config.rs create mode 100644 src/error.rs diff --git a/src/config.rs b/src/config.rs new file mode 100644 index 0000000..2dd6a82 --- /dev/null +++ b/src/config.rs @@ -0,0 +1,247 @@ +use std::{ + collections::HashMap, + fmt::Display, + fs::{self}, + mem::take, + path::PathBuf, +}; + +use crate::error::{self, AppError, EResult}; + +#[derive(Clone)] +pub struct Mac { + pub mac_bytes: [u8; 6], + original_string: String, +} + +impl Display for Mac { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "{}", self.original_string) + } +} + +impl Mac { + pub fn create(mac: &str) -> EResult { + let mut mac_bytes = [0u8; 6]; + let mut mac_bytes_index = 0; + let split: Vec<&str> = mac.split(':').collect(); + if split.len() != 6 { + return Err(AppError::NotAMac); + } + + for i in &split { + mac_bytes[mac_bytes_index] = u8::from_str_radix(i, 16).map_err(|e| { + AppError::ParseMacError(format!( + "'{} is not a valid mac string, failed to parse '{}' as hex due to {}", + &mac, &i, e + )) + })?; + mac_bytes_index += 1; + } + + Ok(Self { + mac_bytes, + original_string: mac.to_string(), + }) + } +} + +#[derive(Eq, PartialEq, Hash)] +pub struct Nick { + name: String, +} + +impl Nick { + pub fn create(nick: &str) -> EResult { + for b in nick.bytes() { + if !b.is_ascii_alphanumeric() { + return Err(AppError::NickNotAlphanumeric(format!( + "Nickname should be alphanumeric for a valid nick, Invalid char {}", + &b + ))); + } + } + + Ok(Self { + name: nick.to_string(), + }) + } +} + +pub enum Device { + Mac(Mac), + Nick(Nick), +} + +impl Device { + pub fn create(mac_or_nick: &str) -> EResult { + let as_mac = Mac::create(mac_or_nick); + let Err(AppError::NotAMac) = as_mac else { + return as_mac.map(|m| Device::mac(m.mac_bytes, mac_or_nick.to_string())); + }; + + let as_nick = Nick::create(mac_or_nick)?; + Ok(Device::Nick(as_nick)) + } + + fn mac(mac_bytes: [u8; 6], original_string: String) -> Self { + Self::Mac(Mac { + mac_bytes, + original_string, + }) + } +} + +pub struct Config { + nick_map: HashMap, +} + +impl Config { + pub fn create() -> EResult { + let mut config_parent = Config::get_parent()?; + config_parent.push("woly.conf"); + let config_file = config_parent; + if config_file.exists() { + return Config::read_existing_conf(config_file); + } + + std::fs::File::create_new(&config_file).map_err(|e| { + AppError::ConfigDirInitError(format!( + "Failed to create file at {} due to {}", + config_file.display(), + e + )) + })?; + + Ok(Config { + nick_map: HashMap::default(), + }) + } + + pub fn get_mac(&self, mac_or_nick: &str) -> EResult { + match Device::create(mac_or_nick)? { + Device::Mac(mac) => Ok(mac), + Device::Nick(nick) => self + .nick_map + .get(&nick) + .ok_or(AppError::NickNotMapped(format!( + "Nickname {} is not mapped to any MAC address in the config file", + &nick.name + ))) + .map(|mac| mac.clone()), + } + } + + fn read_existing_conf(config_file: PathBuf) -> EResult { + let bytes = fs::read(&config_file).map_err(|e| { + AppError::ConfigDirInitError(format!( + "Unable to read the contents of {} because {}", + config_file.display(), + e + )) + })?; + + // name=mac, one per line + // Eg: potato=32:1a:ff:3e:10:ef + let mut key = Vec::new(); + let mut value = Vec::new(); + let mut parse_state = 0; + let mut line_number = 1; + let mut nick_map = HashMap::new(); + for b in bytes.into_iter() { + if b == b'=' { + parse_state = 1; + continue; + } + + // TODO: Also consider \r and \r\n + if b == b'\n' { + let key_str = String::from_utf8(take(&mut key)) + .map_err(|e| error::map_config_error_utf8(&e, line_number))?; + let value_str = String::from_utf8(take(&mut value)) + .map_err(|e| error::map_config_error_utf8(&e, line_number))?; + let mac = Mac::create(&value_str) + .map_err(|e| error::map_config_error(&e, line_number))?; + + let nick = Nick { name: key_str }; + if nick_map.contains_key(&nick) { + return Err(AppError::ConfigFileFormatError(format!( + "Duplicate key {} at line number {}", + nick.name, line_number + ))); + } + nick_map.insert(nick, mac); + + parse_state = 0; + line_number += 1; + continue; + } + + match parse_state { + 0 => key.push(b), + _ => value.push(b), + } + } + + Ok(Config { nick_map }) + } + + fn get_parent() -> EResult { + let config_dir_env_var_name = "WOLY_CONFIG_HOME"; + let config_dir_name = std::env::var(config_dir_env_var_name); + if let Ok(config_dir_name) = config_dir_name { + let path = PathBuf::from(config_dir_name); + if !path.is_dir() { + return Err(AppError::InvalidConfigDir(format!( + "{} does not point to a valid directory {}", + config_dir_env_var_name, + path.display() + ))); + } + + return Ok(path); + } + + // Only works for linux + // Read $XDG_CONFIG_DIR, then fallback to $HOME/.config + if let Ok(xdg_config_dir) = std::env::var("XDG_CONFIG_HOME") { + return Config::validate_config_dir(PathBuf::from(xdg_config_dir), "$XDG_CONFIG_HOME"); + } + + if let Ok(home_dir) = std::env::var("HOME") { + let mut config_dir = PathBuf::from(home_dir); + config_dir.push(".config"); + return Config::validate_config_dir(config_dir, "$HOME/.config"); + } + + return Err(AppError::ConfigDirInitError(format!( + "Set one of {}, {} or {} environment variable to determine the config directory", + "$WOLY_CONFIG_HOME", "$XDG_CONFIG_HOME", "$HOME" + ))); + } + + fn validate_config_dir(config_dir: PathBuf, config_dir_source: &str) -> EResult { + if !config_dir.exists() { + fs::create_dir(&config_dir).map_err(|e| { + AppError::ConfigDirInitError(format!( + "Could not create the config directory at {} due to {}. Config source {}", + config_dir.display(), + e, + config_dir_source, + )) + })?; + + return Ok(config_dir); + } + + if !config_dir.is_dir() { + return Err(AppError::ConfigDirInitError(format!( + "config directory is not a valid directory: {}. Config source {}", + config_dir.display(), + config_dir_source, + ))); + } + + return Ok(config_dir); + } +} diff --git a/src/error.rs b/src/error.rs new file mode 100644 index 0000000..20b96ce --- /dev/null +++ b/src/error.rs @@ -0,0 +1,38 @@ +use std::{fmt::Display, string::FromUtf8Error}; + +pub type EResult = std::result::Result; + +#[derive(Debug)] +pub enum AppError { + ParseMacError(String), + NotAMac, + NickNotAlphanumeric(String), + NickNotMapped(String), + InvalidConfigDir(String), + ConfigDirInitError(String), + ConfigFileFormatError(String), +} + +impl Display for AppError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + let to_write = match self { + AppError::ParseMacError(s) => format!("{}", s), + AppError::NotAMac => "Input is not a valid MAC".to_string(), + AppError::NickNotAlphanumeric(s) => format!("{}", s), + AppError::NickNotMapped(s) => format!("{}", s), + AppError::InvalidConfigDir(s) => format!("{}", s), + AppError::ConfigDirInitError(s) => format!("{}", s), + AppError::ConfigFileFormatError(s) => format!("{}", s), + }; + + write!(f, "{}", to_write) + } +} + +pub fn map_config_error_utf8(error: &FromUtf8Error, line_number: usize) -> AppError { + AppError::ConfigFileFormatError(format!("On line {}, {}", line_number, error)) +} + +pub fn map_config_error(error: &AppError, line_number: usize) -> AppError { + AppError::ConfigFileFormatError(format!("On line {}, {}", line_number, error)) +} diff --git a/src/main.rs b/src/main.rs index 485a41f..2c11907 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,25 +1,22 @@ +mod config; +mod error; + use std::env::args; use std::net::UdpSocket; -fn get_magic_pkt(mac: &str, magic_bytes: &mut [u8; 102]) { - let mut mac_bytes = [0u8; 6]; - for (i, part) in mac.split(':').enumerate() { - mac_bytes[i] = u8::from_str_radix(part, 16) - .expect(&format!("'{}' is not a valid hex between 00-ff", part)); - } +use crate::config::{Config, Mac}; +use crate::error::EResult; - magic_bytes[..6].fill(0xFF); - for i in 0..16 { - let start = 6 + i * 6; - magic_bytes[start..start + 6].copy_from_slice(&mac_bytes); - } -} +fn main() -> EResult<()> { + let config = Config::create()?; -fn main() { - let mac = args() - .nth(1) - .expect("no mac given; Usage woly aa:bb:cc:dd:ee:ff"); + let arg_1 = args().nth(1); + let Some(arg_1) = arg_1 else { + print_usage(); + return Ok(()); + }; + let mac = config.get_mac(&arg_1)?; let mut magic_bytes = [0u8; 102]; get_magic_pkt(&mac, &mut magic_bytes); @@ -28,17 +25,54 @@ fn main() { socket .send_to(&magic_bytes, "255.255.255.255:9") .expect("couldn't send data"); - println!("Magic packet sent to 255.255.255.255:9 for Mac {:?}", mac); + println!("Magic packet sent to 255.255.255.255:9 for Mac {:}", mac); + Ok(()) +} + +fn get_magic_pkt(mac: &Mac, magic_bytes: &mut [u8; 102]) { + magic_bytes[..6].fill(0xFF); + for i in 0..16 { + let start = 6 + i * 6; + magic_bytes[start..start + 6].copy_from_slice(&mac.mac_bytes); + } +} + +fn print_usage() { + println!( + "{}", + r#"usage: woly ARG + +Options: + ARG Either the mac address of the device to send the wol magic + packet to, or a nickname specified in woly.conf to lookup + the mac. + Eg: woly aa:bb:cc:dd:ee:ff + Eg: woly potato + where potato is configured in the woly.conf file as + potato=aa:bb:cc:dd:ee:ff + +Config file: + Filename woly.conf + Location $WOLY_CONFIG_DIR/woly.conf, then $XDG_CONFIG_DIR/woly.conf, + then $HOME/woly.conf + Format one KEY=VALUE per line without spaces + Example potato=aa:bb:cc:dd:ee:ff + idly=12:34:23:e2:3d:ff +"# + ); } #[cfg(test)] mod tests { - use crate::get_magic_pkt; + use crate::{config::Mac, get_magic_pkt}; #[test] fn test_magic_bytes_for_mac() { let mut magic_bytes = [0u8; 102]; - get_magic_pkt("32:1a:ff:3e:10:ef", &mut magic_bytes); + let mac = Mac::create("32:1a:ff:3e:10:ef").expect("Test data is curated to be valid"); + get_magic_pkt(&mac, &mut magic_bytes); + + #[rustfmt::skip] let expected = [ 255, 255, 255, 255, 255, 255, 50, 26, 255, 62, 16, 239, @@ -59,6 +93,9 @@ mod tests { 50, 26, 255, 62, 16, 239, ]; - assert_eq!(magic_bytes, expected, "Magic bytes are not created correctly"); + assert_eq!( + magic_bytes, expected, + "Magic bytes are not created correctly" + ); } }