|
1 | | -use maxminddb::{geoip2, MaxMindDBError}; |
2 | | - |
| 1 | +use crate::mmdb::types::mmdb_asn_entry::MmdbAsnEntry; |
3 | 2 | use crate::mmdb::types::mmdb_reader::MmdbReader; |
4 | 3 | use crate::networking::types::asn::Asn; |
5 | 4 |
|
6 | 5 | pub const ASN_MMDB: &[u8] = include_bytes!("../../resources/DB/GeoLite2-ASN.mmdb"); |
7 | 6 |
|
8 | 7 | #[allow(clippy::module_name_repetitions)] |
9 | | -pub fn get_asn(address_to_lookup: &str, asn_db_reader: &MmdbReader) -> Asn { |
10 | | - let asn_result: Result<geoip2::Asn, MaxMindDBError> = match asn_db_reader { |
11 | | - MmdbReader::Default(reader) => reader.lookup(address_to_lookup.parse().unwrap()), |
12 | | - MmdbReader::Custom(reader) => reader.lookup(address_to_lookup.parse().unwrap()), |
13 | | - }; |
14 | | - if let Ok(res) = asn_result { |
15 | | - if res.autonomous_system_number.is_some() && res.autonomous_system_organization.is_some() { |
16 | | - return Asn { |
17 | | - number: res.autonomous_system_number.unwrap(), |
18 | | - name: res.autonomous_system_organization.unwrap().to_string(), |
19 | | - }; |
20 | | - } |
| 8 | +pub fn get_asn(address: &str, asn_db_reader: &MmdbReader) -> Asn { |
| 9 | + if let Ok(res) = asn_db_reader.lookup::<MmdbAsnEntry>(address.parse().unwrap()) { |
| 10 | + return res.get_asn(); |
21 | 11 | } |
22 | 12 | Asn::default() |
23 | 13 | } |
| 14 | + |
| 15 | +#[cfg(test)] |
| 16 | +mod tests { |
| 17 | + use crate::mmdb::asn::{get_asn, ASN_MMDB}; |
| 18 | + use crate::mmdb::types::mmdb_reader::MmdbReader; |
| 19 | + |
| 20 | + #[test] |
| 21 | + fn test_get_asn_with_default_reader() { |
| 22 | + let reader_1 = MmdbReader::from(&String::from("unknown path"), ASN_MMDB); |
| 23 | + assert!(matches!(reader_1, MmdbReader::Default(_))); |
| 24 | + let reader_2 = MmdbReader::from(&String::new(), ASN_MMDB); |
| 25 | + assert!(matches!(reader_2, MmdbReader::Default(_))); |
| 26 | + let reader_3 = MmdbReader::from(&String::from("resources/repository/hr.png"), ASN_MMDB); |
| 27 | + assert!(matches!(reader_3, MmdbReader::Default(_))); |
| 28 | + let reader_4 = MmdbReader::from(&String::from("resources/DB/GeoLite2-ASN.mmdb"), ASN_MMDB); |
| 29 | + assert!(matches!(reader_4, MmdbReader::Custom(_))); |
| 30 | + let reader_5 = MmdbReader::from(&String::from("resources/DB/GeoLite2-ASN.mmdb"), &[]); |
| 31 | + assert!(matches!(reader_5, MmdbReader::Custom(_))); |
| 32 | + |
| 33 | + for reader in vec![reader_1, reader_2, reader_3, reader_4, reader_5] { |
| 34 | + // known IP |
| 35 | + let res = get_asn("8.8.8.8", &reader); |
| 36 | + assert_eq!(res.code, "15169"); |
| 37 | + assert_eq!(res.name, "GOOGLE"); |
| 38 | + |
| 39 | + // another known IP |
| 40 | + let res = get_asn("78.35.248.93", &reader); |
| 41 | + assert_eq!(res.code, "8422"); |
| 42 | + assert_eq!( |
| 43 | + res.name, |
| 44 | + "NetCologne Gesellschaft fur Telekommunikation mbH" |
| 45 | + ); |
| 46 | + |
| 47 | + // known IPv6 |
| 48 | + let res = get_asn("2806:230:2057::", &reader); |
| 49 | + assert_eq!(res.code, "11888"); |
| 50 | + assert_eq!(res.name, "Television Internacional, S.A. de C.V."); |
| 51 | + |
| 52 | + // unknown IP |
| 53 | + let res = get_asn("127.0.0.1", &reader); |
| 54 | + assert_eq!(res.code, ""); |
| 55 | + assert_eq!(res.name, ""); |
| 56 | + |
| 57 | + // unknown IPv6 |
| 58 | + let res = get_asn("::1", &reader); |
| 59 | + assert_eq!(res.code, ""); |
| 60 | + assert_eq!(res.name, ""); |
| 61 | + } |
| 62 | + } |
| 63 | + |
| 64 | + #[test] |
| 65 | + fn test_get_asn_with_custom_ipinfo_single_reader() { |
| 66 | + let reader_1 = MmdbReader::from( |
| 67 | + &String::from("resources/test/ipinfo_asn_sample.mmdb"), |
| 68 | + ASN_MMDB, |
| 69 | + ); |
| 70 | + let reader_2 = |
| 71 | + MmdbReader::from(&String::from("resources/test/ipinfo_asn_sample.mmdb"), &[]); |
| 72 | + |
| 73 | + for reader in vec![reader_1, reader_2] { |
| 74 | + assert!(matches!(reader, MmdbReader::Custom(_))); |
| 75 | + |
| 76 | + // known IP |
| 77 | + let res = get_asn("61.8.0.0", &reader); |
| 78 | + assert_eq!(res.code, "AS1221"); |
| 79 | + assert_eq!(res.name, "Telstra Limited"); |
| 80 | + |
| 81 | + // another known IP |
| 82 | + let res = get_asn("206.180.34.99", &reader); |
| 83 | + assert_eq!(res.code, "AS63344"); |
| 84 | + assert_eq!(res.name, "The Reynolds and Reynolds Company"); |
| 85 | + |
| 86 | + // known IPv6 |
| 87 | + let res = get_asn("2806:230:2057::", &reader); |
| 88 | + assert_eq!(res.code, "AS11888"); |
| 89 | + assert_eq!(res.name, "Television Internacional, S.A. de C.V."); |
| 90 | + |
| 91 | + // unknown IP |
| 92 | + let res = get_asn("127.0.0.1", &reader); |
| 93 | + assert_eq!(res.code, ""); |
| 94 | + assert_eq!(res.name, ""); |
| 95 | + |
| 96 | + // unknown IPv6 |
| 97 | + let res = get_asn("::1", &reader); |
| 98 | + assert_eq!(res.code, ""); |
| 99 | + assert_eq!(res.name, ""); |
| 100 | + } |
| 101 | + } |
| 102 | + |
| 103 | + #[test] |
| 104 | + fn test_get_asn_with_custom_ipinfo_combined_reader() { |
| 105 | + let reader_1 = MmdbReader::from( |
| 106 | + &String::from("resources/test/ipinfo_country_asn_sample.mmdb"), |
| 107 | + ASN_MMDB, |
| 108 | + ); |
| 109 | + let reader_2 = MmdbReader::from( |
| 110 | + &String::from("resources/test/ipinfo_country_asn_sample.mmdb"), |
| 111 | + &[], |
| 112 | + ); |
| 113 | + |
| 114 | + for reader in vec![reader_1, reader_2] { |
| 115 | + assert!(matches!(reader, MmdbReader::Custom(_))); |
| 116 | + |
| 117 | + // known IP |
| 118 | + let res = get_asn("31.171.144.141", &reader); |
| 119 | + assert_eq!(res.code, "AS197742"); |
| 120 | + assert_eq!(res.name, "IBB Energie AG"); |
| 121 | + |
| 122 | + // another known IP |
| 123 | + let res = get_asn("103.112.220.111", &reader); |
| 124 | + assert_eq!(res.code, "AS134077"); |
| 125 | + assert_eq!(res.name, "Magik Pivot Company Limited"); |
| 126 | + |
| 127 | + // known IPv6 |
| 128 | + let res = get_asn("2a02:6ea0:f001::", &reader); |
| 129 | + assert_eq!(res.code, "AS60068"); |
| 130 | + assert_eq!(res.name, "Datacamp Limited"); |
| 131 | + |
| 132 | + // unknown IP |
| 133 | + let res = get_asn("127.0.0.1", &reader); |
| 134 | + assert_eq!(res.code, ""); |
| 135 | + assert_eq!(res.name, ""); |
| 136 | + |
| 137 | + // unknown IPv6 |
| 138 | + let res = get_asn("::1", &reader); |
| 139 | + assert_eq!(res.code, ""); |
| 140 | + assert_eq!(res.name, ""); |
| 141 | + } |
| 142 | + } |
| 143 | +} |
0 commit comments