From a5403a5de3b95bf8ec51e96ca857fdff3160bb32 Mon Sep 17 00:00:00 2001 From: cool-mist Date: Sat, 11 Jul 2026 21:10:24 +0530 Subject: [PATCH] implement shutdown --- src/config.rs | 4 +- src/daemon.rs | 2 +- src/error.rs | 6 + src/mac.rs | 11 +- src/main.rs | 1 - src/netif.rs | 349 +++++++++++++++++++++++++++++------------ src/test_data/arp_file | 4 + src/wolz.rs | 27 +++- 8 files changed, 285 insertions(+), 119 deletions(-) create mode 100644 src/test_data/arp_file diff --git a/src/config.rs b/src/config.rs index d32a1de..1042551 100644 --- a/src/config.rs +++ b/src/config.rs @@ -7,7 +7,8 @@ use std::{ }; use crate::{ - error::{WError, WResult}, mac::{Device, Mac, Nick}, + error::{WError, WResult}, + mac::{Device, Mac, Nick}, }; pub struct Config { @@ -47,7 +48,6 @@ impl Config { &nick.name ))) .map(|mac| mac.clone()), - Device::IpV4(_) => todo!(), } } diff --git a/src/daemon.rs b/src/daemon.rs index b91fa38..fe01955 100644 --- a/src/daemon.rs +++ b/src/daemon.rs @@ -65,9 +65,9 @@ fn read_command(stream: &mut TcpStream) -> WResult { Ok(command) } +// TODO: Use DBus to send a message to systemd instead of spawning a new process fn build_shutdown_cmd() -> Command { let mut cmd = Command::new("systemctl"); - cmd.arg("-k"); cmd.arg("poweroff"); cmd diff --git a/src/error.rs b/src/error.rs index d97666b..15951ad 100644 --- a/src/error.rs +++ b/src/error.rs @@ -13,8 +13,11 @@ pub enum WError { ConfigDirInitError(String), ConfigFileFormatError(String), UdpSocketBindError(std::io::Error), + TcpSocketBindError(std::io::Error), + TcpSocketWriteError(std::io::Error), InvalidCommandString(String), FfiFailed(String), + MacNotInNetwork(String), DaemonSocketBindError(std::io::Error), DaemonConnectionReadError(std::io::Error), @@ -34,8 +37,11 @@ impl Display for WError { WError::ConfigDirInitError(s) => format!("{}", s), WError::ConfigFileFormatError(s) => format!("{}", s), WError::UdpSocketBindError(e) => format!("Couldn't bind and send udp packet because {}", e), + WError::TcpSocketBindError(e) => format!("Couldn't bind and send tcp packet because {}", e), + WError::TcpSocketWriteError(e) => format!("Couldn't write to TcpStream because {}", e), WError::InvalidCommandString(s) => format!("Command string '{}' is unknown", s), WError::FfiFailed(s) => format!("FFI invocation failed {}", s), + WError::MacNotInNetwork(s) => format!("Mac {} was not found in the local network", s), WError::DaemonCommandError(s) => format!("Failed to execute command {}", s), WError::DaemonCommandExecutionError(output) => format!( diff --git a/src/mac.rs b/src/mac.rs index b45da63..8094d93 100644 --- a/src/mac.rs +++ b/src/mac.rs @@ -40,6 +40,10 @@ impl Mac { magic_bytes[start..start + 6].copy_from_slice(&self.mac_bytes); } } + + pub(crate) fn is_same(&self, other: &Mac) -> bool { + self.mac_bytes == other.mac_bytes + } } impl Display for Mac { @@ -68,16 +72,9 @@ impl Nick { } } -pub(crate) struct IpV4 { - original_string: String, - ip: [u8; 4], - port: u32, -} - pub(crate) enum Device { Mac(Mac), Nick(Nick), - IpV4(IpV4), } impl Device { diff --git a/src/main.rs b/src/main.rs index fc900ec..397c401 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,4 +1,3 @@ -mod arp; mod config; mod daemon; mod error; diff --git a/src/netif.rs b/src/netif.rs index 23454f1..402c5eb 100644 --- a/src/netif.rs +++ b/src/netif.rs @@ -1,11 +1,196 @@ use std::collections::HashMap; use std::ffi::{CStr, c_uchar}; +use std::fs; +use std::mem::replace; use std::net::Ipv4Addr; use std::os::raw::{c_char, c_int, c_uint, c_ushort, c_void}; +use std::str::FromStr; use crate::error::{WError, WResult}; use crate::mac::Mac; +pub struct LocalNetwork { + // 192.168.0.7 + pub source_ip_hint: Ipv4Addr, + // 192.168.0.255 = source_ip_hint | ~netmask + pub subnet_broadcast: Ipv4Addr, +} + +pub fn get_local_networks() -> WResult> { + let mut networks = Vec::new(); + let interfaces = get_network_interfaces()?; + let mut unique_broadcast_addresses = Vec::new(); + for i in &interfaces { + if unique_broadcast_addresses.contains(&i.broadcast) { + continue; + } + + // man systemd.net-naming-scheme + // Only broadcast to local network + let iname = &i.name; + if iname.starts_with("en") || iname.starts_with("wl") || iname.starts_with("eth") { + unique_broadcast_addresses.push(i.broadcast); + let local_network = LocalNetwork { + source_ip_hint: i.ip, + subnet_broadcast: i.broadcast, + }; + networks.push(local_network); + } + } + + Ok(networks) +} + +pub fn find_ip_with_mac(mac: Mac) -> WResult { + let Some(ip) = find_ip_from_arp(&mac, "/proc/net/arp") else { + return Err(WError::MacNotInNetwork(mac.to_string())); + }; + + Ok(ip) +} + +fn find_ip_from_arp(mac: &Mac, arp_file_name: &str) -> Option { + let Ok(contents) = fs::read(arp_file_name) else { + eprintln!("ARP file not found at {}", arp_file_name); + return None; + }; + + // 0 = save word + // 1 = skip whitespace + // 4 = skip line + // Start with 4 to skip header + let mut state = 4; + let mut line: Vec = Vec::new(); + let mut current = Vec::new(); + for b in contents { + if state == 4 { + if b != b'\n' { + continue; + } + + if line.len() >= 4 { + let this_mac = Mac::create(&line[3]).expect(&format!( + "arp file is corrupted, tried to pase {} as a mac address", + line[3] + )); + + if this_mac.is_same(&mac) { + return Some(Ipv4Addr::from_str(&line[0]).expect(&format!( + "arp file is corrupted, tried to parse {} as an ip address", + line[0] + ))); + } + } + + line = Vec::new(); + current = Vec::new(); + state = 0; + continue; + } + + if state == 0 { + if b != b' ' { + current.push(b); + continue; + } + + let c = replace(&mut current, Vec::new()); + let word = String::from_utf8_lossy(&c).to_string(); + line.push(word); + + if line.len() == 4 { + state = 4; + } else { + state = 1; + } + + continue; + } + + if state == 1 { + if b == b' ' { + continue; + } + + current.push(b); + state = 0; + continue; + } + } + + None +} + +// via libc getifaddr // +fn get_network_interfaces() -> WResult> { + let mut network_addresses = HashMap::new(); + let mut link_addresses = HashMap::new(); + let mut addrs: *mut IfAddrs = std::ptr::null_mut(); + + unsafe { + let stat = getifaddrs(&mut addrs); + if stat != 0 { + return Err(WError::FfiFailed(format!( + "Failed to call libc:getifaddrs, returned {}", + stat + ))); + } + + let mut cur = addrs; + while !cur.is_null() { + let ifa = &*cur; + let name = CStr::from_ptr(ifa.ifa_name).to_string_lossy().to_string(); + if ifa.ifa_addr.is_null() { + cur = ifa.ifa_next; + continue; + } + + let family = (*ifa.ifa_addr).sa_family; + if family == AF_INET { + let sa_ip = ifa.ifa_addr as *const SockAddrIn; + let sa_netmask = ifa.ifa_netmask as *const SockAddrIn; + let ip = read_sa_ip(sa_ip); + let netmask = read_sa_ip(sa_netmask); + let broadcast = calculate_broadcast_address(ip, netmask); + let network_address = NetworkAddress { + source_ip_hint: ip, + broadcast, + }; + + network_addresses.insert(name, network_address); + } else if family == AF_PACKET { + let sock_addr_ll = ifa.ifa_addr as *const SockAddrLl; + let mac = read_sa_ll(sock_addr_ll); + + let link_address = LinkAddress { mac }; + + link_addresses.insert(name, link_address); + } + + cur = ifa.ifa_next; + } + + freeifaddrs(addrs); + } + + let mut network_interfaces = Vec::new(); + + // Assume all link_addresses have unique names for now + for (n, m) in link_addresses { + if let Some(n_addr) = network_addresses.get(&n) { + let network_interface = NetworkInterface::create( + n, + n_addr.source_ip_hint, + n_addr.broadcast, + m.mac, + ); + network_interfaces.push(network_interface); + }; + } + + Ok(network_interfaces) +} + #[repr(C)] struct IfAddrs { ifa_next: *mut IfAddrs, @@ -51,8 +236,7 @@ unsafe extern "C" { } struct NetworkAddress { - ip: Ipv4Addr, - netmask: Ipv4Addr, + source_ip_hint: Ipv4Addr, broadcast: Ipv4Addr, } @@ -63,119 +247,26 @@ struct LinkAddress { struct NetworkInterface { name: String, ip: Ipv4Addr, - _netmask: Ipv4Addr, broadcast: Ipv4Addr, _mac: [u8; 6], } impl NetworkInterface { - fn create(name: String, ip: Ipv4Addr, netmask: Ipv4Addr, broadcast: Ipv4Addr, mac: [u8; 6]) -> Self { + fn create( + name: String, + ip: Ipv4Addr, + broadcast: Ipv4Addr, + mac: [u8; 6], + ) -> Self { Self { name, ip, - _netmask: netmask, broadcast, _mac: mac, } } } -pub struct LocalNetwork { - pub source_ip_hint: Ipv4Addr, - pub subnet_broadcast: Ipv4Addr, -} - -pub fn get_local_networks() -> WResult> { - let mut networks = Vec::new(); - let interfaces = get_network_interfaces()?; - let mut unique_broadcast_addresses = Vec::new(); - for i in &interfaces { - if unique_broadcast_addresses.contains(&i.broadcast) { - continue; - } - - // man systemd.net-naming-scheme - // Only broadcast to local network - let iname = &i.name; - if iname.starts_with("en") || iname.starts_with("wl") || iname.starts_with("eth") { - unique_broadcast_addresses.push(i.broadcast); - let local_network = LocalNetwork { - source_ip_hint: i.ip, - subnet_broadcast: i.broadcast, - }; - networks.push(local_network); - } - } - - Ok(networks) -} - -pub(crate) fn find_local_networks_with_mac(mac: Mac) -> WResult> { - todo!() -} - -fn get_network_interfaces() -> WResult> { - let mut network_addresses = HashMap::new(); - let mut link_addresses = HashMap::new(); - let mut addrs: *mut IfAddrs = std::ptr::null_mut(); - - unsafe { - let stat = getifaddrs(&mut addrs); - if stat != 0 { - return Err(WError::FfiFailed(format!( - "Failed to call libc:getifaddrs, returned {}", - stat - ))); - } - - let mut cur = addrs; - while !cur.is_null() { - let ifa = &*cur; - let name = CStr::from_ptr(ifa.ifa_name).to_string_lossy().to_string(); - if ifa.ifa_addr.is_null() { - cur = ifa.ifa_next; - continue; - } - - let family = (*ifa.ifa_addr).sa_family; - if family == AF_INET { - let sa_ip = ifa.ifa_addr as *const SockAddrIn; - let sa_netmask = ifa.ifa_netmask as *const SockAddrIn; - let ip = read_sa_ip(sa_ip); - let netmask = read_sa_ip(sa_netmask); - let broadcast = calculate_broadcast_address(ip, netmask); - - let network_address = NetworkAddress { ip, netmask, broadcast }; - - network_addresses.insert(name, network_address); - } else if family == AF_PACKET { - let sock_addr_ll = ifa.ifa_addr as *const SockAddrLl; - let mac = read_sa_ll(sock_addr_ll); - - let link_address = LinkAddress { mac }; - - link_addresses.insert(name, link_address); - } - - cur = ifa.ifa_next; - } - - freeifaddrs(addrs); - } - - let mut network_interfaces = Vec::new(); - - // Assume all link_addresses have unique names for now - for (n, m) in link_addresses { - if let Some(n_addr) = network_addresses.get(&n) { - let network_interface = NetworkInterface::create(n, n_addr.ip, n_addr.netmask, n_addr.broadcast, m.mac); - network_interfaces.push(network_interface); - }; - } - - Ok(network_interfaces) -} - fn calculate_broadcast_address(ip: Ipv4Addr, netmask: Ipv4Addr) -> Ipv4Addr { let ip_octets = ip.octets(); let netmask_octets = netmask.octets(); @@ -202,3 +293,59 @@ fn read_sa_ll(ifa_addr: *const SockAddrLl) -> [u8; 6] { [mac[0], mac[1], mac[2], mac[3], mac[4], mac[5]] } } + +#[cfg(test)] +mod tests { + use std::{net::Ipv4Addr, str::FromStr}; + + use super::find_ip_from_arp; + use crate::mac::Mac; + + #[test] + fn arp_file_parse() { + validate_mapping("ae:a7:f1:0a:f2:ed", "192.168.0.193"); + validate_mapping("04:7c:16:6e:c0:8c", "192.168.0.8"); + validate_mapping("58:04:4f:67:b3:7c", "192.168.0.1"); + } + + #[test] + fn arp_file_parse_not_found() { + let mac = Mac::create("ae:a7:f1:0a:f2:ee").unwrap(); + let ip = find_ip_from_mac_against_test_data(&mac); + assert!(ip.is_none()); + } + + #[test] + fn get_local_networks_should_work() { + assert!(super::get_local_networks().is_ok(), "Check the FFI bindings code") + } + + #[test] + fn calculate_broadcast_address_correctly() { + validate_broadcast_address_calculation("192.168.0.8", "255.255.255.0", "192.168.0.255"); + validate_broadcast_address_calculation("192.168.0.8", "255.255.0.0", "192.168.255.255"); + validate_broadcast_address_calculation("192.168.1.8", "255.255.0.0", "192.168.255.255"); + validate_broadcast_address_calculation("192.168.1.8", "255.255.240.0", "192.168.15.255"); + } + + fn validate_mapping(mac: &str, expected_ip: &str) { + let mac = Mac::create(mac).unwrap(); + let ip = find_ip_from_mac_against_test_data(&mac) + .expect(&format!("IP mapping for {} is available in test data", &mac)); + + assert_eq!(expected_ip, ip.to_string()); + } + + fn find_ip_from_mac_against_test_data(mac: &Mac) -> Option { + let path = std::path::Path::new(env!("CARGO_MANIFEST_DIR")).join("src/test_data/arp_file"); + find_ip_from_arp(mac, &path.to_string_lossy().to_string()) + } + + fn validate_broadcast_address_calculation(ip: &str, netmask: &str, expected_broadcast: &str) { + let ip = Ipv4Addr::from_str(ip).unwrap(); + let netmask = Ipv4Addr::from_str(netmask).unwrap(); + let expected_broadcast = Ipv4Addr::from_str(expected_broadcast).unwrap(); + let actual_broadcast = super::calculate_broadcast_address(ip, netmask); + assert_eq!(expected_broadcast, actual_broadcast); + } +} diff --git a/src/test_data/arp_file b/src/test_data/arp_file new file mode 100644 index 0000000..6523b0e --- /dev/null +++ b/src/test_data/arp_file @@ -0,0 +1,4 @@ +IP address HW type Flags HW address Mask Device +192.168.0.8 0x1 0x2 04:7c:16:6e:c0:8c * wlan0 +192.168.0.1 0x1 0x2 58:04:4f:67:b3:7c * wlan0 +192.168.0.193 0x1 0x2 ae:a7:f1:0a:f2:ed * wlan0 diff --git a/src/wolz.rs b/src/wolz.rs index 9d1193e..8094a88 100644 --- a/src/wolz.rs +++ b/src/wolz.rs @@ -1,7 +1,9 @@ -use std::net::UdpSocket; +use std::{ + io::Write, + net::{TcpStream, UdpSocket}, +}; use crate::{ - arp, config::Config, error::{WError, WResult}, mac::Mac, @@ -55,10 +57,18 @@ impl Wolz { Ok(()) } - pub fn shut_down(&self, mac_or_nick: &str) -> WResult<()> { + // TODO: Fallback to full network scan if /proc/net/arp file lookup fails + fn shut_down(&self, mac_or_nick: &str) -> WResult<()> { let mac = self.config.get_mac(mac_or_nick)?; - let ip_from_arp = netif::find_local_networks_with_mac(mac)?; - todo!() + let ip = netif::find_ip_with_mac(mac)?; + + println!("Sending shutdown signal to {}:21367", ip); + let mut stream = TcpStream::connect(format!("{}:21367", ip)).map_err(WError::TcpSocketBindError)?; + stream + .write("shutdown".as_bytes()) + .map_err(WError::TcpSocketWriteError)?; + + Ok(()) } } @@ -67,10 +77,13 @@ fn try_send_magic_bytes(magic_bytes: &[u8; 102], mac: &Mac, local_network: &Loca let broadcast_addr = format!("{}:9", local_network.subnet_broadcast); match send_magic_bytes(&magic_bytes, &send_addr, &broadcast_addr) { Ok(()) => { - println!("Sent magic pkt {} -> {} for Mac {}", send_addr, broadcast_addr, mac); + println!( + "Sent WoL magic packet {} -> {} for mac {}", + send_addr, broadcast_addr, mac + ); } Err(e) => println!( - "Unable to send magic pkt {} -> {} for Mac {}, because {}", + "Unable to send WoL magic packet {} -> {} for mac {}, because {}", send_addr, broadcast_addr, mac, e ), }