Add support for IPinfo MMDB format

CHANGELOG.md

@@ -4,9 +4,9 @@ All Sniffnet releases with the relative changes are documented in this file.
 
 ## [UNRELEASED]
 - Thumbnail mode improvements ([#512](https://github.com/GyulyVGC/sniffnet/pull/512))
+- Support IPinfo ASN and Country databases ([#558](https://github.com/GyulyVGC/sniffnet/pull/558) — fixes [#533](https://github.com/GyulyVGC/sniffnet/issues/533))
 - Added keyboard shortcuts to change zoom level (fixes [#554](https://github.com/GyulyVGC/sniffnet/issues/554))
 - Increased the range of selectable zoom values (fixes [#542](https://github.com/GyulyVGC/sniffnet/issues/542))
-- Reduced `String` allocations in translation code ([#524](https://github.com/GyulyVGC/sniffnet/pull/524))
 - Updated some of the existing translations to v1.3: 
   - French - [#494](https://github.com/GyulyVGC/sniffnet/pull/494)
   - German - [#495](https://github.com/GyulyVGC/sniffnet/pull/495)
@@ -16,7 +16,8 @@ All Sniffnet releases with the relative changes are documented in this file.
   - Japanese - [#504](https://github.com/GyulyVGC/sniffnet/pull/504)
   - Uzbek - [#510](https://github.com/GyulyVGC/sniffnet/pull/510)
   - Swedish - [#522](https://github.com/GyulyVGC/sniffnet/pull/522)
-- Fixed bug causing impossibility to exit thumbnail mode on Ubuntu (fixes [#505](https://github.com/GyulyVGC/sniffnet/pull/505))
+- Reduced `String` allocations in translation code ([#524](https://github.com/GyulyVGC/sniffnet/pull/524))
+- Fixed impossibility to exit thumbnail mode in some Linux distributions (fixes [#505](https://github.com/GyulyVGC/sniffnet/pull/505))
 
 ## [1.3.0] - 2024-04-08
 - Introduced thumbnail mode, enabling users to keep an eye on Sniffnet while doing other tasks ([#484](https://github.com/GyulyVGC/sniffnet/pull/484))

src/gui/pages/connection_details_page.rs

@@ -270,7 +270,8 @@ fn get_host_info_col(
     language: Language,
 ) -> Column<'static, Message, StyleType> {
     let mut host_info_col = Column::new().spacing(4);
-    if r_dns.parse::<IpAddr>().is_err() || (!host.asn.name.is_empty() && host.asn.number > 0) {
+    if r_dns.parse::<IpAddr>().is_err() || (!host.asn.name.is_empty() && !host.asn.code.is_empty())
+    {
         host_info_col = host_info_col.push(Rule::horizontal(10.0));
     }
     if r_dns.parse::<IpAddr>().is_err() {
@@ -280,10 +281,10 @@ fn get_host_info_col(
             font,
         ));
     }
-    if !host.asn.name.is_empty() && host.asn.number > 0 {
+    if !host.asn.name.is_empty() && !host.asn.code.is_empty() {
         host_info_col = host_info_col.push(TextType::highlighted_subtitle_with_desc(
             administrative_entity_translation(language),
-            &format!("{} (ASN {})", host.asn.name, host.asn.number),
+            &format!("{} (ASN {})", host.asn.name, host.asn.code),
             font,
         ));
     }

src/gui/pages/thumbnail_page.rs

@@ -166,7 +166,7 @@ mod tests {
             domain: domain.to_string(),
             asn: Asn {
                 name: asn.to_string(),
-                number: 512,
+                code: "512".to_string(),
             },
             country: Default::default(),
         }

src/mmdb/asn.rs

@@ -1,23 +1,143 @@
-use maxminddb::{geoip2, MaxMindDBError};
-
+use crate::mmdb::types::mmdb_asn_entry::MmdbAsnEntry;
 use crate::mmdb::types::mmdb_reader::MmdbReader;
 use crate::networking::types::asn::Asn;
 
 pub const ASN_MMDB: &[u8] = include_bytes!("../../resources/DB/GeoLite2-ASN.mmdb");
 
 #[allow(clippy::module_name_repetitions)]
-pub fn get_asn(address_to_lookup: &str, asn_db_reader: &MmdbReader) -> Asn {
-    let asn_result: Result<geoip2::Asn, MaxMindDBError> = match asn_db_reader {
-        MmdbReader::Default(reader) => reader.lookup(address_to_lookup.parse().unwrap()),
-        MmdbReader::Custom(reader) => reader.lookup(address_to_lookup.parse().unwrap()),
-    };
-    if let Ok(res) = asn_result {
-        if res.autonomous_system_number.is_some() && res.autonomous_system_organization.is_some() {
-            return Asn {
-                number: res.autonomous_system_number.unwrap(),
-                name: res.autonomous_system_organization.unwrap().to_string(),
-            };
-        }
+pub fn get_asn(address: &str, asn_db_reader: &MmdbReader) -> Asn {
+    if let Ok(res) = asn_db_reader.lookup::<MmdbAsnEntry>(address.parse().unwrap()) {
+        return res.get_asn();
     }
     Asn::default()
 }
+
+#[cfg(test)]
+mod tests {
+    use crate::mmdb::asn::{get_asn, ASN_MMDB};
+    use crate::mmdb::types::mmdb_reader::MmdbReader;
+
+    #[test]
+    fn test_get_asn_with_default_reader() {
+        let reader_1 = MmdbReader::from(&String::from("unknown path"), ASN_MMDB);
+        assert!(matches!(reader_1, MmdbReader::Default(_)));
+        let reader_2 = MmdbReader::from(&String::new(), ASN_MMDB);
+        assert!(matches!(reader_2, MmdbReader::Default(_)));
+        let reader_3 = MmdbReader::from(&String::from("resources/repository/hr.png"), ASN_MMDB);
+        assert!(matches!(reader_3, MmdbReader::Default(_)));
+        let reader_4 = MmdbReader::from(&String::from("resources/DB/GeoLite2-ASN.mmdb"), ASN_MMDB);
+        assert!(matches!(reader_4, MmdbReader::Custom(_)));
+        let reader_5 = MmdbReader::from(&String::from("resources/DB/GeoLite2-ASN.mmdb"), &[]);
+        assert!(matches!(reader_5, MmdbReader::Custom(_)));
+
+        for reader in vec![reader_1, reader_2, reader_3, reader_4, reader_5] {
+            // known IP
+            let res = get_asn("8.8.8.8", &reader);
+            assert_eq!(res.code, "15169");
+            assert_eq!(res.name, "GOOGLE");
+
+            // another known IP
+            let res = get_asn("78.35.248.93", &reader);
+            assert_eq!(res.code, "8422");
+            assert_eq!(
+                res.name,
+                "NetCologne Gesellschaft fur Telekommunikation mbH"
+            );
+
+            // known IPv6
+            let res = get_asn("2806:230:2057::", &reader);
+            assert_eq!(res.code, "11888");
+            assert_eq!(res.name, "Television Internacional, S.A. de C.V.");
+
+            // unknown IP
+            let res = get_asn("127.0.0.1", &reader);
+            assert_eq!(res.code, "");
+            assert_eq!(res.name, "");
+
+            // unknown IPv6
+            let res = get_asn("::1", &reader);
+            assert_eq!(res.code, "");
+            assert_eq!(res.name, "");
+        }
+    }
+
+    #[test]
+    fn test_get_asn_with_custom_ipinfo_single_reader() {
+        let reader_1 = MmdbReader::from(
+            &String::from("resources/test/ipinfo_asn_sample.mmdb"),
+            ASN_MMDB,
+        );
+        let reader_2 =
+            MmdbReader::from(&String::from("resources/test/ipinfo_asn_sample.mmdb"), &[]);
+
+        for reader in vec![reader_1, reader_2] {
+            assert!(matches!(reader, MmdbReader::Custom(_)));
+
+            // known IP
+            let res = get_asn("61.8.0.0", &reader);
+            assert_eq!(res.code, "AS1221");
+            assert_eq!(res.name, "Telstra Limited");
+
+            // another known IP
+            let res = get_asn("206.180.34.99", &reader);
+            assert_eq!(res.code, "AS63344");
+            assert_eq!(res.name, "The Reynolds and Reynolds Company");
+
+            // known IPv6
+            let res = get_asn("2806:230:2057::", &reader);
+            assert_eq!(res.code, "AS11888");
+            assert_eq!(res.name, "Television Internacional, S.A. de C.V.");
+
+            // unknown IP
+            let res = get_asn("127.0.0.1", &reader);
+            assert_eq!(res.code, "");
+            assert_eq!(res.name, "");
+
+            // unknown IPv6
+            let res = get_asn("::1", &reader);
+            assert_eq!(res.code, "");
+            assert_eq!(res.name, "");
+        }
+    }
+
+    #[test]
+    fn test_get_asn_with_custom_ipinfo_combined_reader() {
+        let reader_1 = MmdbReader::from(
+            &String::from("resources/test/ipinfo_country_asn_sample.mmdb"),
+            ASN_MMDB,
+        );
+        let reader_2 = MmdbReader::from(
+            &String::from("resources/test/ipinfo_country_asn_sample.mmdb"),
+            &[],
+        );
+
+        for reader in vec![reader_1, reader_2] {
+            assert!(matches!(reader, MmdbReader::Custom(_)));
+
+            // known IP
+            let res = get_asn("31.171.144.141", &reader);
+            assert_eq!(res.code, "AS197742");
+            assert_eq!(res.name, "IBB Energie AG");
+
+            // another known IP
+            let res = get_asn("103.112.220.111", &reader);
+            assert_eq!(res.code, "AS134077");
+            assert_eq!(res.name, "Magik Pivot Company Limited");
+
+            // known IPv6
+            let res = get_asn("2a02:6ea0:f001::", &reader);
+            assert_eq!(res.code, "AS60068");
+            assert_eq!(res.name, "Datacamp Limited");
+
+            // unknown IP
+            let res = get_asn("127.0.0.1", &reader);
+            assert_eq!(res.code, "");
+            assert_eq!(res.name, "");
+
+            // unknown IPv6
+            let res = get_asn("::1", &reader);
+            assert_eq!(res.code, "");
+            assert_eq!(res.name, "");
+        }
+    }
+}

src/mmdb/country.rs

@@ -1,22 +1,131 @@
-use maxminddb::{geoip2, MaxMindDBError};
-
 use crate::countries::types::country::Country;
+use crate::mmdb::types::mmdb_country_entry::MmdbCountryEntry;
 use crate::mmdb::types::mmdb_reader::MmdbReader;
 
 pub const COUNTRY_MMDB: &[u8] = include_bytes!("../../resources/DB/GeoLite2-Country.mmdb");
 
 #[allow(clippy::module_name_repetitions)]
-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()),
-        MmdbReader::Custom(reader) => 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 {
-                return Country::from_str(res3);
-            }
-        }
+pub fn get_country(address: &str, country_db_reader: &MmdbReader) -> Country {
+    if let Ok(res) = country_db_reader.lookup::<MmdbCountryEntry>(address.parse().unwrap()) {
+        return res.get_country();
     }
     Country::ZZ // unknown
 }
+
+#[cfg(test)]
+mod tests {
+    use crate::countries::types::country::Country;
+    use crate::mmdb::country::{get_country, COUNTRY_MMDB};
+    use crate::mmdb::types::mmdb_reader::MmdbReader;
+
+    #[test]
+    fn test_get_country_with_default_reader() {
+        let reader_1 = MmdbReader::from(&String::from("unknown path"), COUNTRY_MMDB);
+        assert!(matches!(reader_1, MmdbReader::Default(_)));
+        let reader_2 = MmdbReader::from(&String::new(), COUNTRY_MMDB);
+        assert!(matches!(reader_2, MmdbReader::Default(_)));
+        let reader_3 = MmdbReader::from(&String::from("resources/repository/hr.png"), COUNTRY_MMDB);
+        assert!(matches!(reader_3, MmdbReader::Default(_)));
+        let reader_4 = MmdbReader::from(
+            &String::from("resources/DB/GeoLite2-Country.mmdb"),
+            COUNTRY_MMDB,
+        );
+        assert!(matches!(reader_4, MmdbReader::Custom(_)));
+        let reader_5 = MmdbReader::from(&String::from("resources/DB/GeoLite2-Country.mmdb"), &[]);
+        assert!(matches!(reader_5, MmdbReader::Custom(_)));
+
+        for reader in vec![reader_1, reader_2, reader_3, reader_4, reader_5] {
+            // known IP
+            let res = get_country("8.8.8.8", &reader);
+            assert_eq!(res, Country::US);
+
+            // another known IP
+            let res = get_country("78.35.248.93", &reader);
+            assert_eq!(res, Country::DE);
+
+            // known IPv6
+            let res = get_country("2806:230:2057::", &reader);
+            assert_eq!(res, Country::MX);
+
+            // unknown IP
+            let res = get_country("127.0.0.1", &reader);
+            assert_eq!(res, Country::ZZ);
+
+            // unknown IPv6
+            let res = get_country("::1", &reader);
+            assert_eq!(res, Country::ZZ);
+        }
+    }
+
+    #[test]
+    fn test_get_country_with_custom_ipinfo_single_reader() {
+        let reader_1 = MmdbReader::from(
+            &String::from("resources/test/ipinfo_country_sample.mmdb"),
+            COUNTRY_MMDB,
+        );
+        let reader_2 = MmdbReader::from(
+            &String::from("resources/test/ipinfo_country_sample.mmdb"),
+            &[],
+        );
+
+        for reader in vec![reader_1, reader_2] {
+            assert!(matches!(reader, MmdbReader::Custom(_)));
+
+            // known IP
+            let res = get_country("2.2.146.0", &reader);
+            assert_eq!(res, Country::GB);
+
+            // another known IP
+            let res = get_country("23.193.112.81", &reader);
+            assert_eq!(res, Country::US);
+
+            // known IPv6
+            let res = get_country("2a0e:1d80::", &reader);
+            assert_eq!(res, Country::RO);
+
+            // unknown IP
+            let res = get_country("127.0.0.1", &reader);
+            assert_eq!(res, Country::ZZ);
+
+            // unknown IPv6
+            let res = get_country("::1", &reader);
+            assert_eq!(res, Country::ZZ);
+        }
+    }
+
+    #[test]
+    fn test_get_country_with_custom_ipinfo_combined_reader() {
+        let reader_1 = MmdbReader::from(
+            &String::from("resources/test/ipinfo_country_asn_sample.mmdb"),
+            COUNTRY_MMDB,
+        );
+        let reader_2 = MmdbReader::from(
+            &String::from("resources/test/ipinfo_country_asn_sample.mmdb"),
+            &[],
+        );
+
+        for reader in vec![reader_1, reader_2] {
+            assert!(matches!(reader, MmdbReader::Custom(_)));
+
+            // known IP
+            let res = get_country("31.171.144.141", &reader);
+            assert_eq!(res, Country::IT);
+
+            // another known IP
+            let res = get_country("103.112.220.111", &reader);
+            assert_eq!(res, Country::TH);
+
+            // known IPv6
+            let res = get_country("2a02:6ea0:f001::", &reader);
+            assert_eq!(res, Country::AR);
+
+            // unknown IP
+            let res = get_country("127.0.0.1", &reader);
+            assert_eq!(res, Country::ZZ);
+
+            // unknown IPv6
+            let res = get_country("::1", &reader);
+            assert_eq!(res, Country::ZZ);
+        }
+    }
+}

src/mmdb/types/mmdb_asn_entry.rs

@@ -0,0 +1,36 @@
+use crate::networking::types::asn::Asn;
+use serde::Deserialize;
+
+#[derive(Deserialize)]
+pub struct MmdbAsnEntry<'a> {
+    #[serde(alias = "autonomous_system_number", alias = "asn")]
+    code: MmdbAsnCode<'a>,
+    #[serde(alias = "autonomous_system_organization", alias = "as_name")]
+    name: Option<&'a str>,
+}
+
+impl MmdbAsnEntry<'_> {
+    pub fn get_asn(&self) -> Asn {
+        Asn {
+            code: self.code.get_code(),
+            name: self.name.unwrap_or_default().to_string(),
+        }
+    }
+}
+
+#[derive(Deserialize)]
+#[serde(untagged)]
+enum MmdbAsnCode<'a> {
+    Int(Option<u32>),
+    Str(Option<&'a str>),
+}
+
+impl MmdbAsnCode<'_> {
+    fn get_code(&self) -> String {
+        match self {
+            Self::Int(Some(code)) => code.to_string(),
+            Self::Str(Some(code)) => code.to_string(),
+            _ => String::new(),
+        }
+    }
+}