use crate::countries::types::country::Country;
use crate::mmdb::types::mmdb_country_entry::MmdbCountryEntry;
use crate::mmdb::types::mmdb_reader::MmdbReader;
use std::net::IpAddr;
pub const COUNTRY_MMDB: &[u8] = include_bytes!("../../resources/DB/GeoLite2-Country.mmdb");
#[allow(clippy::module_name_repetitions)]
pub fn get_country(address: &IpAddr, country_db_reader: &MmdbReader) -> Country {
if let Some(res) = country_db_reader.lookup::<MmdbCountryEntry>(*address) {
return res.get_country();
}
Country::ZZ }
#[cfg(test)]
mod tests {
use crate::countries::types::country::Country;
use crate::mmdb::country::{COUNTRY_MMDB, get_country};
use crate::mmdb::types::mmdb_reader::MmdbReader;
use std::net::IpAddr;
use std::str::FromStr;
#[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] {
let res = get_country(&IpAddr::from([8, 8, 8, 8]), &reader);
assert_eq!(res, Country::US);
let res = get_country(&IpAddr::from([78, 35, 248, 93]), &reader);
assert_eq!(res, Country::DE);
let res = get_country(&IpAddr::from_str("2806:230:2057::").unwrap(), &reader);
assert_eq!(res, Country::MX);
let res = get_country(&IpAddr::from([127, 0, 0, 1]), &reader);
assert_eq!(res, Country::ZZ);
let res = get_country(&IpAddr::from_str("::1").unwrap(), &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_location_sample.mmdb"),
COUNTRY_MMDB,
);
let reader_2 = MmdbReader::from(
&String::from("resources/test/ipinfo_location_sample.mmdb"),
&[],
);
for reader in vec![reader_1, reader_2] {
assert!(matches!(reader, MmdbReader::Custom(_)));
let res = get_country(&IpAddr::from([1, 0, 6, 99]), &reader);
assert_eq!(res, Country::AU);
let res = get_country(&IpAddr::from([1, 0, 8, 0]), &reader);
assert_eq!(res, Country::CN);
let res = get_country(&IpAddr::from([127, 0, 0, 1]), &reader);
assert_eq!(res, Country::ZZ);
let res = get_country(&IpAddr::from_str("::1").unwrap(), &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_lite_sample.mmdb"),
COUNTRY_MMDB,
);
let reader_2 =
MmdbReader::from(&String::from("resources/test/ipinfo_lite_sample.mmdb"), &[]);
for reader in vec![reader_1, reader_2] {
assert!(matches!(reader, MmdbReader::Custom(_)));
let res = get_country(&IpAddr::from([1, 0, 65, 1]), &reader);
assert_eq!(res, Country::JP);
let res = get_country(&IpAddr::from([1, 6, 230, 0]), &reader);
assert_eq!(res, Country::IN);
let res = get_country(&IpAddr::from([127, 0, 0, 1]), &reader);
assert_eq!(res, Country::ZZ);
let res = get_country(&IpAddr::from_str("::1").unwrap(), &reader);
assert_eq!(res, Country::ZZ);
}
}
}