Merge branch 'database-address-information'

# Conflicts:
#	src/database.rs

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

README.md

@@ -51,6 +51,7 @@ ServerSeekerV2 is written in Rust allowing it to be blazingly fast and memory sa
 - Player and mod tracking. Find servers that have specific players online or servers running specific forge mods. (or both at the same time!)
 - Self Hostable. Host your own scanning instance and find your own servers! (See below for warnings against running this
   on a residential network)
+- Country tracking. If enabled, tracks which country and Autonomous System a server is from.
 
 # For people just looking to not be scanned anymore
 

config-example.toml

@@ -8,12 +8,6 @@ port_range_start = 25565
 # End of port range
 port_range_end = 25565
 
-[player_tracking]
-# Enable player tracking?
-enabled = false
-# List of players to track
-players = [""]
-
 [masscan]
 # Location of masscans config file
 config_file = "masscan.conf"
@@ -28,4 +22,18 @@ table = "postgres"
 # Database user
 user = "postgres"
 # Database user password
-password = "SuperSuperSecretPassword"
+password = "SuperSuperSecretPassword"
+
+[player_tracking]
+# Enable player tracking?
+enabled = false
+# List of players to track
+players = [""]
+
+[country_tracking]
+# Enable tracking of country information (requires a valid https://ipinfo.io account token)
+enabled = false
+# How often to update the country information table in hours (minimum 24 hours)
+update_frequency = 48
+# IPInfo token
+ipinfo_token = ":3"

src/config.rs

@@ -4,15 +4,16 @@ use std::fs::File;
 use std::io::{ErrorKind, Read};
 use tracing::error;
 
-#[derive(Deserialize, Debug)]
+#[derive(Deserialize, Clone, Debug)]
 pub struct Config {
 	pub database: Database,
-	pub player_tracking: PlayerTracking,
 	pub scanner: ScannerConfig,
 	pub masscan: Masscan,
+	pub player_tracking: PlayerTracking,
+	pub country_tracking: CountryTracking,
 }
 
-#[derive(Deserialize, Debug)]
+#[derive(Deserialize, Clone, Debug)]
 pub struct Database {
 	pub host: String,
 	pub port: u16,
@@ -21,13 +22,7 @@ pub struct Database {
 	pub password: String,
 }
 
-#[derive(Deserialize, Debug)]
-pub struct PlayerTracking {
-	pub enabled: bool,
-	pub players: Vec<String>,
-}
-
-#[derive(Deserialize, Debug)]
+#[derive(Deserialize, Clone, Debug)]
 pub struct ScannerConfig {
 	pub repeat: bool,
 	pub scan_delay: u64,
@@ -35,11 +30,24 @@ pub struct ScannerConfig {
 	pub port_range_end: u16,
 }
 
-#[derive(Deserialize, Debug)]
+#[derive(Deserialize, Clone, Debug)]
 pub struct Masscan {
 	pub config_file: String,
 }
 
+#[derive(Deserialize, Clone, Debug)]
+pub struct PlayerTracking {
+	pub enabled: bool,
+	pub players: Vec<String>,
+}
+
+#[derive(Deserialize, Clone, Debug)]
+pub struct CountryTracking {
+	pub enabled: bool,
+	pub update_frequency: u64,
+	pub ipinfo_token: String,
+}
+
 impl Default for Config {
 	fn default() -> Self {
 		Config {
@@ -50,10 +58,6 @@ impl Default for Config {
 				user: "postgres".to_string(),
 				password: "password".to_string(),
 			},
-			player_tracking: PlayerTracking {
-				enabled: false,
-				players: vec![],
-			},
 			scanner: ScannerConfig {
 				repeat: true,
 				scan_delay: 60,
@@ -63,6 +67,15 @@ impl Default for Config {
 			masscan: Masscan {
 				config_file: "masscan.conf".to_string(),
 			},
+			player_tracking: PlayerTracking {
+				enabled: false,
+				players: vec![],
+			},
+			country_tracking: CountryTracking {
+				enabled: false,
+				update_frequency: 48,
+				ipinfo_token: "".to_string(),
+			},
 		}
 	}
 }

src/country_tracking.rs

@@ -0,0 +1,199 @@
+use crate::config::Config;
+use anyhow::bail;
+use flate2::read::GzDecoder;
+use futures_util::StreamExt;
+use indicatif::{ProgressBar, ProgressIterator, ProgressStyle};
+use serde::Deserialize;
+use sqlx::types::ipnet::IpNet;
+use sqlx::PgPool;
+use std::fs::File;
+use std::io::{Read, Write};
+use std::str::FromStr;
+use std::time::Duration;
+use tracing::{debug, info};
+
+const DOWNLOAD_URL: &str = "https://ipinfo.io/data/ipinfo_lite.json.gz?token=";
+
+#[derive(Deserialize, Debug)]
+struct CountryRow {
+	network: String,
+	country: String,
+	country_code: String,
+	asn: Option<String>,
+	#[serde(rename = "as_name")]
+	company: Option<String>,
+}
+
+pub async fn country_tracking(pool: PgPool, config: Config) -> anyhow::Result<()> {
+	loop {
+		download_database(&config).await?;
+		insert_json_to_table(&pool).await?;
+
+		// Sleep
+		tokio::time::sleep(Duration::from_secs(
+			config.country_tracking.update_frequency * 60 * 60,
+		))
+		.await;
+	}
+}
+
+async fn download_database(config: &Config) -> anyhow::Result<()> {
+	let url = format!("{}{}", DOWNLOAD_URL, config.country_tracking.ipinfo_token);
+	let response = reqwest::get(url).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 header 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} {msg}",
+		)
+		.expect("failed to create progress bar style")
+		.progress_chars("=>-");
+
+		let bar = ProgressBar::new(content_length).with_style(style);
+		bar.set_message("Downloading the latest version of the IPInfo database...");
+
+		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!");
+		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(())
+	} else {
+		bail!(
+			"IPInfo download failed: {} {:?}",
+			response.status(),
+			response.status().canonical_reason()
+		)
+	}
+}
+
+async fn parse_json_to_vec(string: String) -> serde_json::Result<Vec<CountryRow>> {
+	serde_json::from_str(&format!(
+		"[{}]",
+		string
+			// Split at the end of every object
+			.split("}\n{")
+			// Skip all IPv6 netblocks
+			.map_while(|x| {
+				match x.contains("::") {
+					true => None,
+					false => Some(x),
+				}
+			})
+			.map(|s| s.trim_matches(&['\n', '{', '}'][..]))
+			.map(|s| format!("{{{}}}", s))
+			// Collect everything
+			.collect::<Vec<_>>()
+			// Join everything with commas
+			.join(",")
+	))
+}
+
+pub async fn create_tables(pool: &PgPool) -> Result<(), sqlx::Error> {
+	// Create table
+	sqlx::query(
+		"CREATE TABLE IF NOT EXISTS countries (
+    		network CIDR,
+    		country VARCHAR(255),
+    		country_code VARCHAR(2),
+    		asn VARCHAR(16),
+    		company VARCHAR(255),
+    		PRIMARY KEY(network)
+    	)",
+	)
+	.execute(pool)
+	.await?;
+
+	// Use a GIST inet_ops index for the network column
+	// This allows for really fast lookup times from my testing
+	//
+	// Countries table needs to exist before we can create an index on it
+	sqlx::query("CREATE INDEX IF NOT EXISTS countries_table_index ON countries USING GIST (network inet_ops);")
+		.execute(pool)
+		.await?;
+
+	Ok(())
+}
+
+async fn insert_json_to_table(pool: &PgPool) -> anyhow::Result<()> {
+	let mut file = File::open("ipinfo.json")?;
+	let mut string = String::new();
+	file.read_to_string(&mut string)?;
+
+	let json = parse_json_to_vec(string).await?;
+	info!("JSON Parsed successfully.");
+
+	let mut transaction = pool.begin().await?;
+
+	let style = ProgressStyle::with_template(
+		"[{elapsed_precise}] [{bar:40.white/blue}] {human_pos}/{human_len} {msg}",
+	)
+	.expect("failed to create progress bar style")
+	.progress_chars("=>-");
+
+	let bar = ProgressBar::new(json.len() as u64).with_style(style);
+	bar.set_message("Inserting rows to countries table...");
+
+	for netblock in json.into_iter().progress_with(bar) {
+		if let Ok(cidr) = IpNet::from_str(&netblock.network) {
+			let result = sqlx::query(
+				"INSERT INTO countries VALUES ($1, $2, $3, $4, $5) 
+					ON CONFLICT (network, country, country_code) DO UPDATE 
+					network = EXCLUDED.network,
+					country = EXCLUDED.country,
+					country_code = EXCLUDED.country_code,
+					asn = EXCLUDED.asn,
+					company = EXCLUDED.company",
+			)
+			.bind(cidr)
+			.bind(netblock.country)
+			.bind(netblock.country_code)
+			.bind(netblock.asn)
+			.bind(netblock.company)
+			.execute(&mut *transaction)
+			.await;
+
+			if let Err(e) = result {
+				debug!("Error while updating row in countries table {e}");
+			}
+		};
+	}
+
+	transaction.commit().await?;
+
+	info!("All done!");
+	Ok(())
+}

src/database.rs

@@ -1,12 +1,29 @@
 use crate::response::Server;
 use crate::utils::RunError;
-use sqlx::postgres::PgQueryResult;
+use sqlx::postgres::{PgQueryResult, PgRow};
 use sqlx::types::ipnet::{IpNet, Ipv4Net};
 use sqlx::types::Uuid;
-use sqlx::{PgPool, Row};
+use sqlx::{FromRow, PgPool, Row};
 use std::net::SocketAddrV4;
 use std::time::{SystemTime, UNIX_EPOCH};
 
+#[derive(Debug)]
+struct AddressInfo {
+	country: String,
+	country_code: String,
+	asn: String,
+}
+
+impl FromRow<'_, PgRow> for AddressInfo {
+	fn from_row(row: &PgRow) -> Result<Self, sqlx::Error> {
+		Ok(Self {
+			country: row.try_get("country")?,
+			country_code: row.try_get("country_code")?,
+			asn: row.try_get("asn")?,
+		})
+	}
+}
+
 #[derive(Debug, Clone)]
 pub struct Database(pub PgPool);
 
@@ -26,13 +43,20 @@ impl Database {
 	}
 
 	/// Deletes a server from the database
-	pub async fn delete_server(&self, address: IpNet) -> Result<PgQueryResult, sqlx::Error> {
+	async fn delete_server(&self, address: IpNet) -> Result<PgQueryResult, sqlx::Error> {
 		sqlx::query("DELETE FROM servers WHERE address = $1")
 			.bind(address)
 			.execute(&self.0)
 			.await
 	}
 
+	async fn get_address_info(&self, address: &IpNet) -> Result<AddressInfo, sqlx::Error> {
+		sqlx::query_as("SELECT country, country_code, asn FROM countries WHERE $1 <<= network")
+			.bind(address)
+			.fetch_one(&self.0)
+			.await
+	}
+
 	/// Updates a single server in the database, this includes all mods
 	/// and players that come with it. Will also remove a server from the
 	/// database if it has requested to be removed
@@ -52,6 +76,8 @@ impl Database {
 			return Err(RunError::ServerOptOut)?;
 		}
 
+		let address_information = self.get_address_info(&address).await?;
+
 		sqlx::query(
 			"INSERT INTO servers (
 		address,
@@ -67,7 +93,9 @@ impl Database {
 		first_seen,
 		last_seen,
 		online_players,
-		max_players) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14)
+		max_players,
+        country,
+    	asn) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15, $16)
     	ON CONFLICT (address, port) DO UPDATE SET
     	software = EXCLUDED.software,
     	version = EXCLUDED.version,
@@ -79,7 +107,9 @@ impl Database {
     	enforces_secure_chat = EXCLUDED.enforces_secure_chat,
     	last_seen = EXCLUDED.last_seen,
     	online_players = EXCLUDED.online_players,
-    	max_players = EXCLUDED.max_players",
+    	max_players = EXCLUDED.max_players,
+    	country = EXCLUDED.country,
+    	asn = EXCLUDED.asn",
 		)
 		.bind(address)
 		.bind(socket.port() as i32)
@@ -99,6 +129,8 @@ impl Database {
 		.bind(timestamp)
 		.bind(server.players.online)
 		.bind(server.players.max)
+		.bind(address_information.country_code)
+		.bind(address_information.asn)
 		.execute(&self.0)
 		.await?;
 

src/main.rs

@@ -1,6 +1,7 @@
 #![feature(let_chains)]
 
 mod config;
+mod country_tracking;
 mod database;
 mod protocol;
 mod response;
@@ -15,7 +16,7 @@ use sqlx::postgres::{PgConnectOptions, PgPoolOptions};
 use sqlx::ConnectOptions;
 use std::time::Duration;
 use tracing::log::LevelFilter;
-use tracing::{error, info};
+use tracing::{debug, error, info};
 
 #[derive(Parser, Debug)]
 #[clap(about = "Scans the internet for minecraft servers and indexes them")]
@@ -64,6 +65,28 @@ async fn main() {
 		.await
 		.ok();
 
+	if pool.is_none() {
+		error!("Failed to connect to database");
+		std::process::exit(1);
+	}
+
+	// Spawn a task to update the country info database everyday
+	if let Some(pool) = &pool
+		&& config.country_tracking.enabled
+	{
+		if country_tracking::create_tables(pool).await.is_err() {
+			error!("failed to create tables");
+			std::process::exit(1);
+		}
+
+		debug!("Spawning country tracking task");
+
+		tokio::task::spawn(country_tracking::country_tracking(
+			pool.clone(),
+			config.clone(),
+		));
+	}
+
 	Scanner::new()
 		.config(config)
 		.mode(arguments.mode)

src/scanner.rs

@@ -100,7 +100,7 @@ impl Scanner {
 			let (tx, mut rx) = tokio::sync::mpsc::channel::<SocketAddrV4>(10);
 
 			let mut stream = sqlx::query(
-				"SELECT (address - '0.0.0.0'::inet) AS address FROM servers ORDER BY last_seen ASC",
+				"SELECT (address - '0.0.0.0'::inet) AS address FROM servers ORDER BY last_seen DESC LIMIT 1000",
 			)
 			.fetch(&self.database.0);
 
@@ -135,9 +135,7 @@ impl Scanner {
 			.expect("failed to create progress bar style")
 			.progress_chars("=>-");
 
-			let bar =
-				ProgressBar::new((total_servers as u64) * self.config.scanner.total_ports() as u64)
-					.with_style(style);
+			let bar = ProgressBar::new(1000).with_style(style);
 
 			// Consume values from the receiver
 			while let Some(socket) = rx.recv().await {