#[macro_use]
extern crate serde_derive;
use std::collections::HashMap;
use std::io::Read;
use std::net::Ipv4Addr;
use ipnet::Ipv4Net;
use crate::datasets::{Block, Location};
mod datasets;
#[derive(Debug)]
pub struct GeoIPDB {
locations: HashMap<u32, Location>,
blocks: HashMap<u32, Vec<Block>>,
}
impl GeoIPDB {
fn expand_network(network: &Ipv4Net) -> Vec<u32> {
let prefix = network.prefix_len();
let expanded_networks;
if prefix < 16 {
expanded_networks = network
.subnets(16)
.unwrap()
.map(|network| GeoIPDB::ipnet_to_map_key(&network))
.collect();
} else {
expanded_networks = vec![GeoIPDB::ipnet_to_map_key(network)];
}
expanded_networks
}
fn ipnet_to_map_key(ip_address: &Ipv4Net) -> u32 {
GeoIPDB::ipaddr_to_map_key(&ip_address.addr())
}
fn ipaddr_to_map_key(ip_address: &Ipv4Addr) -> u32 {
ip_address.octets()[0..2]
.iter()
.map(|n| u32::from(*n))
.scan(1_000, |state, value| {
let res = *state * value;
*state = *state / 1000;
Some(res)
})
.map(|n| u32::from(n))
.sum()
}
pub fn new<R: Read + Sized>(blocks_csv_file: R, locations_csv_file: R) -> Self {
let mut blocks = HashMap::new();
datasets::parse_blocks_csv(blocks_csv_file)
.map(|block| {
let networks = GeoIPDB::expand_network(&block.network);
(block, networks)
})
.for_each(|(block, networks)| {
networks.iter().for_each(|network| {
let blocks = blocks.entry(*network).or_insert(Vec::new());
blocks.push(block.clone());
});
});
let mut locations = HashMap::new();
datasets::parse_locations_csv(locations_csv_file).for_each(|location| {
locations.insert(location.geoname_id, location);
});
GeoIPDB { locations, blocks }
}
pub fn resolve(&self, ip_address: &str) -> Option<&Block> {
let ip_address = ip_address.parse::<Ipv4Addr>().unwrap();
let candidates = self.blocks.get(&GeoIPDB::ipaddr_to_map_key(&ip_address));
candidates.and_then(|candidates| {
candidates
.iter()
.find(|block| block.network.contains(&ip_address))
})
}
pub fn get_location(&self, geoname_id: u32) -> &Location {
self.locations.get(&geoname_id).unwrap()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
#[ignore]
fn ipaddress_expansion() {
let ip1 = "172.16.0.0/26".parse::<Ipv4Net>().unwrap();
ip1.subnets(16).unwrap().for_each(|ip| println!("{:?}", ip));
let ip1 = "172.16.0.128/26".parse::<Ipv4Net>().unwrap();
ip1.subnets(16).unwrap().for_each(|ip| println!("{:?}", ip));
}
#[test]
fn ip_to_number() {
assert_eq!(
255255,
GeoIPDB::ipnet_to_map_key(&"255.255.255.0/24".parse::<Ipv4Net>().unwrap())
);
assert_eq!(
255255,
GeoIPDB::ipaddr_to_map_key(&"255.255.255.12".parse::<Ipv4Addr>().unwrap())
);
assert_eq!(
1000,
GeoIPDB::ipaddr_to_map_key(&"1.0.0.1".parse::<Ipv4Addr>().unwrap())
);
assert_eq!(
81030,
GeoIPDB::ipaddr_to_map_key(&"81.30.9.30".parse::<Ipv4Addr>().unwrap())
);
}
#[test]
fn can_resolve_ip() {
let blocks = "network,geoname_id,registered_country_geoname_id,represented_country_geoname_id,is_anonymous_proxy,is_satellite_provider,postal_code,latitude,longitude,accuracy_radius
1.0.0.0/24,2077456,2077456,,0,0,,-33.4940,143.2104,1000
1.0.1.0/24,1811017,1814991,,0,0,,24.4798,118.0819,50
1.3.0.0/16,1809935,1814991,,0,0,,23.1167,113.2500,50";
let locations = "geoname_id,locale_code,continent_code,continent_name,country_iso_code,country_name,subdivision_1_iso_code,subdivision_1_name,subdivision_2_iso_code,subdivision_2_name,city_name,metro_code,time_zone,is_in_european_union
1809935,en,AS,Asia,CN,China,GD,Guangdong,,,,,Asia/Shanghai,0
49518,en,AF,Africa,RW,Rwanda,,,,,,,Africa/Kigali,0";
let geoip_db = GeoIPDB::new(blocks.as_bytes(), locations.as_bytes());
let block = geoip_db.resolve("1.3.4.2").unwrap();
assert_eq!("1.3.0.0/16", block.network.to_string());
assert_eq!(1809935, block.geoname_id);
assert_eq!("", block.postal_code);
assert_eq!(23.1167, block.latitude);
assert_eq!(113.25, block.longitude);
let location = geoip_db.get_location(block.geoname_id);
assert_eq!(1809935, location.geoname_id);
assert_eq!("AS", location.continent_code);
assert_eq!("Asia", location.continent_name);
assert_eq!("CN", location.country_code);
assert_eq!("China", location.country_name);
assert_eq!("GD", location.region_code);
assert_eq!("Guangdong", location.region_name);
assert_eq!("", location.province_code);
assert_eq!("", location.province_name);
assert_eq!("", location.city_name);
assert_eq!("Asia/Shanghai", location.timezone);
}
#[test]
fn cannot_resolve_ip() {
let geoip_db = GeoIPDB::new("".as_bytes(), "".as_bytes());
assert_eq!(true, geoip_db.resolve("1.2.3.4").is_none());
}
}