Skip to content

Commit

Permalink
MMDB readers are now part of Sniffer struct
Browse files Browse the repository at this point in the history
  • Loading branch information
@GyulyVGC
GyulyVGC committed Sep 10, 2023
1 parent 78e2c54 commit 3b36289
Show file tree
Hide file tree
Showing 5 changed files with 43 additions and 31 deletions.
13 changes: 0 additions & 13 deletions src/countries/country_utils.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -30,19 +30,6 @@ use crate::{Language, StyleType};

pub const COUNTRY_MMDB: &[u8] = include_bytes!("../../resources/DB/GeoLite2-Country.mmdb");

pub fn mmdb_country_reader(mmdb_country_path: String) -> MmdbReader {
let default_reader = maxminddb::Reader::from_source(COUNTRY_MMDB).unwrap();
if mmdb_country_path.is_empty() {
MmdbReader::Default(default_reader)
} else {
let custom_reader_result = maxminddb::Reader::open_readfile(mmdb_country_path);
if let Ok(custom_reader) = custom_reader_result {
return MmdbReader::Custom(custom_reader);
}
MmdbReader::Default(default_reader)
}
}

pub fn get_country(address_to_lookup: &str, country_db_reader: &MmdbReader) -> Country {
let country_result: Result<geoip2::Country, MaxMindDBError> = match country_db_reader {
MmdbReader::Default(reader) => reader.lookup(address_to_lookup.parse().unwrap()),
Expand Down
14 changes: 13 additions & 1 deletion src/gui/pages/settings_advanced_page.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@ use crate::translations::translations_3::{
mmdb_paths_translation, params_not_editable_translation, restore_defaults_translation,
scale_factor_translation,
};
use crate::utils::asn::MmdbReader;
use crate::utils::formatted_strings::get_default_report_directory;
use crate::utils::types::icon::Icon;
use crate::{ConfigAdvancedSettings, Language, Sniffer, Status, StyleType};
Expand All @@ -24,6 +25,7 @@ use iced::widget::{
use iced::Length::Fixed;
use iced::{Alignment, Font, Length, Renderer};
use std::path::PathBuf;
use std::sync::Arc;

pub fn settings_advanced_page(sniffer: &Sniffer) -> Container<Message, Renderer<StyleType>> {
let font = get_font(sniffer.style);
Expand Down Expand Up @@ -79,6 +81,8 @@ pub fn settings_advanced_page(sniffer: &Sniffer) -> Container<Message, Renderer<
font,
&sniffer.advanced_settings.mmdb_country,
&sniffer.advanced_settings.mmdb_asn,
&sniffer.country_mmdb_reader,
&sniffer.asn_mmdb_reader,
));

Container::new(content)
Expand Down Expand Up @@ -194,6 +198,8 @@ fn mmdb_settings(
font: Font,
country_path: &str,
asn_path: &str,
country_reader: &Arc<MmdbReader>,
asn_reader: &Arc<MmdbReader>,
) -> Column<'static, Message, Renderer<StyleType>> {
Column::new()
.spacing(5)
Expand Down Expand Up @@ -233,13 +239,15 @@ fn mmdb_settings(
font,
Message::CustomCountryDb,
country_path,
country_reader,
country_translation(language),
))
.push(mmdb_input(
is_editable,
font,
Message::CustomAsnDb,
asn_path,
asn_reader,
"ASN",
)),
)
Expand All @@ -250,12 +258,16 @@ fn mmdb_input(
font: Font,
message: fn(String) -> Message,
custom_path: &str,
mmdb_reader: &Arc<MmdbReader>,
caption: &str,
) -> Row<'static, Message, Renderer<StyleType>> {
let is_error = if custom_path.is_empty() {
false
} else {
maxminddb::Reader::open_readfile(custom_path.clone()).is_err()
match **mmdb_reader {
MmdbReader::Default(_) => true,
MmdbReader::Custom(_) => false,
}
};

let mut input = TextInput::new("-", custom_path)
Expand Down
29 changes: 23 additions & 6 deletions src/gui/types/sniffer.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@ use pcap::Device;

use crate::chart::manage_chart_data::update_charts_data;
use crate::configs::types::config_window::ConfigWindow;
use crate::countries::country_utils::COUNTRY_MMDB;
use crate::gui::components::types::my_modal::MyModal;
use crate::gui::pages::types::running_page::RunningPage;
use crate::gui::pages::types::settings_page::SettingsPage;
Expand All @@ -30,6 +31,7 @@ use crate::report::get_report_entries::get_searched_entries;
use crate::report::types::report_sort_type::ReportSortType;
use crate::secondary_threads::parse_packets::parse_packets;
use crate::translations::types::language::Language;
use crate::utils::asn::{mmdb_reader, MmdbReader, ASN_MMDB};
use crate::utils::formatted_strings::{get_default_report_directory, push_pcap_file_name};
use crate::utils::types::web_page::WebPage;
use crate::{
Expand Down Expand Up @@ -93,6 +95,10 @@ pub struct Sniffer {
pub advanced_settings: ConfigAdvancedSettings,
/// Position and size of the app window
pub window: ConfigWindow,
/// MMDB reader for countries
pub country_mmdb_reader: Arc<MmdbReader>,
/// MMDB reader for ASN
pub asn_mmdb_reader: Arc<MmdbReader>,
}

impl Sniffer {
Expand Down Expand Up @@ -130,6 +136,11 @@ impl Sniffer {
last_focus_time: std::time::Instant::now(),
advanced_settings: configs.advanced_settings.clone(),
window: configs.window,
country_mmdb_reader: Arc::new(mmdb_reader(
&configs.advanced_settings.mmdb_country,
COUNTRY_MMDB,
)),
asn_mmdb_reader: Arc::new(mmdb_reader(&configs.advanced_settings.mmdb_asn, ASN_MMDB)),
}
}

Expand Down Expand Up @@ -262,8 +273,14 @@ impl Sniffer {
(f64::from(height) * self.advanced_settings.scale_factor) as u32;
self.window.size = (scaled_width, scaled_height);
}
Message::CustomCountryDb(db) => self.advanced_settings.mmdb_country = db,
Message::CustomAsnDb(db) => self.advanced_settings.mmdb_asn = db,
Message::CustomCountryDb(db) => {
self.advanced_settings.mmdb_country = db.clone();
self.country_mmdb_reader = Arc::new(mmdb_reader(&db, COUNTRY_MMDB));
}
Message::CustomAsnDb(db) => {
self.advanced_settings.mmdb_asn = db.clone();
self.asn_mmdb_reader = Arc::new(mmdb_reader(&db, ASN_MMDB));
}
Message::CustomReportDirectory(directory) => {
let path = push_pcap_file_name(PathBuf::from(directory));
self.advanced_settings.output_path = path;
Expand Down Expand Up @@ -371,8 +388,8 @@ 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();
let mmdb_asn_path = self.advanced_settings.mmdb_asn.clone();
let country_mmdb_reader = self.country_mmdb_reader.clone();
let asn_mmdb_reader = self.country_mmdb_reader.clone();
self.status_pair.1.notify_all();
thread::Builder::new()
.name("thread_parse_packets".to_string())
Expand All @@ -383,8 +400,8 @@ impl Sniffer {
cap.unwrap(),
filters,
&info_traffic_mutex,
mmdb_country_path,
mmdb_asn_path,
&country_mmdb_reader,
&asn_mmdb_reader,
);
})
.unwrap();
Expand Down
14 changes: 5 additions & 9 deletions src/secondary_threads/parse_packets.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -7,15 +7,14 @@ use std::thread;
use etherparse::PacketHeaders;
use pcap::{Active, Capture};

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,
};
use crate::networking::types::data_info::DataInfo;
use crate::networking::types::filters::Filters;
use crate::networking::types::info_address_port_pair::InfoAddressPortPair;
use crate::networking::types::my_device::MyDevice;
use crate::utils::asn::mmdb_asn_reader;
use crate::utils::asn::MmdbReader;
use crate::InfoTraffic;

/// The calling thread enters in a loop in which it waits for network packets, parses them according
Expand All @@ -26,14 +25,11 @@ pub fn parse_packets(
mut cap: Capture<Active>,
filters: Filters,
info_traffic_mutex: &Arc<Mutex<InfoTraffic>>,
mmdb_country_path: String,
mmdb_asn_path: String,
country_mmdb_reader: &Arc<MmdbReader>,
asn_mmdb_reader: &Arc<MmdbReader>,
) {
let capture_id = *current_capture_id.lock().unwrap();

let country_db_reader = Arc::new(mmdb_country_reader(mmdb_country_path));
let asn_db_reader = Arc::new(mmdb_asn_reader(mmdb_asn_path));

loop {
match cap.next_packet() {
Err(_) => {
Expand Down Expand Up @@ -125,8 +121,8 @@ pub fn parse_packets(
let key2 = key.clone();
let info_traffic2 = info_traffic_mutex.clone();
let device2 = device.clone();
let country_db_reader_2 = country_db_reader.clone();
let asn_db_reader_2 = asn_db_reader.clone();
let country_db_reader_2 = country_mmdb_reader.clone();
let asn_db_reader_2 = asn_mmdb_reader.clone();
thread::Builder::new()
.name("thread_reverse_dns_lookup".to_string())
.spawn(move || {
Expand Down
4 changes: 2 additions & 2 deletions src/utils/asn.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -9,8 +9,8 @@ pub enum MmdbReader {
Custom(Reader<Vec<u8>>),
}

pub fn mmdb_asn_reader(mmdb_asn_path: String) -> MmdbReader {
let default_reader = maxminddb::Reader::from_source(ASN_MMDB).unwrap();
pub fn mmdb_reader(mmdb_asn_path: &String, default_mmdb: &'static [u8]) -> MmdbReader {
let default_reader = maxminddb::Reader::from_source(default_mmdb).unwrap();
if mmdb_asn_path.is_empty() {
MmdbReader::Default(default_reader)
} else {
Expand Down

0 comments on commit 3b36289

Please sign in to comment.