diff --git a/src/cli/mod.rs b/src/cli/mod.rs index 5015d3ec..47ce0615 100644 --- a/src/cli/mod.rs +++ b/src/cli/mod.rs @@ -96,6 +96,7 @@ pub fn get_boot_task_chain(&self) -> Task { #[cfg(test)] mod tests { use serial_test::serial; + use std::collections::HashSet; use crate::gui::pages::types::running_page::RunningPage; use crate::gui::pages::types::settings_page::SettingsPage; @@ -103,11 +104,13 @@ mod tests { use crate::gui::types::conf::Conf; use crate::gui::types::config_window::ConfigWindow; use crate::gui::types::export_pcap::ExportPcap; + use crate::gui::types::favorite::FavoriteKey; use crate::gui::types::filters::Filters; use crate::gui::types::settings::Settings; use crate::networking::types::capture_context::CaptureSourcePicklist; use crate::networking::types::config_device::ConfigDevice; use crate::networking::types::data_representation::DataRepr; + use crate::networking::types::service::Service; use crate::notifications::types::notifications::Notifications; use crate::report::types::sort_type::SortType; use crate::{Language, Sniffer, StyleType}; @@ -137,6 +140,7 @@ fn test_restore_default_configs() { }, style: StyleType::DraculaDark, ip_blacklist: "some-path".to_string(), + favorites: HashSet::from([FavoriteKey::Service(Service::Name("https"))]), }, device: ConfigDevice { device_name: "hey-hey".to_string(), diff --git a/src/countries/types/country.rs b/src/countries/types/country.rs index 21b6283e..bb28835e 100644 --- a/src/countries/types/country.rs +++ b/src/countries/types/country.rs @@ -1,7 +1,8 @@ +use serde::{Deserialize, Serialize}; use std::fmt; use std::fmt::Formatter; -#[derive(Debug, Copy, Clone, PartialEq, Eq, Hash, Default)] +#[derive(Debug, Copy, Clone, PartialEq, Eq, Hash, Default, Serialize, Deserialize)] pub enum Country { AD, AE, diff --git a/src/gui/pages/notifications_page.rs b/src/gui/pages/notifications_page.rs index 4753f9d2..72ca1a97 100644 --- a/src/gui/pages/notifications_page.rs +++ b/src/gui/pages/notifications_page.rs @@ -255,7 +255,7 @@ fn blacklisted_notification_log<'a>( let blacklisted_bar = item_bar( icon, host.to_blacklist_string(logged_notification.ip), - &data_info_host.data_info_fav.data_info, + &data_info_host.data_info, data_repr, first_entry_data_info, ); @@ -403,14 +403,13 @@ fn data_notification_extra<'a>( .first() .unwrap_or(&(Host::default(), DataInfoHost::default())) .1 - .data_info_fav .data_info; for (host, data_info_host) in &logged_notification.hosts { let icon = get_flag_tooltip(host.country, data_info_host, language, false); let host_bar = item_bar( icon, host.to_entry_string(), - &data_info_host.data_info_fav.data_info, + &data_info_host.data_info, logged_notification.data_repr, first_data_info, ); diff --git a/src/gui/pages/overview_page.rs b/src/gui/pages/overview_page.rs index 824de5a2..48b2486d 100644 --- a/src/gui/pages/overview_page.rs +++ b/src/gui/pages/overview_page.rs @@ -112,7 +112,7 @@ fn col_favorite_item( .unwrap_or_default(); for fi in &entries { - let star_button = fi.star_button(); + let star_button = fi.star_button(&sniffer.conf.settings.favorites); let icon = fi.icon(language, program_lookup, false); let item_bar = item_bar( diff --git a/src/gui/sniffer.rs b/src/gui/sniffer.rs index 197bf778..5d320766 100644 --- a/src/gui/sniffer.rs +++ b/src/gui/sniffer.rs @@ -66,7 +66,7 @@ use iced::{Element, Point, Size, Subscription, Task, window}; use listeners::Process; use rfd::FileHandle; -use std::collections::{HashMap, HashSet}; +use std::collections::HashMap; use std::net::IpAddr; use std::path::PathBuf; use std::sync::Arc; @@ -97,8 +97,6 @@ pub struct Sniffer { pub info_traffic: InfoTraffic, /// Map of the resolved addresses with their full rDNS value and the corresponding host pub addresses_resolved: HashMap, - /// Collection of the favorite hosts, services and programs - pub favorites: HashSet, /// Log of the displayed notifications, with the total number of notifications for this capture pub logged_notifications: LoggedNotifications, /// Reports if a newer release of the software is available on GitHub @@ -165,7 +163,6 @@ pub fn new(conf: Conf) -> Self { preview_captures_rx: None, info_traffic: InfoTraffic::default(), addresses_resolved: HashMap::new(), - favorites: HashSet::new(), logged_notifications: LoggedNotifications::default(), newer_release_available: None, capture_source, @@ -903,7 +900,7 @@ fn refresh_data(&mut self, mut msg: InfoTraffic, no_more_packets: bool) { &mut self.logged_notifications, &self.conf.settings.notifications, &msg, - &self.favorites, + &self.conf.settings.favorites, &self.capture_source, &self.addresses_resolved, ); @@ -1048,7 +1045,6 @@ fn reset(&mut self) -> Task { self.current_capture_rx = (self.current_capture_rx.0 + 1, None); self.info_traffic = InfoTraffic::default(); self.addresses_resolved = HashMap::new(); - self.favorites = HashSet::new(); self.logged_notifications = LoggedNotifications::default(); self.pcap_error = None; self.traffic_chart = TrafficChart::new(style, language, self.conf.data_repr); @@ -1101,28 +1097,9 @@ fn update_waiting_dots(&mut self) { fn add_or_remove_favorite(&mut self, fav: &FavoriteKey, add: bool) { if add { - self.favorites.insert(fav.clone()); + self.conf.settings.favorites.insert(fav.clone()); } else { - self.favorites.remove(fav); - } - - match fav { - FavoriteKey::Host(host) => { - if let Some(host_info) = self.info_traffic.hosts.get_mut(host) { - host_info.data_info_fav.is_favorite = add; - } - } - FavoriteKey::Service(service) => { - if let Some(service_info) = self.info_traffic.services.get_mut(service) { - service_info.is_favorite = add; - } - } - FavoriteKey::Program(program) => { - let Some(program_lookup) = &mut self.program_lookup else { - return; - }; - program_lookup.edit_fav(program, add); - } + self.conf.settings.favorites.remove(fav); } } @@ -1639,43 +1616,55 @@ fn test_add_or_remove_favorite() { // remove host sniffer.update(Message::AddOrRemoveFavorite(fav_host.clone(), false)); - assert_eq!(sniffer.favorites, HashSet::new()); + assert_eq!(sniffer.conf.settings.favorites, HashSet::new()); // remove service sniffer.update(Message::AddOrRemoveFavorite(fav_service.clone(), false)); - assert_eq!(sniffer.favorites, HashSet::new()); + assert_eq!(sniffer.conf.settings.favorites, HashSet::new()); // add service sniffer.update(Message::AddOrRemoveFavorite(fav_service.clone(), true)); - assert_eq!(sniffer.favorites, HashSet::from([fav_service.clone()])); + assert_eq!( + sniffer.conf.settings.favorites, + HashSet::from([fav_service.clone()]) + ); // remove host sniffer.update(Message::AddOrRemoveFavorite(fav_host.clone(), false)); - assert_eq!(sniffer.favorites, HashSet::from([fav_service.clone()])); + assert_eq!( + sniffer.conf.settings.favorites, + HashSet::from([fav_service.clone()]) + ); // add service sniffer.update(Message::AddOrRemoveFavorite(fav_service.clone(), true)); - assert_eq!(sniffer.favorites, HashSet::from([fav_service.clone()])); + assert_eq!( + sniffer.conf.settings.favorites, + HashSet::from([fav_service.clone()]) + ); // add host sniffer.update(Message::AddOrRemoveFavorite(fav_host.clone(), true)); assert_eq!( - sniffer.favorites, + sniffer.conf.settings.favorites, HashSet::from([fav_host.clone(), fav_service.clone()]) ); // add program sniffer.update(Message::AddOrRemoveFavorite(fav_program.clone(), true)); assert_eq!( - sniffer.favorites, + sniffer.conf.settings.favorites, HashSet::from([fav_host.clone(), fav_service.clone(), fav_program.clone()]) ); // remove service sniffer.update(Message::AddOrRemoveFavorite(fav_service.clone(), false)); assert_eq!( - sniffer.favorites, + sniffer.conf.settings.favorites, HashSet::from([fav_host.clone(), fav_program.clone()]) ); // remove program sniffer.update(Message::AddOrRemoveFavorite(fav_program.clone(), false)); - assert_eq!(sniffer.favorites, HashSet::from([fav_host.clone()])); + assert_eq!( + sniffer.conf.settings.favorites, + HashSet::from([fav_host.clone()]) + ); // remove host sniffer.update(Message::AddOrRemoveFavorite(fav_host.clone(), false)); - assert_eq!(sniffer.favorites, HashSet::new()); + assert_eq!(sniffer.conf.settings.favorites, HashSet::new()); } #[test] @@ -2139,6 +2128,10 @@ fn test_conf() { sniffer.update(Message::ChangeRunningPage(RunningPage::Notifications)); sniffer.update(Message::DataReprSelection(DataRepr::Bits)); sniffer.update(Message::LoadIpBlacklist("blacklist_file.csv".to_string())); + sniffer.update(Message::AddOrRemoveFavorite( + FavoriteKey::Service(Service::Name("https")), + true, + )); // force saving configs by quitting the app sniffer.welcome = Some((false, 0)); @@ -2167,6 +2160,7 @@ fn test_conf() { }, style: StyleType::DraculaDark, ip_blacklist: "blacklist_file.csv".to_string(), + favorites: HashSet::from([FavoriteKey::Service(Service::Name("https"))]), }, window: ConfigWindow::new((1000.0, 999.0), (-5.0, 277.5), (20.0, 20.0)), device: ConfigDevice::default(), diff --git a/src/gui/types/favorite.rs b/src/gui/types/favorite.rs index 04a6d2e9..83f0d053 100644 --- a/src/gui/types/favorite.rs +++ b/src/gui/types/favorite.rs @@ -6,7 +6,6 @@ use crate::gui::types::conf::Conf; use crate::gui::types::message::Message; use crate::networking::types::data_info::DataInfo; -use crate::networking::types::data_info_fav::DataInfoFav; use crate::networking::types::data_info_host::DataInfoHost; use crate::networking::types::data_representation::DataRepr; use crate::networking::types::host::Host; @@ -23,7 +22,9 @@ use crate::utils::types::icon::Icon; use iced::widget::{Button, Container, Space, button}; use iced::{Alignment, Element}; +use serde::{Deserialize, Serialize}; use std::cmp::min; +use std::collections::HashSet; #[derive(Debug, Clone, Copy, Eq, PartialEq)] pub enum Favorite { @@ -96,29 +97,25 @@ pub fn sort_arrows<'a>(self, conf: &Conf) -> Container<'a, Message, StyleType> { #[derive(Clone)] pub enum FavoriteItem { Host((Host, DataInfoHost)), - Service((Service, DataInfoFav)), - Program((Program, DataInfoFav)), + Service((Service, DataInfo)), + Program((Program, DataInfo)), } impl FavoriteItem { pub fn data_info(&self) -> DataInfo { match self { - FavoriteItem::Host((_, data_info_host)) => data_info_host.data_info_fav.data_info, - FavoriteItem::Service((_, data_info_fav)) - | FavoriteItem::Program((_, data_info_fav)) => data_info_fav.data_info, + FavoriteItem::Host((_, data_info_host)) => data_info_host.data_info, + FavoriteItem::Service((_, data_info)) | FavoriteItem::Program((_, data_info)) => { + *data_info + } } } - fn is_favorite(&self) -> bool { - match self { - FavoriteItem::Host((_, data_info_host)) => data_info_host.data_info_fav.is_favorite, - FavoriteItem::Service((_, data_info_fav)) - | FavoriteItem::Program((_, data_info_fav)) => data_info_fav.is_favorite, - } - } - - pub fn star_button<'a>(&self) -> Button<'a, Message, StyleType> { - let is_favorite = self.is_favorite(); + pub fn star_button<'a>( + &self, + favorites: &HashSet, + ) -> Button<'a, Message, StyleType> { + let is_favorite = favorites.contains(&FavoriteKey::from(self.clone())); let (icon, class) = if is_favorite { (Icon::StarFull, ButtonType::Starred) @@ -187,7 +184,7 @@ pub fn new_entry_search(&self) -> SearchParameters { } } -#[derive(Clone, Hash, Eq, PartialEq, Debug)] +#[derive(Clone, Hash, Eq, PartialEq, Debug, Serialize, Deserialize)] pub enum FavoriteKey { Host(Host), Service(Service), @@ -213,11 +210,7 @@ fn get_host_entries( ) -> Vec { let mut sorted_vec: Vec<(&Host, &DataInfoHost)> = info_traffic.hosts.iter().collect(); - sorted_vec.sort_by(|&(_, a), &(_, b)| { - a.data_info_fav - .data_info - .compare(&b.data_info_fav.data_info, sort_type, data_repr) - }); + sorted_vec.sort_by(|&(_, a), &(_, b)| a.data_info.compare(&b.data_info, sort_type, data_repr)); let n_entry = min(sorted_vec.len(), 30); sorted_vec[0..n_entry] @@ -233,13 +226,13 @@ fn get_service_entries( data_repr: DataRepr, sort_type: SortType, ) -> Vec { - let mut sorted_vec: Vec<(&Service, &DataInfoFav)> = info_traffic + let mut sorted_vec: Vec<(&Service, &DataInfo)> = info_traffic .services .iter() .filter(|(service, _)| service != &&Service::NotApplicable) .collect(); - sorted_vec.sort_by(|&(_, a), &(_, b)| a.data_info.compare(&b.data_info, sort_type, data_repr)); + sorted_vec.sort_by(|&(_, a), &(_, b)| a.compare(b, sort_type, data_repr)); let n_entry = min(sorted_vec.len(), 30); sorted_vec[0..n_entry] @@ -257,22 +250,22 @@ fn get_program_entries( return Vec::new(); }; - let mut sorted_vec: Vec<(&Program, &DataInfoFav)> = program_lookup + let mut sorted_vec: Vec<(&Program, &DataInfo)> = program_lookup .programs() .iter() // Unknown may be inserted, and then all of its data could be reassigned to known programs - .filter(|(_, d)| d.data_info.tot_data(DataRepr::Packets) > 0) + .filter(|(_, d)| d.tot_data(DataRepr::Packets) > 0) .collect(); sorted_vec.sort_by(|&(p1, a), &(p2, b)| { - if sort_type == SortType::Neutral && a.data_info.is_within_same_second(&b.data_info) { + if sort_type == SortType::Neutral && a.is_within_same_second(b) { if p1.is_unknown() { return std::cmp::Ordering::Greater; } else if p2.is_unknown() { return std::cmp::Ordering::Less; } } - a.data_info.compare(&b.data_info, sort_type, data_repr) + a.compare(b, sort_type, data_repr) }); let n_entry = min(sorted_vec.len(), 30); diff --git a/src/gui/types/settings.rs b/src/gui/types/settings.rs index 1ac3a65f..2f53c3f5 100644 --- a/src/gui/types/settings.rs +++ b/src/gui/types/settings.rs @@ -1,7 +1,9 @@ use serde::{Deserialize, Serialize}; +use std::collections::HashSet; use crate::gui::styles::types::gradient_type::GradientType; use crate::gui::types::conf::deserialize_or_default; +use crate::gui::types::favorite::FavoriteKey; use crate::notifications::types::notifications::Notifications; use crate::{Language, StyleType}; @@ -27,6 +29,8 @@ pub struct Settings { pub notifications: Notifications, #[serde(deserialize_with = "deserialize_or_default")] pub style: StyleType, + #[serde(deserialize_with = "deserialize_or_default")] + pub favorites: HashSet, } impl Default for Settings { @@ -41,6 +45,7 @@ fn default() -> Self { style_path: String::new(), notifications: Notifications::default(), style: StyleType::default(), + favorites: HashSet::new(), } } } diff --git a/src/networking/parse_packets.rs b/src/networking/parse_packets.rs index f10a407a..f23a3379 100644 --- a/src/networking/parse_packets.rs +++ b/src/networking/parse_packets.rs @@ -14,7 +14,6 @@ use crate::networking::types::bogon::is_bogon; use crate::networking::types::capture_context::{CaptureContext, CaptureSource, CaptureType}; use crate::networking::types::data_info::DataInfo; -use crate::networking::types::data_info_fav::DataInfoFav; use crate::networking::types::data_info_host::DataInfoHost; use crate::networking::types::host::{Host, HostMessage}; use crate::networking::types::icmp_type::IcmpType; @@ -244,7 +243,6 @@ pub fn parse_packets( .entry(host) .and_modify(|data_info_host| { data_info_host - .data_info_fav .data_info .add_packet(exchanged_bytes, traffic_direction); }) @@ -262,11 +260,10 @@ pub fn parse_packets( ); let is_bogon = is_bogon(&address_to_lookup); DataInfoHost { - data_info_fav: DataInfo::new_with_first_packet( + data_info: DataInfo::new_with_first_packet( exchanged_bytes, traffic_direction, - ) - .into(), + ), is_loopback, is_local, is_bogon, @@ -280,14 +277,11 @@ pub fn parse_packets( info_traffic_msg .services .entry(service) - .and_modify(|data_info_fav| { - data_info_fav - .data_info - .add_packet(exchanged_bytes, traffic_direction); + .and_modify(|data_info| { + data_info.add_packet(exchanged_bytes, traffic_direction); }) .or_insert_with(|| { DataInfo::new_with_first_packet(exchanged_bytes, traffic_direction) - .into() }); // update dropped packets number @@ -404,7 +398,7 @@ fn reverse_dns_lookups( }; let data_info_host = DataInfoHost { - data_info_fav: DataInfoFav::default(), + data_info: DataInfo::default(), is_local, is_bogon, is_loopback, @@ -459,7 +453,7 @@ fn new_hosts_to_send(&mut self) -> Vec { .remove(&address_to_lookup) .unwrap_or_default(); // overwrite the host message with the collected data - host_msg.data_info_host.data_info_fav.data_info = other_data; + host_msg.data_info_host.data_info = other_data; // insert the newly resolved host in the collection of resolved addresses self.addresses_resolved .insert(address_to_lookup, host_msg.host.clone()); diff --git a/src/networking/types/asn.rs b/src/networking/types/asn.rs index 37fd751b..b61d4702 100644 --- a/src/networking/types/asn.rs +++ b/src/networking/types/asn.rs @@ -1,5 +1,7 @@ +use serde::{Deserialize, Serialize}; + /// Struct to represent an Autonomous System -#[derive(Default, Clone, PartialEq, Eq, Hash, Debug)] +#[derive(Default, Clone, PartialEq, Eq, Hash, Debug, Serialize, Deserialize)] pub struct Asn { /// Autonomous System number pub code: String, diff --git a/src/networking/types/data_info_fav.rs b/src/networking/types/data_info_fav.rs deleted file mode 100644 index 64dd468c..00000000 --- a/src/networking/types/data_info_fav.rs +++ /dev/null @@ -1,17 +0,0 @@ -use crate::networking::types::data_info::DataInfo; - -#[derive(Clone, Copy, Debug, Eq, PartialEq, Hash, Default)] -pub struct DataInfoFav { - pub data_info: DataInfo, - /// Determine if this item is one of the favorites - pub is_favorite: bool, -} - -impl From for DataInfoFav { - fn from(data_info: DataInfo) -> Self { - Self { - data_info, - is_favorite: false, - } - } -} diff --git a/src/networking/types/data_info_host.rs b/src/networking/types/data_info_host.rs index d04001ae..30aceed8 100644 --- a/src/networking/types/data_info_host.rs +++ b/src/networking/types/data_info_host.rs @@ -1,13 +1,13 @@ //! Module defining the `DataInfoHost` struct related to hosts. -use crate::networking::types::data_info_fav::DataInfoFav; +use crate::networking::types::data_info::DataInfo; use crate::networking::types::traffic_type::TrafficType; /// Host-related information. #[derive(Clone, Copy, Default, Debug, Eq, PartialEq, Hash)] pub struct DataInfoHost { /// Incoming and outgoing packets and bytes - pub data_info_fav: DataInfoFav, + pub data_info: DataInfo, /// Determine if the connection is loopback (the "remote" is loopback) pub is_loopback: bool, /// Determine if the connection with this host is local @@ -20,9 +20,7 @@ pub struct DataInfoHost { impl DataInfoHost { pub fn refresh(&mut self, other: &Self) { - self.data_info_fav - .data_info - .refresh(other.data_info_fav.data_info); + self.data_info.refresh(other.data_info); self.is_loopback = other.is_loopback; self.is_local = other.is_local; self.is_bogon = other.is_bogon; diff --git a/src/networking/types/host.rs b/src/networking/types/host.rs index e4c34fb5..e2544aae 100644 --- a/src/networking/types/host.rs +++ b/src/networking/types/host.rs @@ -2,10 +2,11 @@ use crate::networking::types::asn::Asn; use crate::networking::types::data_info_host::DataInfoHost; use crate::utils::formatted_strings::clip_text; +use serde::{Deserialize, Serialize}; use std::net::IpAddr; /// Struct to represent a network host -#[derive(Default, PartialEq, Eq, Hash, Clone, Debug)] +#[derive(Default, PartialEq, Eq, Hash, Clone, Debug, Serialize, Deserialize)] pub struct Host { /// Hostname (domain). Obtained from the reverse DNS. pub domain: String, diff --git a/src/networking/types/info_traffic.rs b/src/networking/types/info_traffic.rs index 561205b6..6a9b8150 100644 --- a/src/networking/types/info_traffic.rs +++ b/src/networking/types/info_traffic.rs @@ -2,7 +2,6 @@ use crate::networking::manage_packets::get_local_port; use crate::networking::types::address_port_pair::AddressPortPair; use crate::networking::types::data_info::DataInfo; -use crate::networking::types::data_info_fav::DataInfoFav; use crate::networking::types::data_info_host::DataInfoHost; use crate::networking::types::data_representation::DataRepr; use crate::networking::types::host::Host; @@ -24,7 +23,7 @@ pub struct InfoTraffic { /// Map of the traffic pub map: HashMap, /// Map of the upper layer services with their data info - pub services: HashMap, + pub services: HashMap, /// Map of the hosts with their data info pub hosts: HashMap, } @@ -78,7 +77,7 @@ pub fn refresh(&mut self, msg: &mut Self, program_lookup_opt: &mut Option, picon_rx: Receiver<(String, IconHandle)>, state: HashMap<(u16, Protocol), LookedUpProgram>, - programs: HashMap, + programs: HashMap, picons: HashMap, } @@ -61,8 +60,8 @@ pub fn lookup_and_add_data( let res = Program::from_proc(proc.as_ref()); self.programs .entry(res.clone()) - .and_modify(|d| d.data_info.refresh(new_data)) - .or_insert(new_data.into()); + .and_modify(|d| d.refresh(new_data)) + .or_insert(new_data); res } @@ -161,8 +160,8 @@ pub fn update( // assign to known self.programs .entry(program) - .and_modify(|d| d.data_info.refresh(reassigned_data)) - .or_insert(reassigned_data.into()); + .and_modify(|d| d.refresh(reassigned_data)) + .or_insert(reassigned_data); // remove from Unknown // NOTE: subtracting reassigned_data from Unknown wouldn't correctly reassign final_instant, // so let's just reiterate through all the Unknown connections @@ -176,7 +175,7 @@ pub fn update( // or_insert not needed: Unknown is already in the map since reassigned data came from it self.programs .entry(Program::Unknown) - .and_modify(|d| d.data_info = unknown_data); + .and_modify(|d| *d = unknown_data); } } @@ -193,16 +192,10 @@ pub fn handle_pending_icons(&mut self) { } } - pub fn programs(&self) -> &HashMap { + pub fn programs(&self) -> &HashMap { &self.programs } - pub fn edit_fav(&mut self, program: &Program, add: bool) { - if let Some(info) = self.programs.get_mut(program) { - info.is_favorite = add; - } - } - pub fn picon_tooltip<'a>( &self, icon_key: &str, diff --git a/src/networking/types/service.rs b/src/networking/types/service.rs index 5dcced84..4eabf677 100644 --- a/src/networking/types/service.rs +++ b/src/networking/types/service.rs @@ -1,5 +1,8 @@ +use serde::de::{self, Deserializer, VariantAccess}; +use serde::{Deserialize, Serialize}; + /// Upper layer services. -#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Default)] +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Default, Serialize)] pub enum Service { /// One of the known services. Name(&'static str), @@ -10,6 +13,55 @@ pub enum Service { NotApplicable, } +impl<'de> Deserialize<'de> for Service { + fn deserialize(deserializer: D) -> Result + where + D: Deserializer<'de>, + { + struct ServiceVisitor; + + impl<'de> de::Visitor<'de> for ServiceVisitor { + type Value = Service; + + fn expecting(&self, formatter: &mut std::fmt::Formatter) -> std::fmt::Result { + formatter.write_str("a Service enum") + } + + fn visit_enum(self, data: A) -> Result + where + A: de::EnumAccess<'de>, + { + let (variant, access) = data.variant::()?; + match variant.as_str() { + "Name" => { + let s: String = access.newtype_variant()?; + let leaked: &'static str = Box::leak(s.into_boxed_str()); + Ok(Service::Name(leaked)) + } + "Unknown" => { + access.unit_variant()?; + Ok(Service::Unknown) + } + "NotApplicable" => { + access.unit_variant()?; + Ok(Service::NotApplicable) + } + other => Err(de::Error::unknown_variant( + other, + &["Name", "Unknown", "NotApplicable"], + )), + } + } + } + + deserializer.deserialize_enum( + "Service", + &["Name", "Unknown", "NotApplicable"], + ServiceVisitor, + ) + } +} + impl Service { pub fn to_string_with_equal_prefix(self) -> String { format!("={self}") @@ -53,4 +105,31 @@ fn test_service_to_string_with_equal_prefix() { assert_eq!(Service::NotApplicable.to_string_with_equal_prefix(), "=-"); assert_eq!(Service::Unknown.to_string_with_equal_prefix(), "=?"); } + + #[test] + fn test_deserialize_name() { + let json = serde_json::to_string(&Service::Name("https")).unwrap(); + let deserialized: Service = serde_json::from_str(&json).unwrap(); + assert_eq!(deserialized, Service::Name("https")); + } + + #[test] + fn test_deserialize_unknown() { + let json = serde_json::to_string(&Service::Unknown).unwrap(); + let deserialized: Service = serde_json::from_str(&json).unwrap(); + assert_eq!(deserialized, Service::Unknown); + } + + #[test] + fn test_deserialize_not_applicable() { + let json = serde_json::to_string(&Service::NotApplicable).unwrap(); + let deserialized: Service = serde_json::from_str(&json).unwrap(); + assert_eq!(deserialized, Service::NotApplicable); + } + + #[test] + fn test_deserialize_invalid_variant() { + let json = r#""InvalidVariant""#; + assert!(serde_json::from_str::(json).is_err()); + } } diff --git a/src/notifications/notify_and_log.rs b/src/notifications/notify_and_log.rs index d03776df..0e64e10e 100644 --- a/src/notifications/notify_and_log.rs +++ b/src/notifications/notify_and_log.rs @@ -2,7 +2,6 @@ use crate::networking::manage_packets::get_address_to_lookup; use crate::networking::types::capture_context::CaptureSource; use crate::networking::types::data_info::DataInfo; -use crate::networking::types::data_info_fav::DataInfoFav; use crate::networking::types::data_info_host::DataInfoHost; use crate::networking::types::data_representation::DataRepr; use crate::networking::types::host::Host; @@ -109,14 +108,13 @@ pub fn notify_and_log( .get(&host) .copied() .unwrap_or_default(); - data_info_host.data_info_fav.data_info = v.data_info(); + data_info_host.data_info = v.data_info(); blacklisted_last_interval .entry(*address_to_lookup) .and_modify(|(_, existing_data_info_host)| { existing_data_info_host - .data_info_fav .data_info - .refresh(data_info_host.data_info_fav.data_info); + .refresh(data_info_host.data_info); }) .or_insert((host, data_info_host)); } @@ -163,11 +161,8 @@ fn threshold_hosts( .map(|(h, data)| (h.clone(), *data)) .collect(); hosts.sort_by(|(_, a), (_, b)| { - a.data_info_fav.data_info.compare( - &b.data_info_fav.data_info, - SortType::Descending, - data_repr, - ) + a.data_info + .compare(&b.data_info, SortType::Descending, data_repr) }); hosts.truncate(4); hosts @@ -181,7 +176,7 @@ fn threshold_services( .services .iter() .filter(|(service, _)| service != &&Service::NotApplicable) - .map(|(s, data_info_fav)| (*s, data_info_fav.data_info)) + .map(|(s, data_info)| (*s, *data_info)) .collect(); services.sort_by(|(_, a), (_, b)| a.compare(b, SortType::Descending, data_repr)); services.truncate(4); @@ -211,13 +206,7 @@ fn favorites_last_interval( .filter(|v| v.program.eq(p)) .for_each(|v| data_info.refresh(v.data_info())); if data_info.tot_data(DataRepr::Packets) > 0 { - Some(FavoriteItem::Program(( - p.clone(), - DataInfoFav { - data_info, - is_favorite: true, - }, - ))) + Some(FavoriteItem::Program((p.clone(), data_info))) } else { None } diff --git a/src/notifications/types/logged_notification.rs b/src/notifications/types/logged_notification.rs index 9e8f2ea3..3fe011f0 100644 --- a/src/notifications/types/logged_notification.rs +++ b/src/notifications/types/logged_notification.rs @@ -82,9 +82,7 @@ pub fn data_info(&self) -> DataInfo { match self { LoggedNotification::DataThresholdExceeded(d) => d.data_info, LoggedNotification::FavoriteTransmitted(f) => f.favorite.data_info(), - LoggedNotification::BlacklistedTransmitted(b) => { - b.data_info_host.data_info_fav.data_info - } + LoggedNotification::BlacklistedTransmitted(b) => b.data_info_host.data_info, } } @@ -180,7 +178,7 @@ fn to_json(&self) -> String { "domain": self.host.domain, "asn": self.host.asn.name, }, - "data": DataRepr::Bytes.formatted_string(self.data_info_host.data_info_fav.data_info.tot_data(DataRepr::Bytes)), + "data": DataRepr::Bytes.formatted_string(self.data_info_host.data_info.tot_data(DataRepr::Bytes)), }) .to_string() } @@ -191,7 +189,6 @@ mod tests { use super::*; use crate::countries::types::country::Country; use crate::networking::types::asn::Asn; - use crate::networking::types::data_info_fav::DataInfoFav; use crate::networking::types::program::Program; use crate::networking::types::traffic_direction::TrafficDirection; use crate::networking::types::traffic_type::TrafficType; @@ -228,14 +225,7 @@ fn test_favorite_host_transmitted_to_json() { }, }; let data_info_host = DataInfoHost { - data_info_fav: DataInfoFav { - data_info: { - let mut di = DataInfo::default(); - di.add_packets(5, 500, TrafficDirection::Outgoing, Instant::now()); - di - }, - is_favorite: true, - }, + data_info: DataInfo::new_for_tests(0, 5, 0, 500), is_loopback: false, is_local: false, is_bogon: None, @@ -258,13 +248,7 @@ fn test_favorite_service_transmitted_to_json() { let service = Service::Name("https"); let mut data_info = DataInfo::default(); data_info.add_packets(12, 1500, TrafficDirection::Incoming, Instant::now()); - let favorite_item = FavoriteItem::Service(( - service, - DataInfoFav { - data_info, - is_favorite: true, - }, - )); + let favorite_item = FavoriteItem::Service((service, data_info)); let notification = FavoriteTransmitted { id: 3, favorite: favorite_item, @@ -284,13 +268,7 @@ fn test_favorite_program_transmitted_to_json() { )); let mut data_info = DataInfo::default(); data_info.add_packets(20, 2_500_000, TrafficDirection::Outgoing, Instant::now()); - let favorite_item = FavoriteItem::Program(( - program, - DataInfoFav { - data_info, - is_favorite: true, - }, - )); + let favorite_item = FavoriteItem::Program((program, data_info)); let notification = FavoriteTransmitted { id: 4, favorite: favorite_item, @@ -313,14 +291,7 @@ fn test_blacklisted_transmitted_to_json() { }, }; let data_info_host = DataInfoHost { - data_info_fav: DataInfoFav { - data_info: { - let mut di = DataInfo::default(); - di.add_packets(50, 10_000, TrafficDirection::Incoming, Instant::now()); - di - }, - is_favorite: false, - }, + data_info: DataInfo::new_for_tests(50, 0, 10_000, 0), is_loopback: false, is_local: false, is_bogon: None, diff --git a/src/report/get_report_entries.rs b/src/report/get_report_entries.rs index 81295b38..0d65b524 100644 --- a/src/report/get_report_entries.rs +++ b/src/report/get_report_entries.rs @@ -1,12 +1,10 @@ -use std::cmp::min; - use crate::Sniffer; +use crate::gui::types::favorite::FavoriteKey; use crate::networking::manage_packets::get_address_to_lookup; use crate::networking::types::address_port_pair::AddressPortPair; use crate::networking::types::data_info::DataInfo; -use crate::networking::types::data_info_fav::DataInfoFav; -use crate::networking::types::data_info_host::DataInfoHost; use crate::networking::types::info_address_port_pair::InfoAddressPortPair; +use std::cmp::min; /// Return the elements that satisfy the search constraints and belong to the given page, /// and the total number of elements which satisfy the search constraints, @@ -20,6 +18,7 @@ pub fn get_searched_entries( ) { let mut agglomerate = DataInfo::default(); let info_traffic = &sniffer.info_traffic; + let favorites = &sniffer.conf.settings.favorites; let mut all_results: Vec<(&AddressPortPair, &InfoAddressPortPair)> = info_traffic .map .iter() @@ -28,27 +27,15 @@ pub fn get_searched_entries( let r_dns_host = sniffer.addresses_resolved.get(address_to_lookup); // is this a favorite host? let is_favorite_host = if let Some(e) = r_dns_host { - info_traffic - .hosts - .get(&e.1) - .unwrap_or(&DataInfoHost::default()) - .data_info_fav - .is_favorite + favorites.contains(&FavoriteKey::Host(e.1.clone())) } else { false }; // is this a favorite service? - let is_favorite_service = info_traffic - .services - .get(&value.service) - .unwrap_or(&DataInfoFav::default()) - .is_favorite; + let is_favorite_service = favorites.contains(&FavoriteKey::Service(value.service)); // is this a favorite program? - let is_favorite_program = if let Some(pl) = sniffer.program_lookup.as_ref() { - pl.programs() - .get(&value.program) - .unwrap_or(&DataInfoFav::default()) - .is_favorite + let is_favorite_program = if sniffer.program_lookup.is_some() { + favorites.contains(&FavoriteKey::Program(value.program.clone())) } else { false };