feat: custom country MMDB

src/configs/types/config_advanced_settings.rs

@@ -3,9 +3,11 @@
 
 use serde::{Deserialize, Serialize};
 
-#[derive(Serialize, Deserialize, Copy, Clone, PartialEq)]
+#[derive(Serialize, Deserialize, Clone, PartialEq)]
 pub struct ConfigAdvancedSettings {
     pub scale_factor: f64,
+    pub mmdb_country: String,
+    // pub mmdb_asn: String,
 }
 
 impl ConfigAdvancedSettings {
@@ -28,6 +30,9 @@ impl ConfigAdvancedSettings {
 
 impl Default for ConfigAdvancedSettings {
     fn default() -> Self {
-        ConfigAdvancedSettings { scale_factor: 1.0 }
+        ConfigAdvancedSettings {
+            scale_factor: 1.0,
+            mmdb_country: String::new(),
+        }
     }
 }

src/countries/country_utils.rs

@@ -29,9 +29,34 @@ use crate::{Language, StyleType};
 
 pub const COUNTRY_MMDB: &[u8] = include_bytes!("../../resources/DB/GeoLite2-Country.mmdb");
 
-pub fn get_country(address_to_lookup: &str, country_db_reader: &Reader<&[u8]>) -> Country {
-    let country_result: Result<geoip2::Country, MaxMindDBError> =
-        country_db_reader.lookup(address_to_lookup.parse().unwrap());
+pub fn mmdb_country_reader(
+    mmdb_country_path: String,
+) -> (Reader<&'static [u8]>, Option<Reader<Vec<u8>>>) {
+    let default_reader = maxminddb::Reader::from_source(COUNTRY_MMDB).unwrap();
+    if mmdb_country_path.is_empty() {
+        (default_reader, None)
+    } else {
+        let custom_reader_result = maxminddb::Reader::open_readfile(mmdb_country_path);
+        if let Ok(custom_reader) = custom_reader_result {
+            return (default_reader, Some(custom_reader));
+        }
+        (default_reader, None)
+    }
+}
+
+pub fn get_country(
+    address_to_lookup: &str,
+    country_db_readers: &(Reader<&[u8]>, Option<Reader<Vec<u8>>>),
+) -> Country {
+    let (default_reader, custom_reader) = country_db_readers;
+    let country_result: Result<geoip2::Country, MaxMindDBError> = if custom_reader.is_some() {
+        custom_reader
+            .as_ref()
+            .unwrap()
+            .lookup(address_to_lookup.parse().unwrap())
+    } else {
+        default_reader.lookup(address_to_lookup.parse().unwrap())
+    };
     if let Ok(res1) = country_result {
         if let Some(res2) = res1.country {
             if let Some(res3) = res2.iso_code {

src/gui/pages/settings_advanced_page.rs

@@ -4,6 +4,7 @@ use crate::gui::pages::types::settings_page::SettingsPage;
 use crate::gui::styles::container::ContainerType;
 use crate::gui::styles::style_constants::{get_font, get_font_headers, FONT_SIZE_SUBTITLE};
 use crate::gui::styles::text::TextType;
+use crate::gui::styles::text_input::TextInputType;
 use crate::gui::types::message::Message;
 use crate::translations::translations_3::{
     advanced_settings_translation, restore_defaults_translation, scale_factor_translation,
@@ -13,7 +14,7 @@ use crate::{ConfigAdvancedSettings, Language, Sniffer, StyleType};
 use iced::advanced::widget::Text;
 use iced::alignment::{Horizontal, Vertical};
 use iced::widget::tooltip::Position;
-use iced::widget::{button, vertical_space, Column, Container, Row, Slider, Tooltip};
+use iced::widget::{button, vertical_space, Column, Container, Row, Slider, TextInput, Tooltip};
 use iced::Length::Fixed;
 use iced::{Alignment, Font, Length, Renderer};
 
@@ -36,12 +37,20 @@ pub fn settings_advanced_page(sniffer: &Sniffer) -> Container<Message, Renderer<
             sniffer.language,
         ))
         .push(vertical_space(Fixed(15.0)))
-        .push(title_row(sniffer.language, font, sniffer.advanced_settings))
+        .push(title_row(
+            sniffer.language,
+            font,
+            &sniffer.advanced_settings,
+        ))
         .push(vertical_space(Fixed(5.0)))
         .push(scale_factor_slider(
             sniffer.language,
             font,
             sniffer.advanced_settings.scale_factor,
+        ))
+        .push(mmdb_country_input(
+            font,
+            &sniffer.advanced_settings.mmdb_country,
         ));
 
     Container::new(content)
@@ -53,7 +62,7 @@ pub fn settings_advanced_page(sniffer: &Sniffer) -> Container<Message, Renderer<
 fn title_row(
     language: Language,
     font: Font,
-    advanced_settings: ConfigAdvancedSettings,
+    advanced_settings: &ConfigAdvancedSettings,
 ) -> Row<'static, Message, Renderer<StyleType>> {
     let mut ret_val = Row::new().spacing(10).align_items(Alignment::Center).push(
         Text::new(advanced_settings_translation(language))
@@ -70,7 +79,7 @@ fn title_row(
                         .to_text()
                         .vertical_alignment(Vertical::Center)
                         .horizontal_alignment(Horizontal::Center)
-                        .size(15),
+                        .size(17),
                 )
                 .padding(2)
                 .height(Fixed(25.0))
@@ -115,3 +124,29 @@ fn scale_factor_slider(
     .align_x(Horizontal::Center)
     .align_y(Vertical::Center)
 }
+
+fn mmdb_country_input(
+    font: Font,
+    custom_path: &str,
+) -> Container<'static, Message, Renderer<StyleType>> {
+    let is_error = if custom_path.is_empty() {
+        false
+    } else {
+        maxminddb::Reader::open_readfile(custom_path.clone()).is_err()
+    };
+
+    let input = TextInput::new("-", custom_path)
+        .on_input(Message::CustomCountryDb)
+        .padding([0, 5])
+        .font(font)
+        .width(Length::Fixed(200.0))
+        .style(if is_error {
+            TextInputType::Error
+        } else {
+            TextInputType::Standard
+        });
+
+    Container::new(input)
+        .padding(5)
+        .style(ContainerType::Neutral)
+}

src/gui/styles/text_input.rs

@@ -14,6 +14,7 @@ pub enum TextInputType {
     #[default]
     Standard,
     Badge,
+    Error,
 }
 
 impl iced::widget::text_input::StyleSheet for StyleType {
@@ -24,28 +25,38 @@ impl iced::widget::text_input::StyleSheet for StyleType {
         Appearance {
             background: Background::Color(match style {
                 TextInputType::Badge => Color::TRANSPARENT,
-                TextInputType::Standard => colors.buttons,
+                _ => colors.buttons,
             }),
             border_radius: 0.0.into(),
-            border_width: 1.0,
+            border_width: match style {
+                TextInputType::Error => 2.0,
+                _ => 1.0,
+            },
             border_color: match style {
                 TextInputType::Badge => Color::TRANSPARENT,
                 TextInputType::Standard => Color {
                     a: get_alpha_round_borders(*self),
                     ..colors.buttons
                 },
+                TextInputType::Error => Color::new(0.8, 0.15, 0.15, 1.0),
             },
             icon_color: colors.text_body,
         }
     }
 
-    fn focused(&self, _: &Self::Style) -> iced::widget::text_input::Appearance {
+    fn focused(&self, style: &Self::Style) -> iced::widget::text_input::Appearance {
         let colors = get_colors(*self);
         Appearance {
             background: Background::Color(colors.primary),
             border_radius: 0.0.into(),
-            border_width: 1.0,
-            border_color: colors.secondary,
+            border_width: match style {
+                TextInputType::Error => 2.0,
+                _ => 1.0,
+            },
+            border_color: match style {
+                TextInputType::Error => Color::new(0.8, 0.15, 0.15, 1.0),
+                _ => colors.secondary,
+            },
             icon_color: colors.text_body,
         }
     }
@@ -79,11 +90,17 @@ impl iced::widget::text_input::StyleSheet for StyleType {
         Appearance {
             background: Background::Color(match style {
                 TextInputType::Badge => Color::TRANSPARENT,
-                TextInputType::Standard => colors.buttons,
+                _ => colors.buttons,
             }),
             border_radius: 0.0.into(),
-            border_width: 1.0,
-            border_color: colors.secondary,
+            border_width: match style {
+                TextInputType::Error => 2.0,
+                _ => 1.0,
+            },
+            border_color: match style {
+                TextInputType::Error => Color::new(0.8, 0.15, 0.15, 1.0),
+                _ => colors.secondary,
+            },
             icon_color: colors.text_body,
         }
     }

src/gui/types/message.rs

@@ -97,4 +97,6 @@ pub enum Message {
     WindowMoved(i32, i32),
     /// The app window size has been changed
     WindowResized(u32, u32),
+    /// The country MMDB custom path has been updated
+    CustomCountryDb(String),
 }

src/gui/types/sniffer.rs

@@ -130,7 +130,7 @@ impl Sniffer {
             page_number: 1,
             selected_connection: 0,
             last_focus_time: std::time::Instant::now(),
-            advanced_settings: configs.advanced_settings,
+            advanced_settings: configs.advanced_settings.clone(),
             window: configs.window,
         }
     }
@@ -250,6 +250,7 @@ impl Sniffer {
                 self.window.size = (scaled_width, scaled_height);
                 confy::store("sniffnet", "window", self.window).unwrap_or(());
             }
+            Message::CustomCountryDb(db) => self.advanced_settings.mmdb_country = db,
             _ => {}
         }
         Command::none()
@@ -357,6 +358,7 @@ impl Sniffer {
             // no pcap error
             let current_capture_id = self.current_capture_id.clone();
             let filters = self.filters;
+            let mmdb_country_path = self.advanced_settings.mmdb_country.clone();
             self.status_pair.1.notify_all();
             thread::Builder::new()
                 .name("thread_parse_packets".to_string())
@@ -367,6 +369,7 @@ impl Sniffer {
                         cap.unwrap(),
                         filters,
                         &info_traffic_mutex,
+                        mmdb_country_path,
                     );
                 })
                 .unwrap();
@@ -432,7 +435,12 @@ impl Sniffer {
                 color_gradient: self.color_gradient,
             };
             confy::store("sniffnet", "settings", settings).unwrap_or(());
-            confy::store("sniffnet", "advanced_settings", self.advanced_settings).unwrap_or(());
+            confy::store(
+                "sniffnet",
+                "advanced_settings",
+                self.advanced_settings.clone(),
+            )
+            .unwrap_or(());
         }
     }
 

src/networking/manage_packets.rs

@@ -247,7 +247,7 @@ pub fn reverse_dns_lookup(
     key: &AddressPortPair,
     traffic_direction: TrafficDirection,
     my_device: &MyDevice,
-    country_db_reader: &Reader<&[u8]>,
+    country_db_readers: &(Reader<&[u8]>, Option<Reader<Vec<u8>>>),
     asn_db_reader: &Reader<&[u8]>,
 ) {
     let address_to_lookup = get_address_to_lookup(key, traffic_direction);
@@ -263,7 +263,7 @@ pub fn reverse_dns_lookup(
         traffic_direction,
     );
     let is_local = is_local_connection(&address_to_lookup, &my_interface_addresses);
-    let country = get_country(&address_to_lookup, country_db_reader);
+    let country = get_country(&address_to_lookup, country_db_readers);
     let asn = asn(&address_to_lookup, asn_db_reader);
     let r_dns = if let Ok(result) = lookup_result {
         if result.is_empty() {

src/secondary_threads/parse_packets.rs

@@ -7,7 +7,7 @@ use std::thread;
 use etherparse::PacketHeaders;
 use pcap::{Active, Capture};
 
-use crate::countries::country_utils::COUNTRY_MMDB;
+use crate::countries::country_utils::mmdb_country_reader;
 use crate::networking::manage_packets::{
     analyze_headers, get_address_to_lookup, modify_or_insert_in_map, reverse_dns_lookup,
 };
@@ -26,10 +26,11 @@ pub fn parse_packets(
     mut cap: Capture<Active>,
     filters: Filters,
     info_traffic_mutex: &Arc<Mutex<InfoTraffic>>,
+    mmdb_country_path: String,
 ) {
     let capture_id = *current_capture_id.lock().unwrap();
 
-    let country_db_reader = Arc::new(maxminddb::Reader::from_source(COUNTRY_MMDB).unwrap());
+    let country_db_readers = Arc::new(mmdb_country_reader(mmdb_country_path));
     let asn_db_reader = Arc::new(maxminddb::Reader::from_source(ASN_MMDB).unwrap());
 
     loop {
@@ -123,7 +124,7 @@ pub fn parse_packets(
                                     let key2 = key.clone();
                                     let info_traffic2 = info_traffic_mutex.clone();
                                     let device2 = device.clone();
-                                    let country_db_reader2 = country_db_reader.clone();
+                                    let country_db_readers_2 = country_db_readers.clone();
                                     let asn_db_reader2 = asn_db_reader.clone();
                                     thread::Builder::new()
                                         .name("thread_reverse_dns_lookup".to_string())
@@ -133,7 +134,7 @@ pub fn parse_packets(
                                                 &key2,
                                                 new_info.traffic_direction,
                                                 &device2,
-                                                &country_db_reader2,
+                                                &country_db_readers_2,
                                                 &asn_db_reader2,
                                             );
                                         })

src/utils/types/icon.rs

@@ -67,7 +67,7 @@ impl Icon {
             Icon::Notification => "7",
             Icon::Overview => "d",
             Icon::PacketsThreshold => "e",
-            Icon::Restore => "E",
+            Icon::Restore => "k",
             Icon::Rocket => "S",
             Icon::Settings => "a",
             Icon::Sniffnet => "A",