diff --git a/src/main.rs b/src/main.rs index 121ac2c..683ffa1 100644 --- a/src/main.rs +++ b/src/main.rs @@ -47,15 +47,6 @@ async fn main() { info!("Using config file: {}", arguments.config_file); - let database_url = format!( - "postgresql://{}:{}@{}:{}/{}", - config.database.user, - config.database.password, - config.database.host, - config.database.port, - config.database.table - ); - let options = PgConnectOptions::new() .username(&config.database.user) .password(&config.database.password) diff --git a/src/scanner.rs b/src/scanner.rs index 6600ad7..24ee58d 100644 --- a/src/scanner.rs +++ b/src/scanner.rs @@ -1,5 +1,6 @@ use crate::config::Config; use crate::response::Server; +use crate::utils::RunError; use crate::{database, ping}; use futures_util::future; use indicatif::{ProgressBar, ProgressStyle}; @@ -9,9 +10,9 @@ use std::io::{BufRead, BufReader}; use std::sync::Arc; use std::time::Duration; use tokio::process::Command; -use tokio::sync::Semaphore; +use tokio::sync::{Mutex, Semaphore}; use tracing::log::warn; -use tracing::{error, info}; +use tracing::{debug, error, info}; static PERMITS: Semaphore = Semaphore::const_new(1000); const TIMEOUT_SECS: Duration = Duration::from_secs(5); @@ -84,7 +85,8 @@ impl Scanner { } } - pub async fn scan_servers_from_vec(servers: Vec<(String, u16)>) { + /// Takes a vector of pingable servers + pub async fn scan_servers_from_vec(servers: Vec<(String, u16)>) -> Vec { let style = ProgressStyle::with_template( "[{elapsed}] [{bar:40.white/blue}] {pos:>7}/{len:7} ETA {eta}", ) @@ -95,18 +97,19 @@ impl Scanner { let handles = servers .into_iter() - .map(|s| tokio::task::spawn(Self::run(s.0, s.1))) + .map(|s| Scanner::run(s.0, s.1, bar.clone())) .collect::>(); // Wait for all tasks to finish let results = future::join_all(handles).await; - let completed_servers = results - .into_iter() - .flatten() - .map(|s| s.ok()) - .collect::>(); bar.finish_and_clear(); + + results + .into_iter() + // TODO! Don't flatten here, errors are needed later + .flatten() + .collect::>() } /// Rescan servers already found in the database @@ -129,10 +132,10 @@ impl Scanner { let servers = database::fetch_servers(&self.pool) .await .into_iter() - .map(|row| (row.get::(0), row.get::(1))) + .map(|row| (row.get::(0), row.get::(1) as u16)) .collect::>(); - // Self::scan_servers_from_iterator(servers.into_iter()).await; + Self::scan_servers_from_vec(servers).await; // Quit if only one scan is requested in config if !self.config.scanner.repeat { @@ -162,7 +165,6 @@ impl Scanner { loop { let output_file = std::fs::File::open(masscan_output).unwrap(); - // let output_file = std::fs::File::create(masscan_output).unwrap(); let _ = Command::new("sudo") .args(["masscan", "-c", masscan_config]) @@ -174,36 +176,31 @@ impl Scanner { info!("Masscan has completed!"); let reader = BufReader::new(output_file); - let mut handles = Vec::new(); + let mut servers = Vec::new(); for line in reader.lines() { - let line = if let Ok(line) = line { - line - } else { - continue; + let line = match line { + Ok(line) => line, + Err(_) => continue, }; - let mut full_line = line.split_whitespace(); + let mut line = line.split_whitespace(); - let port = if let Some(port) = full_line.nth(2) { - if let Ok(port) = port.parse::() { - port - } else { - continue; - } - } else { - continue; + let port = match line.nth(2).and_then(|s| s.parse::().ok()) { + Some(port) => port, + None => continue, }; - let address = if let Some(address) = full_line.next() { - address.to_owned() - } else { - continue; + let address = match line.next() { + Some(address) => address.to_owned(), + None => continue, }; - handles.push(tokio::task::spawn(Self::run(address, port))); + servers.push((address, port)); } + Self::scan_servers_from_vec(servers).await; + // Quit if only one scan is requested in config if !self.config.scanner.repeat { info!("Exiting"); @@ -221,21 +218,79 @@ impl Scanner { } } - async fn run(address: String, port: u16) -> anyhow::Result { - let host = (&*address, port); + async fn run(address: String, port: u16, bar: Arc) -> anyhow::Result { + async fn inner(address: String, port: u16) -> anyhow::Result { + // Ping server + let permit = PERMITS.acquire().await?; + let pinged_server = tokio::time::timeout( + TIMEOUT_SECS, + // TODO! Fix + ping::ping_server((&*address.split('/').nth(0).unwrap(), port)), + ) + .await??; + drop(permit); - // Ping server - let permit = PERMITS.acquire().await?; - let pinged_server = tokio::time::timeout(TIMEOUT_SECS, ping::ping_server(host)).await??; - drop(permit); + // Parse response + let mut server = serde_json::from_str::(&pinged_server)?; + server.address = address; + server.port = port; - println!("{pinged_server}"); + Ok(server) + } - // Parse response - let mut server: Server = serde_json::from_str(&pinged_server)?; - server.address = address; - server.port = port; + let server = inner(address, port).await; + bar.inc(1); + server + } - Ok(server) + pub async fn complete_scan(&self, results: Vec>) { + let results_len = results.len(); + debug!("results_len = {}", results_len); + + let (servers, errors): (Vec<_>, Vec<_>) = results.into_iter().partition(Result::is_ok); + + let errors_len = errors.len(); + + // Print scan errors + if !errors.is_empty() { + warn!("Scan returned {} total errors!", errors_len); + let mut counts = [0u32; 6]; + for e in errors.into_iter().filter_map(Result::err) { + let i: usize = e.into(); + counts[i] += 1; + } + + warn!("{} errors while parsing addresses", counts[0]); + warn!("{} I/0 Errors", counts[1]); + warn!("{} malformed responses", counts[2]); + warn!("{} errors while parsing responses", counts[3]); + warn!("{} servers timed out", counts[4]); + info!("{} servers removed due to opting out", counts[5]) + } + + // Transactions allow adding multiple statements to a single query + let transaction = Arc::new(Mutex::new( + self.pool + .begin() + .await + .expect("failed to create transaction"), + )); + + let completed_servers = servers + .into_iter() + .filter_map(Result::ok) + .collect::>(); + + info!( + "Commiting {} servers to database...", + completed_servers.len() + ); + + Arc::try_unwrap(transaction) + .unwrap() + .into_inner() + .commit() + .await + .expect("error while commiting to database"); } } diff --git a/src/utils.rs b/src/utils.rs index 4c27896..e589020 100644 --- a/src/utils.rs +++ b/src/utils.rs @@ -1,66 +1,5 @@ -use crate::database; use crate::response::Server; -use indicatif::{ProgressBar, ProgressIterator, ProgressStyle}; -use sqlx::{Pool, Postgres}; -use std::sync::Arc; use thiserror::Error; -use tokio::sync::Mutex; -use tokio::task::JoinError; -use tracing::{debug, info, warn}; - -pub async fn complete_scan( - results: Vec, JoinError>>, - pool: &Pool, -) { - let results_len = results.len(); - debug!("results_len = {}", results_len); - - let (servers, errors): (Vec<_>, Vec<_>) = results - .into_iter() - .filter_map(Result::ok) - .partition(Result::is_ok); - - let errors_len = errors.len(); - - // Print scan errors - if !errors.is_empty() { - warn!("Scan returned {} total errors!", errors_len); - let mut counts = [0u32; 6]; - for e in errors.into_iter().filter_map(Result::err) { - let i: usize = e.into(); - counts[i] += 1; - } - - warn!("{} errors while parsing addresses", counts[0]); - warn!("{} I/0 Errors", counts[1]); - warn!("{} malformed responses", counts[2]); - warn!("{} errors while parsing responses", counts[3]); - warn!("{} servers timed out", counts[4]); - info!("{} servers removed due to opting out", counts[5]) - } - - // Transactions allow adding multiple statements to a single query - let transaction = Arc::new(Mutex::new( - pool.begin().await.expect("failed to create transaction"), - )); - - let completed_servers = servers - .into_iter() - .filter_map(Result::ok) - .collect::>(); - - info!( - "Commiting {} servers to database...", - completed_servers.len() - ); - - Arc::try_unwrap(transaction) - .unwrap() - .into_inner() - .commit() - .await - .expect("error while commiting to database"); -} #[derive(Debug, Error)] pub enum RunError {