echoip-slatecave/src/geoip.rs

/*
 * This module provides an abstraction over the maxmind geoip databases
 * that provides the results ready for templating.
 */

use log::{debug,info,warn,error};
use maxminddb::geoip2;

use maxminddb::MaxMindDBError::AddressNotFoundError;
use parking_lot::RwLock;

use std::collections::BTreeMap;
use std::net::IpAddr;
use std::path::Path;

/* Datatypes */

#[derive(serde::Deserialize, serde::Serialize, Default, Clone)]
pub struct NamedLocation {
	iso_code: Option<String>,
	name: Option<String>,
	geoname_id: Option<u32>,
}

#[derive(serde::Deserialize, serde::Serialize, Default, Copy, Clone)]
pub struct LocationCoordinates {
	lat: f64,
	lon: f64,
}

#[derive(serde::Deserialize, serde::Serialize, Default, Clone)]
pub struct LocationResult {
	continent:           Option<NamedLocation>,
	country:             Option<NamedLocation>,
	registered_country:  Option<NamedLocation>,
	represented_country: Option<NamedLocation>,
	subdivisions:        Option<Vec<NamedLocation>>,
	city:                Option<NamedLocation>,

	coordinates:         Option<LocationCoordinates>,
	accuracy:            Option<u16>,
	postal_code:         Option<String>,
	time_zone:           Option<String>,
}

#[derive(serde::Deserialize, serde::Serialize, Default, Clone)]
pub struct AsnResult {
	asn:  Option<u32>,
	name: Option<String>,
}

pub struct MMDBCarrier {
	pub mmdb: RwLock<Option<maxminddb::Reader<Vec<u8>>>>,
	pub name: String,
	pub path: Option<String>,
}

pub trait QueryLocation {
	fn query_location_for_ip(&self, address: &IpAddr, laguages: &[&str]) -> Option<LocationResult>;
}

pub trait QueryAsn {
	fn query_asn_for_ip(&self, address: &IpAddr) -> Option<AsnResult>;
}

/* Converters */

pub fn extract_localized_name(
names: &Option<BTreeMap<&str, &str>>,
	languages: &[&str])
-> Option<String> {
	match names {
		Some(names) => {
			for language in languages {
				if let Some(name) = names.get(language){
					return Some(name.to_string())
				}
			}
			return None
		},
		None => None
	}	
}

pub fn geoip2_city_to_named_location(item: geoip2::city::City, languages: &[&str]) -> NamedLocation {
	NamedLocation {
		iso_code: None,
		geoname_id: item.geoname_id,
		name: extract_localized_name(&item.names, languages),
	}
}

pub fn geoip2_continent_to_named_location(item: geoip2::country::Continent, languages: &[&str]) -> NamedLocation {
	NamedLocation {
		iso_code: item.code.map(ToString::to_string),
		geoname_id: item.geoname_id,
		name: extract_localized_name(&item.names, languages),
	}
}

pub fn geoip2_country_to_named_location(item: geoip2::country::Country, languages: &[&str]) -> NamedLocation {
	NamedLocation {
		iso_code: item.iso_code.map(ToString::to_string),
		geoname_id: item.geoname_id,
		name: extract_localized_name(&item.names, languages),
	}
}

pub fn geoip2_represented_country_to_named_location(item: geoip2::country::RepresentedCountry, languages: &[&str]) -> NamedLocation {
	NamedLocation {
		iso_code: item.iso_code.map(ToString::to_string),
		geoname_id: item.geoname_id,
		name: extract_localized_name(&item.names, languages),
	}
}

pub fn geoip2_subdivision_to_named_location(item: geoip2::city::Subdivision, languages: &[&str]) -> NamedLocation {
	NamedLocation {
		iso_code: item.iso_code.map(ToString::to_string),
		geoname_id: item.geoname_id,
		name: extract_localized_name(&item.names, languages),
	}
}

/* Implementation */

impl QueryAsn for MMDBCarrier {
	fn query_asn_for_ip(&self, address: &IpAddr) -> Option<AsnResult> {
		let mmdb = self.mmdb.read();
		match &*mmdb {
			Some(mmdb) => {
				match  mmdb.lookup::<geoip2::Asn>(*address) {
					Ok(res) => {
						Some(AsnResult {
							asn: res.autonomous_system_number,
							name: res.autonomous_system_organization.map(ToString::to_string),
						})
					},
					Err(AddressNotFoundError(_)) => {
						// Log to the debug channel.
						// This isn't severe, and shouldn't be logged in production.
						debug!("ASN not found in database for {address}.");
						None
					},
					Err(e) => {
						error!("Error while looking up ASN for {address}: {e}");
						None
					}
				}
			},
			None => None,
		}
	}
}

impl QueryLocation for MMDBCarrier {
	fn query_location_for_ip(&self, address: &IpAddr, languages: &[&str]) -> Option<LocationResult> {
		let mmdb = self.mmdb.read();
		match &*mmdb {
			Some(mmdb) => {
				match  mmdb.lookup::<geoip2::City>(*address) {
					Ok(res) => {
						Some(LocationResult {
							continent:
								res.continent.map(|c| geoip2_continent_to_named_location(c, languages)),
							country:
								res.country.map(|c| geoip2_country_to_named_location(c, languages)),
							registered_country:
								res.registered_country.map(|c| geoip2_country_to_named_location(c, languages)),
							represented_country:
								res.represented_country.map(|c| geoip2_represented_country_to_named_location(c, languages)),
							city:
								res.city.map(|c| geoip2_city_to_named_location(c, languages)),
								
							subdivisions: match res.subdivisions {
								Some(sds) => {
									let mut subdivisions = Vec::new();
									subdivisions.reserve_exact(sds.len());
									for sd in sds {
										subdivisions.push(geoip2_subdivision_to_named_location(sd, languages));
									}
									Some(subdivisions)
								},
								None => None,
							},
							coordinates: match &res.location {
								Some(loc) => {
									if loc.latitude.is_some() && loc.longitude.is_some() {
										Some(LocationCoordinates {
											lat: loc.latitude.unwrap_or(0.0),
											lon: loc.longitude.unwrap_or(0.0),
										})
									} else {
										None
									}
								},
								None => None,
							},
							accuracy: match &res.location {
								Some(loc) => loc.accuracy_radius,
								None => None,
							},
							postal_code: match res.postal {
								Some(p) => p.code.map(ToString::to_string),
								None => None,
							},
							time_zone:   match res.location {
								Some(loc) => loc.time_zone.map(ToString::to_string),
								None => None,
							},
						})
					},
					Err(AddressNotFoundError(_)) => {
						// Log to the debug channel.
						// This isn't severe, and shouldn't be logged in production.
						debug!("IP location not found in database for {address}");
						None
					},
					Err(e) => {
						error!("Error while looking up IP location for {address}: {e}");
						None
					}
				}
			},
			None => None,
		}
	}
}

impl MMDBCarrier {
	pub fn new(name: String, path: Option<String>) -> MMDBCarrier {
		MMDBCarrier {
			mmdb: RwLock::new(None),
			name: name,
			path: path,
		}
	}

	pub fn reload_database(&self) -> Result<(),maxminddb::MaxMindDBError> {
		match &self.path {
			Some(path) => self.load_database_from_path(Path::new(&path)),
			None => Ok(()),
		}
	}

	pub fn load_database_from_path(&self, path: &Path) -> Result<(),maxminddb::MaxMindDBError> {
		let mut mmdb = self.mmdb.write();
		info!("Loading {} from '{}' ...", &self.name, path.display());
		match maxminddb::Reader::open_readfile(path) {
			Ok(reader) => {
				let wording = if mmdb.is_some() {
					"Replaced old"
				} else {
					"Loaded new"
				};
				*mmdb = Some(reader);
				info!("{} {} with new one.", wording, &self.name);
				Ok(())
			},
			Err(e) => {
				error!("Error while reading {}: {}", &self.name, &e);
				if mmdb.is_some() {
					warn!("Not replacing old database.");
				}
				Err(e)
			},
		}
	}
}