echoip-slatecave/src/geoip.rs

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

use maxminddb;
use maxminddb::geoip2;

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 {
	latitude: f64,
	logtitude: 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>,

	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: Option<maxminddb::Reader<Vec<u8>>>,
	pub name: String,
}

pub trait QueryLocation {
	fn query_location_for_ip(&self, address: &IpAddr, laguages: &Vec<&String>) -> 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: &Vec<&String>)
-> Option<String> {
	match names {
		Some(names) => {
			for language in languages {
				if let Some(name) = names.get(language.as_str()){
					return Some(name.to_string())
				}
			}
			return None
		},
		None => None
	}	
}

pub fn geoip2_city_to_named_location(item: geoip2::model::City, languages: &Vec<&String>) -> 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::model::Continent, languages: &Vec<&String>) -> 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::model::Country, languages: &Vec<&String>) -> 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::model::RepresentedCountry, languages: &Vec<&String>) -> 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::model::Subdivision, languages: &Vec<&String>) -> 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> {
		match &self.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(e) => {
						println!("Error while looking up ASN for {address}: {e}");
						Default::default()
					}
				}
			},
			None => None,
		}
	}
}

impl QueryLocation for MMDBCarrier {
	fn query_location_for_ip(&self, address: &IpAddr, languages: &Vec<&String>) -> Option<LocationResult> {
		match &self.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,
							},
							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(e) => {
						println!("Error while looking up ASN for {address}: {e}");
						Default::default()
					}
				}
			},
			None => None,
		}
	}
}

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