feature(country_tracking): Part 2: Downloading and decompressing the json file

Cargo.toml

@@ -19,6 +19,8 @@ indicatif = { version = "0.17" }
 futures-util = "0.3"
 anyhow = "1.0.98"
 clap = { version = "4.5.37", features = ["derive"] }
+reqwest = { version = "0.12.19", features = ["stream", "gzip"] }
+flate2 = "1.1.1"
 
 [profile.release]
 strip = true

src/country_tracking.rs

@@ -0,0 +1,85 @@
+use crate::config::Config;
+use anyhow::bail;
+use flate2::read::GzDecoder;
+use futures_util::StreamExt;
+use indicatif::{ProgressBar, ProgressStyle};
+use reqwest::redirect::Policy;
+use std::fs::File;
+use std::io::{Read, Write};
+use std::time::Duration;
+use tracing::info;
+
+const DOWNLOAD_URL: &str = "https://ipinfo.io/data/ipinfo_lite.json.gz";
+
+pub async fn download_database(config: &Config) -> anyhow::Result<File> {
+	let client = reqwest::ClientBuilder::new()
+		.gzip(true)
+		.redirect(Policy::limited(3))
+		.connect_timeout(Duration::from_secs(10))
+		.build()?;
+
+	let url = format!(
+		"{}{}{}",
+		DOWNLOAD_URL, "?token=", config.country_tracking.ipinfo_token
+	);
+
+	// Send request
+	let response = client.get(url).send().await?;
+
+	// If response is OK write to file and unzip
+	if response.status().is_success() {
+		// Content length header is required for the progress bar
+		// IPInfo should always supply this
+		let content_length = match response.content_length() {
+			Some(len) => len,
+			None => bail!("Content-Length was not set!"),
+		};
+
+		let mut downloaded: u64 = 0;
+		let mut output_file = File::create("ipinfo.json.gz")?;
+		let mut reader = response.bytes_stream();
+
+		let style = ProgressStyle::with_template(
+			"[{elapsed_precise}] [{bar:40.white/blue}] {bytes}/{total_bytes}",
+		)
+		.expect("failed to create progress bar style")
+		.progress_chars("=>-");
+
+		let bar = ProgressBar::new(content_length).with_style(style);
+
+		while let Some(Ok(chunk)) = reader.next().await {
+			output_file.write_all(&chunk)?;
+
+			// Update download bar position
+			let new = std::cmp::min(downloaded + (chunk.len() as u64), content_length);
+			downloaded = new;
+			bar.set_position(new);
+		}
+
+		// Done
+		bar.finish_with_message("Finished downloading IPInfo database to ipinfo.json.gz");
+		info!("Decompressing output file...");
+
+		// Decompress file
+		let mut decoder = GzDecoder::new(File::open("ipinfo.json.gz")?);
+		let mut file = File::create("ipinfo.json")?;
+		let mut string = String::new();
+
+		// Write to output file
+		decoder.read_to_string(&mut string)?;
+		let mut output_file = File::create("ipinfo.json")?;
+		output_file.write_all(string.as_bytes())?;
+		file.flush()?;
+
+		// Delete compressed file
+		std::fs::remove_file("ipinfo.json.gz")?;
+
+		Ok(file)
+	} else {
+		bail!(
+			"IPInfo download failed: {} {:?}",
+			response.status(),
+			response.status().canonical_reason()
+		)
+	}
+}

src/main.rs

@@ -1,6 +1,7 @@
 #![feature(let_chains)]
 
 mod config;
+mod country_tracking;
 mod database;
 mod protocol;
 mod response;
@@ -63,11 +64,16 @@ async fn main() {
 		.await
 		.ok();
 
-	Scanner::new()
-		.config(config)
-		.mode(arguments.mode)
-		.pool(pool)
-		.build()
-		.start()
-		.await;
+	if config.country_tracking.enabled {
+		info!("Downloading latest version of IPInfo database...");
+		country_tracking::download_database(&config).await.unwrap();
+	}
+
+	// Scanner::new()
+	// 	.config(config)
+	// 	.mode(arguments.mode)
+	// 	.pool(pool)
+	// 	.build()
+	// 	.start()
+	// 	.await;
 }