From 4412a04760ea34cedf447ea5c88e88176a01d8d9 Mon Sep 17 00:00:00 2001
From: funtimes909 <funtimes90969@proton.me>
Date: Wed, 23 Apr 2025 22:08:09 +1200
Subject: [PATCH] Scan servers found by masscan using a JoinSet

---
 src/main.rs    | 26 +++++++++++++++++++++++---
 src/masscan.rs | 41 +++++++++++++++++++++++++++++++++++++++++
 src/scan.rs    | 15 ++++-----------
 3 files changed, 68 insertions(+), 14 deletions(-)

diff --git a/src/main.rs b/src/main.rs
index 71d2ace..c5bd688 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -5,11 +5,20 @@ mod ping;
 mod response;
 mod scan;
 
-use crate::scan::scan_servers;
 use config::load_config;
+use indicatif::ProgressStyle;
 use sqlx::PgPool;
+use std::time::Duration;
 use tracing::error;
 
+enum Mode {
+	Discovery,
+	Rescan,
+}
+
+// TODO: add to config file
+const TIMEOUT_SECS: Duration = Duration::from_secs(5);
+
 #[tokio::main]
 async fn main() {
 	tracing_subscriber::fmt::init();
@@ -27,11 +36,22 @@ async fn main() {
 		config.database.table
 	);
 
-	match PgPool::connect(&database_url).await {
-		Ok(pool) => scan_servers(pool, config).await,
+	let pool = match PgPool::connect(&database_url).await {
+		Ok(pool) => pool,
 		Err(e) => {
 			error!("Failed to connect to database: {e}");
 			std::process::exit(1);
 		}
 	};
+
+	let style = ProgressStyle::with_template("[{elapsed}] [{bar:40.white/blue}] {pos:>7}/{len:7}")
+		.unwrap()
+		.progress_chars("=>-");
+
+	let mode = Mode::Rescan;
+
+	match mode {
+		Mode::Discovery => masscan::start(pool, config, style).await,
+		Mode::Rescan => scan::rescan_servers(pool, config, style).await,
+	}
 }
diff --git a/src/masscan.rs b/src/masscan.rs
index f502eb3..71365c1 100644
--- a/src/masscan.rs
+++ b/src/masscan.rs
@@ -1,5 +1,13 @@
+use crate::config::Config;
+use crate::scan;
+use indicatif::{ProgressBar, ProgressStyle};
 use serde::Deserialize;
+use sqlx::{Pool, Postgres};
 use std::process::Command;
+use std::sync::Arc;
+use tokio::sync::Mutex;
+use tokio::task::JoinSet;
+use tracing::info;
 
 #[derive(Deserialize)]
 pub struct Masscan {
@@ -13,6 +21,7 @@ pub struct Port {
 }
 
 pub fn start_masscan(masscan_config: String) {
+	info!("Starting masscan...");
 	if cfg!(target_os = "windows") {
 		Command::new("cmd.exe")
 			.arg("/c")
@@ -30,6 +39,8 @@ pub fn start_masscan(masscan_config: String) {
 	}
 	.wait()
 	.expect("failed to wait on child");
+
+	info!("Masscan completed");
 }
 
 pub fn parse_output(masscan_output: String) -> Vec<(String, u16)> {
@@ -44,3 +55,33 @@ pub fn parse_output(masscan_output: String) -> Vec<(String, u16)> {
 		})
 		.collect()
 }
+
+pub async fn start(pool: Pool<Postgres>, config: Config, style: ProgressStyle) {
+	let transaction = Arc::new(Mutex::new(
+		pool.begin().await.expect("failed to create transaction"),
+	));
+
+	loop {
+		// TODO: Config changes needed here
+		start_masscan("masscan.conf".to_string());
+		let servers = parse_output("minecraft-servers.json".to_string());
+		info!("{} servers found", servers.len());
+
+		let progress_bar =
+			Arc::new(ProgressBar::new(servers.len() as u64).with_style(style.clone()));
+
+		let mut join_set = JoinSet::new();
+
+		servers.into_iter().for_each(|(ip, port)| {
+			join_set.spawn(scan::run(
+				(ip, port),
+				transaction.clone(),
+				progress_bar.clone(),
+			));
+		});
+
+		join_set.join_all().await;
+
+		// TODO: Delay here
+	}
+}
diff --git a/src/scan.rs b/src/scan.rs
index 7da0f8e..546a1fc 100644
--- a/src/scan.rs
+++ b/src/scan.rs
@@ -10,7 +10,7 @@ use tokio::sync::{Mutex, Semaphore};
 use tokio::task::JoinSet;
 use tracing::{info, warn};
 
-pub async fn scan_servers(pool: Pool<Postgres>, config: Config) {
+pub async fn rescan_servers(pool: Pool<Postgres>, config: Config, style: ProgressStyle) {
 	let port_start = config.rescanner.port_range_start;
 	let port_end = config.rescanner.port_range_end;
 	let total_ports = config.rescanner.total_ports();
@@ -25,10 +25,6 @@ pub async fn scan_servers(pool: Pool<Postgres>, config: Config) {
 
 	info!("Scanning port range {port_start} - {port_end} ({total_ports} port(s) per host)");
 
-	let style = ProgressStyle::with_template("[{elapsed}] [{bar:40.white/blue}] {pos:>7}/{len:7}")
-		.unwrap()
-		.progress_chars("=>-");
-
 	loop {
 		let mut servers = database::fetch_servers(&pool).await;
 		let length = database::fetch_count(&pool).await as u64;
@@ -119,7 +115,7 @@ pub async fn scan_servers(pool: Pool<Postgres>, config: Config) {
 }
 
 #[derive(Debug, Error)]
-enum RunError {
+pub enum RunError {
 	#[error("Error while pinging server")]
 	PingServer(#[from] ping::PingServerError),
 	#[error("Error while parsing response")]
@@ -140,12 +136,9 @@ impl Into<usize> for RunError {
 	}
 }
 
-// TODO: add to config file
-const TIMEOUT_SECS: Duration = Duration::from_secs(5);
-
 static PERMITS: Semaphore = Semaphore::const_new(2000);
 
-async fn run(
+pub async fn run(
 	host: (String, u16),
 	transaction: Arc<Mutex<PgTransaction<'_>>>,
 	progress_bar: Arc<ProgressBar>,
@@ -155,7 +148,7 @@ async fn run(
 		transaction: Arc<Mutex<PgTransaction<'_>>>,
 	) -> Result<(), RunError> {
 		let permit = PERMITS.acquire().await.unwrap();
-		let results = tokio::time::timeout(TIMEOUT_SECS, ping::ping_server(&host)).await??;
+		let results = tokio::time::timeout(crate::TIMEOUT_SECS, ping::ping_server(&host)).await??;
 		drop(permit);
 
 		let response = response::parse_response(results)?;