diff --git a/appmgr/Cargo.lock b/appmgr/Cargo.lock index f7b465b72..071f1e9dd 100644 --- a/appmgr/Cargo.lock +++ b/appmgr/Cargo.lock @@ -802,6 +802,7 @@ dependencies = [ "http", "hyper-ws-listener", "indexmap", + "isocountry", "itertools 0.10.1", "jsonpath_lib", "lazy_static", @@ -1392,6 +1393,16 @@ version = "2.3.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "68f2d64f2edebec4ce84ad108148e67e1064789bee435edc5b60ad398714a3a9" +[[package]] +name = "isocountry" +version = "0.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1ea1dc4bf0fb4904ba83ffdb98af3d9c325274e92e6e295e4151e86c96363e04" +dependencies = [ + "serde", + "thiserror", +] + [[package]] name = "itertools" version = "0.8.2" diff --git a/appmgr/Cargo.toml b/appmgr/Cargo.toml index 26a325f36..33daeff09 100644 --- a/appmgr/Cargo.toml +++ b/appmgr/Cargo.toml @@ -65,6 +65,7 @@ hex = "0.4.3" http = "0.2.3" hyper-ws-listener = { git = "https://github.com/Start9Labs/hyper-ws-listener.git", branch = "main" } indexmap = { version = "1.6.2", features = ["serde"] } +isocountry = "0.3.2" itertools = "0.10.0" jsonpath_lib = "0.3.0" lazy_static = "1.4" @@ -94,10 +95,10 @@ serde_yaml = "0.8.14" sha2 = "0.9.3" simple-logging = "2.0" sqlx = { version = "0.5", features = [ + "chrono", + "offline", "runtime-tokio-rustls", "sqlite", - "offline", - "chrono", ] } tar = "0.4.35" thiserror = "1.0.24" diff --git a/appmgr/src/error.rs b/appmgr/src/error.rs index d046bf0f3..ddd22713b 100644 --- a/appmgr/src/error.rs +++ b/appmgr/src/error.rs @@ -51,6 +51,7 @@ pub enum ErrorKind { SoundError = 43, ParseTimestamp = 44, ParseSysInfo = 45, + WifiError = 46, } impl ErrorKind { pub fn as_str(&self) -> &'static str { @@ -101,6 +102,7 @@ impl ErrorKind { SoundError => "Sound Interface Error", ParseTimestamp => "Timestamp Parsing Error", ParseSysInfo => "System Info Parsing Error", + WifiError => "Wifi Internal Error", } } } diff --git a/appmgr/src/net/mod.rs b/appmgr/src/net/mod.rs index de5a2a36a..85279740e 100644 --- a/appmgr/src/net/mod.rs +++ b/appmgr/src/net/mod.rs @@ -13,6 +13,7 @@ pub mod interface; #[cfg(feature = "avahi")] pub mod mdns; pub mod tor; +pub mod wifi; pub struct NetController { tor: TorController, diff --git a/appmgr/src/net/wifi.rs b/appmgr/src/net/wifi.rs new file mode 100644 index 000000000..55100ba40 --- /dev/null +++ b/appmgr/src/net/wifi.rs @@ -0,0 +1,536 @@ +use std::collections::HashMap; +use std::time::Duration; + +use clap::ArgMatches; +use isocountry::CountryCode; +use rpc_toolkit::command; +use tokio::process::Command; + +use crate::context::EitherContext; +use crate::util::{display_none, display_serializable, Invoke, IoFormat}; +use crate::{Error, ErrorKind}; + +#[command(subcommands(add, connect, delete, get, set_country))] +pub async fn wifi(#[context] ctx: EitherContext) -> Result { + Ok(ctx) +} + +#[command(display(display_none))] +pub async fn add( + #[context] _ctx: EitherContext, + #[arg] ssid: String, + #[arg] password: String, + #[arg] priority: isize, + #[arg] connect: bool, +) -> Result<(), Error> { + let wpa_supplicant = WpaCli { interface: "wlan0" }; // TODO: pull from config + if !ssid.is_ascii() { + return Err(Error::new( + anyhow::anyhow!("SSID may not have special characters"), + ErrorKind::WifiError, + )); + } + if !password.is_ascii() { + return Err(Error::new( + anyhow::anyhow!("WiFi Password may not have special characters"), + ErrorKind::WifiError, + )); + } + async fn add_procedure<'a>( + wpa_supplicant: WpaCli<'a>, + ssid: &str, + password: &str, + priority: isize, + connect: bool, + ) -> Result<(), Error> { + log::info!("Adding new WiFi network: '{}'", ssid); + wpa_supplicant.add_network(ssid, password, priority).await?; + if connect { + let current = wpa_supplicant.get_current_network().await?; + let connected = wpa_supplicant.select_network(ssid).await?; + if !connected { + log::error!("Faild to add new WiFi network: '{}'", ssid); + wpa_supplicant.remove_network(ssid).await?; + match current { + None => {} + Some(current) => { + wpa_supplicant.select_network(¤t).await?; + } + } + } + } + Ok(()) + } + tokio::spawn(async move { + match add_procedure(wpa_supplicant, &ssid, &password, priority, connect).await { + Err(e) => { + log::error!("Failed to add new WiFi network '{}': {}", ssid, e); + } + Ok(_) => {} + } + }); + Ok(()) +} + +#[command(display(display_none))] +pub async fn connect(#[context] _ctx: EitherContext, #[arg] ssid: String) -> Result<(), Error> { + if !ssid.is_ascii() { + return Err(Error::new( + anyhow::anyhow!("SSID may not have special characters"), + ErrorKind::WifiError, + )); + } + async fn connect_procedure<'a>(wpa_supplicant: WpaCli<'a>, ssid: &String) -> Result<(), Error> { + let current = wpa_supplicant.get_current_network().await?; + let connected = wpa_supplicant.select_network(&ssid).await?; + if connected { + log::info!("Successfully connected to WiFi: '{}'", ssid); + } else { + log::error!("Failed to connect to WiFi: '{}'", ssid); + match current { + None => { + log::warn!("No WiFi to revert to!"); + } + Some(current) => { + wpa_supplicant.select_network(¤t).await?; + } + } + } + Ok(()) + } + let wpa_supplicant = WpaCli { interface: "wlan0" }; + tokio::spawn(async move { + match connect_procedure(wpa_supplicant, &ssid).await { + Err(e) => { + log::error!("Failed to connect to WiFi network '{}': {}", &ssid, e); + } + Ok(_) => {} + } + }); + Ok(()) +} + +#[command(display(display_none))] +pub async fn delete(#[context] _ctx: EitherContext, #[arg] ssid: String) -> Result<(), Error> { + if !ssid.is_ascii() { + return Err(Error::new( + anyhow::anyhow!("SSID may not have special characters"), + ErrorKind::WifiError, + )); + } + let wpa_supplicant = WpaCli { interface: "wlan0" }; + let current = wpa_supplicant.get_current_network().await?; + match current { + None => { + wpa_supplicant.remove_network(&ssid).await?; + } + Some(current) => { + if current == ssid { + if interface_connected("eth0").await? { + wpa_supplicant.remove_network(&ssid).await?; + } else { + return Err(Error::new(anyhow::anyhow!("Forbidden: Deleting this Network would make your Embassy Unreachable. Either connect to ethernet or connect to a different WiFi network to remedy this."), ErrorKind::WifiError)); + } + } + } + } + Ok(()) +} +#[derive(serde::Serialize, serde::Deserialize)] +#[serde(rename_all = "kebab-case")] +pub struct WiFiInfo { + ssids: Vec, + connected: Option, + country: CountryCode, + ethernet: bool, + signal_strength: Option, +} +fn display_wifi_info(info: WiFiInfo, matches: &ArgMatches<'_>) { + use prettytable::*; + + if matches.is_present("format") { + return display_serializable(info, matches); + } + + let mut table_global = Table::new(); + table_global.add_row(row![bc => + "CONNECTED", + "SIGNAL_STRENGTH", + "COUNTRY", + "ETHERNET", + ]); + table_global.add_row(row![ + &info + .connected + .as_ref() + .map_or("[N/A]".to_owned(), |c| format!("{}", c)), + &info + .signal_strength + .as_ref() + .map_or("[N/A]".to_owned(), |ss| format!("{}", ss)), + &format!("{}", info.country.alpha2()), + &format!("{}", info.ethernet) + ]); + table_global.print_tty(false); + + let mut table_ssids = Table::new(); + table_ssids.add_row(row![bc => "SSID",]); + for ssid in &info.ssids { + let mut row = row![ssid]; + if Some(ssid) == info.connected.as_ref() { + row.iter_mut() + .map(|c| { + c.style(Attr::ForegroundColor(match &info.signal_strength { + Some(100) => color::GREEN, + Some(0) => color::RED, + _ => color::YELLOW, + })) + }) + .collect::<()>() + } + table_ssids.add_row(row); + } + table_ssids.print_tty(false); +} + +#[command(display(display_wifi_info))] +pub async fn get( + #[context] _ctx: EitherContext, + #[allow(unused_variables)] + #[arg(long = "format")] + format: Option, +) -> Result { + let wpa_supplicant = WpaCli { interface: "wlan0" }; + let ssids_task = async { + Result::, Error>::Ok( + wpa_supplicant + .list_networks_low() + .await? + .into_keys() + .collect::>(), + ) + }; + let current_task = wpa_supplicant.get_current_network(); + let country_task = wpa_supplicant.get_country_low(); + let ethernet_task = interface_connected("eth0"); // TODO: pull from config + let rssi_task = wpa_supplicant.signal_poll_low(); + let (ssids_res, current_res, country_res, ethernet_res, rssi_res) = tokio::join!( + ssids_task, + current_task, + country_task, + ethernet_task, + rssi_task + ); + let current = current_res?; + let signal_strength = match rssi_res? { + None => None, + Some(x) if x <= -100 => Some(0 as usize), + Some(x) if x >= -50 => Some(100 as usize), + Some(x) => Some(2 * (x + 100) as usize), + }; + Ok(WiFiInfo { + ssids: ssids_res?, + connected: current, + country: country_res?, + ethernet: ethernet_res?, + signal_strength, + }) +} + +#[command(display(display_none))] +pub async fn set_country( + #[context] _ctx: EitherContext, + #[arg(parse(country_code_parse))] country: CountryCode, +) -> Result<(), Error> { + let wpa_supplicant = WpaCli { interface: "wlan0" }; + wpa_supplicant.set_country_low(country.alpha2()).await +} + +pub struct WpaCli<'a> { + interface: &'a str, +} +#[derive(Clone)] +pub struct NetworkId(String); +pub enum NetworkAttr { + Ssid(String), + Psk(String), + Priority(isize), + ScanSsid(bool), +} +impl std::fmt::Display for NetworkAttr { + fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result { + use NetworkAttr::*; + match self { + Ssid(s) => write!(f, "\"{}\"", s), + Psk(s) => write!(f, "\"{}\"", s), + Priority(n) => write!(f, "{}", n), + ScanSsid(b) => { + if *b { + write!(f, "1") + } else { + write!(f, "0") + } + } + } + } +} +impl<'a> WpaCli<'a> { + // Low Level + pub async fn add_network_low(&self) -> Result { + let r = Command::new("wpa_cli") + .arg("-i") + .arg(self.interface) + .arg("add_network") + .invoke(ErrorKind::WifiError) + .await?; + let s = std::str::from_utf8(&r)?; + Ok(NetworkId(s.trim().to_owned())) + } + pub async fn set_network_low(&self, id: &NetworkId, attr: &NetworkAttr) -> Result<(), Error> { + let _ = Command::new("wpa_cli") + .arg("-i") + .arg(self.interface) + .arg("set_network") + .arg(&id.0) + .arg(format!("{}", attr)) + .invoke(ErrorKind::WifiError) + .await?; + Ok(()) + } + pub async fn set_country_low(&self, country_code: &str) -> Result<(), Error> { + let _ = Command::new("wpa_cli") + .arg("-i") + .arg(self.interface) + .arg("set") + .arg("country") + .arg(country_code) + .invoke(ErrorKind::WifiError) + .await?; + Ok(()) + } + pub async fn get_country_low(&self) -> Result { + let r = Command::new("wpa_cli") + .arg("-i") + .arg(self.interface) + .arg("get") + .arg("country") + .invoke(ErrorKind::WifiError) + .await?; + Ok(CountryCode::for_alpha2(&String::from_utf8(r)?).unwrap()) + } + pub async fn enable_network_low(&self, id: &NetworkId) -> Result<(), Error> { + let _ = Command::new("wpa_cli") + .arg("-i") + .arg(self.interface) + .arg("enable_network") + .arg(&id.0) + .invoke(ErrorKind::WifiError) + .await?; + Ok(()) + } + pub async fn save_config_low(&self) -> Result<(), Error> { + let _ = Command::new("wpa_cli") + .arg("-i") + .arg(self.interface) + .arg("save_config") + .invoke(ErrorKind::WifiError) + .await?; + Ok(()) + } + pub async fn remove_network_low(&self, id: NetworkId) -> Result<(), Error> { + let _ = Command::new("wpa_cli") + .arg("-i") + .arg(self.interface) + .arg("remove_network") + .arg(&id.0) + .invoke(ErrorKind::WifiError) + .await?; + Ok(()) + } + pub async fn reconfigure_low(&self) -> Result<(), Error> { + let _ = Command::new("wpa_cli") + .arg("-i") + .arg(self.interface) + .arg("reconfigure") + .invoke(ErrorKind::WifiError) + .await?; + Ok(()) + } + pub async fn list_networks_low(&self) -> Result, Error> { + let r = Command::new("wpa_cli") + .arg("-i") + .arg(self.interface) + .arg("list_networks") + .invoke(ErrorKind::WifiError) + .await?; + Ok(String::from_utf8(r)? + .lines() + .skip(1) + .filter_map(|l| { + let mut cs = l.split("\t"); + let nid = NetworkId(cs.next()?.to_owned()); + let ssid = cs.next()?.to_owned(); + Some((ssid, nid)) + }) + .collect::>()) + } + pub async fn select_network_low(&self, id: &NetworkId) -> Result<(), Error> { + let _ = Command::new("wpa_cli") + .arg("-i") + .arg(self.interface) + .arg("select_network") + .arg(&id.0) + .invoke(ErrorKind::WifiError) + .await?; + Ok(()) + } + pub async fn new_password_low(&self, id: &NetworkId, pass: &str) -> Result<(), Error> { + let _ = Command::new("wpa_cli") + .arg("-i") + .arg(self.interface) + .arg("new_password") + .arg(&id.0) + .arg(pass) + .invoke(ErrorKind::WifiError) + .await?; + Ok(()) + } + pub async fn signal_poll_low(&self) -> Result, Error> { + let r = Command::new("wpa_cli") + .arg("-i") + .arg(self.interface) + .arg("signal_poll") + .invoke(ErrorKind::WifiError) + .await?; + let e = || { + Error::new( + anyhow::anyhow!("Invalid output from wpa_cli signal_poll"), + ErrorKind::WifiError, + ) + }; + let output = String::from_utf8(r)?; + Ok(if output.trim() == "FAIL" { + None + } else { + let l = output.lines().next().ok_or_else(e)?; + let rssi = l.split("=").nth(1).ok_or_else(e)?.parse()?; + Some(rssi) + }) + } + + // High Level + pub async fn check_network(&self, ssid: &str) -> Result, Error> { + Ok(self.list_networks_low().await?.remove(ssid)) + } + pub async fn select_network(&self, ssid: &str) -> Result { + let m_id = self.check_network(ssid).await?; + match m_id { + None => Err(Error::new( + anyhow::anyhow!("SSID Not Found"), + ErrorKind::WifiError, + )), + Some(x) => { + self.select_network_low(&x).await?; + self.save_config_low().await?; + let connect = async { + let mut current; + loop { + current = self.get_current_network().await; + match ¤t { + Ok(Some(_)) => { + break; + } + _ => {} + } + } + current + }; + let res = match tokio::time::timeout(Duration::from_secs(20), connect).await { + Err(_) => None, + Ok(net) => net?, + }; + Ok(match res { + None => false, + Some(net) => net == ssid, + }) + } + } + } + pub async fn get_current_network(&self) -> Result, Error> { + let r = Command::new("iwgetid") + .arg(self.interface) + .arg("--raw") + .invoke(ErrorKind::WifiError) + .await?; + let output = String::from_utf8(r)?; + if output.trim().is_empty() { + Ok(None) + } else { + Ok(Some(output)) + } + } + pub async fn remove_network(&self, ssid: &str) -> Result { + match self.check_network(ssid).await? { + None => Ok(false), + Some(x) => { + self.remove_network_low(x).await?; + self.save_config_low().await?; + self.reconfigure_low().await?; + Ok(true) + } + } + } + pub async fn add_network(&self, ssid: &str, psk: &str, priority: isize) -> Result<(), Error> { + use NetworkAttr::*; + let nid = match self.check_network(ssid).await? { + None => { + let nid = self.add_network_low().await?; + self.set_network_low(&nid, &Ssid(ssid.to_owned())).await?; + self.set_network_low(&nid, &Psk(psk.to_owned())).await?; + self.set_network_low(&nid, &Priority(priority)).await?; + self.set_network_low(&nid, &ScanSsid(true)).await?; + Result::::Ok(nid) + } + Some(nid) => { + self.new_password_low(&nid, psk).await?; + Ok(nid) + } + }?; + self.enable_network_low(&nid).await?; + self.save_config_low().await?; + Ok(()) + } +} + +pub async fn interface_connected(interface: &str) -> Result { + let out = Command::new("ifconfig") + .arg(interface) + .invoke(ErrorKind::WifiError) + .await?; + let v = std::str::from_utf8(&out)? + .lines() + .filter(|s| s.contains("inet")) + .next(); + Ok(!v.is_none()) +} + +pub fn country_code_parse(code: &str, _matches: &ArgMatches<'_>) -> Result { + CountryCode::for_alpha2(code).or(Err(Error::new( + anyhow::anyhow!("Invalid Country Code: {}", code), + ErrorKind::WifiError, + ))) +} + +#[tokio::test] +pub async fn test_interface_connected() { + println!("{}", interface_connected("wlp5s0").await.unwrap()); + println!("{}", interface_connected("enp4s0f1").await.unwrap()); +} + +#[tokio::test] +pub async fn test_signal_strength() { + let wpa = WpaCli { + interface: "wlp5s0", + }; + println!("{}", wpa.signal_poll_low().await.unwrap().unwrap()) +} diff --git a/appmgr/src/util.rs b/appmgr/src/util/mod.rs similarity index 99% rename from appmgr/src/util.rs rename to appmgr/src/util/mod.rs index eef8cbcac..fff8fa570 100644 --- a/appmgr/src/util.rs +++ b/appmgr/src/util/mod.rs @@ -230,8 +230,7 @@ impl Invoke for tokio::process::Command { crate::ensure_code!( res.status.success(), error_kind, - "{}: {}", - error_kind, + "{}", std::str::from_utf8(&res.stderr).unwrap_or("Unknown Error") ); Ok(res.stdout)