From 0f7bccdbbccd548da9a54ff439efa28b0a8b3364 Mon Sep 17 00:00:00 2001 From: funtimes909 Date: Thu, 24 Apr 2025 00:56:01 +1200 Subject: [PATCH] Move scan results to a utils file, more config cleanups --- src/config.rs | 7 +++++++ src/main.rs | 3 ++- src/masscan.rs | 29 +++++++++++++++++++++++++---- src/scan.rs | 33 ++++----------------------------- src/utils.rs | 31 +++++++++++++++++++++++++++++++ 5 files changed, 69 insertions(+), 34 deletions(-) create mode 100644 src/utils.rs diff --git a/src/config.rs b/src/config.rs index 4ef20ec..4ad0292 100644 --- a/src/config.rs +++ b/src/config.rs @@ -9,6 +9,7 @@ pub struct Config { pub database: Database, pub player_tracking: PlayerTracking, pub scanner: Scanner, + pub masscan: Masscan, } #[derive(Deserialize)] @@ -34,6 +35,12 @@ pub struct Scanner { pub port_range_end: u16, } +#[derive(Deserialize)] +pub struct Masscan { + pub config_file: String, + pub output_file: String, +} + impl Default for Scanner { fn default() -> Self { Self { diff --git a/src/main.rs b/src/main.rs index c5bd688..e194a20 100644 --- a/src/main.rs +++ b/src/main.rs @@ -4,6 +4,7 @@ mod masscan; mod ping; mod response; mod scan; +mod utils; use config::load_config; use indicatif::ProgressStyle; @@ -48,7 +49,7 @@ async fn main() { .unwrap() .progress_chars("=>-"); - let mode = Mode::Rescan; + let mode = Mode::Discovery; match mode { Mode::Discovery => masscan::start(pool, config, style).await, diff --git a/src/masscan.rs b/src/masscan.rs index 71365c1..ce251c3 100644 --- a/src/masscan.rs +++ b/src/masscan.rs @@ -1,10 +1,12 @@ use crate::config::Config; use crate::scan; +use crate::utils::scan_results; use indicatif::{ProgressBar, ProgressStyle}; use serde::Deserialize; use sqlx::{Pool, Postgres}; use std::process::Command; use std::sync::Arc; +use std::time::Duration; use tokio::sync::Mutex; use tokio::task::JoinSet; use tracing::info; @@ -57,14 +59,17 @@ pub fn parse_output(masscan_output: String) -> Vec<(String, u16)> { } pub async fn start(pool: Pool, config: Config, style: ProgressStyle) { + let masscan_config = config.masscan.config_file; + let masscan_output = config.masscan.output_file; + 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()); + start_masscan(masscan_config.clone()); + let servers = parse_output(masscan_output.clone()); info!("{} servers found", servers.len()); let progress_bar = @@ -80,8 +85,24 @@ pub async fn start(pool: Pool, config: Config, style: ProgressStyle) { )); }); - join_set.join_all().await; + let results = join_set.join_all().await; - // TODO: Delay here + // Print information about scan + scan_results(results); + + // Quit if only one scan is requested in config + if !config.scanner.repeat { + info!("Exiting"); + std::process::exit(0); + } + + // Wait rescan delay before starting a new scan + if config.scanner.scan_delay > 0 { + info!( + "Waiting {} seconds before starting another scan", + config.scanner.scan_delay + ); + tokio::time::sleep(Duration::from_secs(config.scanner.scan_delay)).await; + } } } diff --git a/src/scan.rs b/src/scan.rs index b20df0f..7746059 100644 --- a/src/scan.rs +++ b/src/scan.rs @@ -1,4 +1,5 @@ use crate::config::Config; +use crate::utils::scan_results; use crate::{database, ping, response}; use futures_util::TryStreamExt; use indicatif::{ProgressBar, ProgressStyle}; @@ -53,35 +54,9 @@ pub async fn rescan_servers(pool: Pool, config: Config, style: Progres } let results = ping_set.join_all().await; - let results_len = results.len(); - // Save all errors for statistics - let errors = results - .into_iter() - .filter_map(Result::err) - .collect::>(); - - let errors_len = errors.len(); - - // Print scan errors - if !errors.is_empty() { - warn!("Scan returned {} errors!", errors.len()); - let mut counts = [0u32; 4]; - for e in errors { - let i: usize = e.into(); - counts[i] += 1; - } - - warn!("{} errors while pinging servers", counts[0]); - warn!("{} errors while parsing responses", counts[1]); - warn!("{} errors while updating the database", counts[2]); - warn!("{} connection timeouts", counts[3]); - } - - info!( - "Commiting {} results to database...", - results_len - errors_len - ); + // Print information about scan + scan_results(results); Arc::try_unwrap(transaction) .unwrap() @@ -99,7 +74,7 @@ pub async fn rescan_servers(pool: Pool, config: Config, style: Progres // Quit if only one scan is requested in config if !config.scanner.repeat { - info!("Exiting..."); + info!("Exiting"); std::process::exit(0); } diff --git a/src/utils.rs b/src/utils.rs new file mode 100644 index 0000000..d74ba51 --- /dev/null +++ b/src/utils.rs @@ -0,0 +1,31 @@ +use crate::scan::RunError; +use tracing::{info, warn}; + +pub fn scan_results(results: Vec>) { + let results_len = results.len(); + + // Save all errors for statistics + let errors = results + .into_iter() + .filter_map(Result::err) + .collect::>(); + + let errors_len = errors.len(); + + // Print scan errors + if !errors.is_empty() { + warn!("Scan returned {} errors!", errors.len()); + let mut counts = [0u32; 4]; + for e in errors { + let i: usize = e.into(); + counts[i] += 1; + } + + warn!("{} errors while pinging servers", counts[0]); + warn!("{} errors while parsing responses", counts[1]); + warn!("{} errors while updating the database", counts[2]); + warn!("{} connection timeouts", counts[3]); + } + + info!("Commiting {} results to database", results_len - errors_len); +}