use std::{io, path::PathBuf, time::Duration}; use color_eyre::eyre::{WrapErr as _, eyre}; use tokio::io::AsyncWriteExt as _; #[cfg(feature = "tui")] use crate::event::{AppEvent, Event}; use crate::{fs::get_cache_path, utils::is_container}; #[derive(Clone, Copy)] pub enum DbType { Asn, Geo, } impl DbType { pub async fn download( self, http_client: reqwest_middleware::ClientWithMiddleware, #[cfg(feature = "tui")] tx: tokio::sync::mpsc::UnboundedSender, ) -> crate::Result<()> { let db_path = self.db_path().await?; let mut headers = reqwest::header::HeaderMap::new(); if tokio::fs::metadata(&db_path).await.is_ok_and(|m| m.is_file()) && let Some(etag) = self.read_etag().await? { headers.insert(reqwest::header::IF_NONE_MATCH, etag); } let request = http_client.get(self.url()).headers(headers).timeout(Duration::MAX); drop(http_client); let response = request.send().await?.error_for_status()?; let status = response.status(); if status == reqwest::StatusCode::NOT_MODIFIED { drop(response); tracing::info!( "Latest {} database is already cached at {}", self.name(), db_path.display() ); return Ok(()); } if status != reqwest::StatusCode::OK { return Err(eyre!( "HTTP status error ({}) for url ({})", status, response.url(), )); } let etag = response.headers().get(reqwest::header::ETAG).cloned(); self.save_db( response, #[cfg(feature = "tui")] tx, ) .await?; if is_container().await { tracing::info!( "Downloaded {} database to Docker volume ({} in container)", self.name(), db_path.display() ); } else { tracing::info!( "Downloaded {} database to {}", self.name(), db_path.display() ); } drop(db_path); if let Some(etag_value) = etag { self.save_etag(etag_value).await } else { self.remove_etag().await } } pub async fn open_mmap( self, ) -> crate::Result> { let path = self.db_path().await?; tokio::task::spawn_blocking(move || { #[expect(clippy::undocumented_unsafe_blocks)] unsafe { maxminddb::Reader::open_mmap(&path) }.wrap_err_with( move || { compact_str::format_compact!( "failed to open IP database: {}", path.display() ) }, ) }) .await? } const fn name(self) -> &'static str { match self { Self::Asn => "ASN", Self::Geo => "geolocation", } } const fn url(self) -> &'static str { match self { Self::Asn => { "https://raw.githubusercontent.com/P3TERX/GeoLite.mmdb/refs/heads/download/GeoLite2-ASN.mmdb" } Self::Geo => { "https://raw.githubusercontent.com/P3TERX/GeoLite.mmdb/refs/heads/download/GeoLite2-City.mmdb" } } } async fn db_path(self) -> crate::Result { let mut cache_path = get_cache_path().await?; cache_path.push(match self { Self::Asn => "asn_database.mmdb", Self::Geo => "geolocation_database.mmdb", }); Ok(cache_path) } async fn etag_path(self) -> crate::Result { let mut db_path = self.db_path().await?; db_path.set_extension("mmdb.etag"); Ok(db_path) } async fn save_db( self, mut response: reqwest::Response, #[cfg(feature = "tui")] tx: tokio::sync::mpsc::UnboundedSender, ) -> crate::Result<()> { #[cfg(feature = "tui")] drop(tx.send(Event::App(AppEvent::IpDbTotal( self, response.content_length(), )))); let db_path = self.db_path().await?; let file = tokio::fs::File::create(&db_path).await.wrap_err_with(|| { compact_str::format_compact!( "failed to create file: {}", db_path.display() ) })?; let mut writer = tokio::io::BufWriter::new(file); while let Some(chunk) = response.chunk().await? { writer.write_all(&chunk).await.wrap_err_with(|| { compact_str::format_compact!( "failed to write to file: {}", db_path.display() ) })?; #[cfg(feature = "tui")] drop( tx.send(Event::App(AppEvent::IpDbDownloaded( self, chunk.len(), ))), ); } drop(response); #[cfg(feature = "tui")] drop(tx); writer.flush().await.wrap_err_with(move || { compact_str::format_compact!( "failed to write to file: {}", db_path.display() ) })?; Ok(()) } async fn save_etag(self, etag: impl AsRef<[u8]>) -> crate::Result<()> { let path = self.etag_path().await?; tokio::fs::write(&path, etag).await.wrap_err_with(move || { compact_str::format_compact!( "failed to write to file: {}", path.display() ) }) } async fn read_etag( self, ) -> crate::Result> { let path = self.etag_path().await?; match tokio::fs::read_to_string(&path).await { Ok(text) => Ok(text.parse().ok()), Err(e) if e.kind() == io::ErrorKind::NotFound => Ok(None), Err(e) => Err(e).wrap_err_with(move || { compact_str::format_compact!( "failed to read file to string: {}", path.display() ) }), } } async fn remove_etag(self) -> crate::Result<()> { let path = self.etag_path().await?; match tokio::fs::remove_file(&path).await { Ok(()) => Ok(()), Err(e) if e.kind() == io::ErrorKind::NotFound => Ok(()), Err(e) => Err(e).wrap_err_with(move || { compact_str::format_compact!( "failed to remove file: {}", path.display() ) }), } } }