use flate2::read::GzDecoder; use log::warn; use maxminddb::geoip2::Country; use std::io::{Cursor, Read}; use std::net::Ipv6Addr; use tar::Archive; use tokio::sync::RwLock; pub struct MaxMindIndexer { pub reader: RwLock>>>, } impl MaxMindIndexer { pub async fn new() -> Result { let reader = MaxMindIndexer::inner_index(false).await.ok().flatten(); Ok(MaxMindIndexer { reader: RwLock::new(reader), }) } pub async fn index(&self) -> Result<(), reqwest::Error> { let reader = MaxMindIndexer::inner_index(false).await?; if let Some(reader) = reader { let mut reader_new = self.reader.write().await; *reader_new = Some(reader); } Ok(()) } async fn inner_index( should_panic: bool, ) -> Result>>, reqwest::Error> { let response = reqwest::get( format!( "https://download.maxmind.com/app/geoip_download?edition_id=GeoLite2-Country&license_key={}&suffix=tar.gz", dotenvy::var("MAXMIND_LICENSE_KEY").unwrap() ) ).await?.bytes().await.unwrap().to_vec(); let tarfile = GzDecoder::new(Cursor::new(response)); let mut archive = Archive::new(tarfile); if let Ok(entries) = archive.entries() { for mut file in entries.flatten() { if let Ok(path) = file.header().path() { if path.extension().and_then(|x| x.to_str()) == Some("mmdb") { let mut buf = Vec::new(); file.read_to_end(&mut buf).unwrap(); let reader = maxminddb::Reader::from_source(buf).unwrap(); return Ok(Some(reader)); } } } } if should_panic { panic!("Unable to download maxmind database- did you get a license key?") } else { warn!("Unable to download maxmind database."); Ok(None) } } pub async fn query(&self, ip: Ipv6Addr) -> Option { let maxmind = self.reader.read().await; if let Some(ref maxmind) = *maxmind { maxmind .lookup::(ip.into()) .ok() .and_then(|x| x.country.and_then(|x| x.iso_code.map(|x| x.to_string()))) } else { None } } }