diff --git a/Cargo.lock b/Cargo.lock index c491b8d..8c8a9ba 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -429,9 +429,9 @@ dependencies = [ [[package]] name = "getrandom" -version = "0.2.15" +version = "0.2.16" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c4567c8db10ae91089c99af84c68c38da3ec2f087c3f82960bcdbf3656b6f4d7" +checksum = "335ff9f135e4384c8150d6f27c6daed433577f86b4750418338c01a1a2528592" dependencies = [ "cfg-if", "libc", @@ -710,9 +710,9 @@ checksum = "d750af042f7ef4f724306de029d18836c26c1765a54a6a3f094cbd23a7267ffa" [[package]] name = "libm" -version = "0.2.11" +version = "0.2.13" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8355be11b20d696c8f18f6cc018c4e372165b1fa8126cef092399c9951984ffa" +checksum = "c9627da5196e5d8ed0b0495e61e518847578da83483c37288316d9b2e03a7f72" [[package]] name = "libsqlite3-sys" @@ -1021,7 +1021,7 @@ version = "0.6.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ec0be4795e2f6a28069bec0b5ff3e2ac9bafc99e6a9a7dc3547996c5c816922c" dependencies = [ - "getrandom 0.2.15", + "getrandom 0.2.16", ] [[package]] @@ -1448,9 +1448,9 @@ checksum = "13c2bddecc57b384dee18652358fb23172facb8a2c51ccc10d74c157bdea3292" [[package]] name = "syn" -version = "2.0.100" +version = "2.0.101" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b09a44accad81e1ba1cd74a32461ba89dee89095ba17b32f5d03683b1b1fc2a0" +checksum = "8ce2b7fc941b3a24138a0a7cf8e858bfc6a992e7978a068a5c760deb0ed43caf" dependencies = [ "proc-macro2", "quote", @@ -1578,9 +1578,9 @@ dependencies = [ [[package]] name = "toml" -version = "0.8.20" +version = "0.8.21" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "cd87a5cdd6ffab733b2f74bc4fd7ee5fff6634124999ac278c35fc78c6120148" +checksum = "900f6c86a685850b1bc9f6223b20125115ee3f31e01207d81655bbcc0aea9231" dependencies = [ "serde", "serde_spanned", @@ -1590,26 +1590,33 @@ dependencies = [ [[package]] name = "toml_datetime" -version = "0.6.8" +version = "0.6.9" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0dd7358ecb8fc2f8d014bf86f6f638ce72ba252a2c3a2572f2a795f1d23efb41" +checksum = "3da5db5a963e24bc68be8b17b6fa82814bb22ee8660f192bb182771d498f09a3" dependencies = [ "serde", ] [[package]] name = "toml_edit" -version = "0.22.24" +version = "0.22.25" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "17b4795ff5edd201c7cd6dca065ae59972ce77d1b80fa0a84d94950ece7d1474" +checksum = "10558ed0bd2a1562e630926a2d1f0b98c827da99fabd3fe20920a59642504485" dependencies = [ "indexmap", "serde", "serde_spanned", "toml_datetime", + "toml_write", "winnow", ] +[[package]] +name = "toml_write" +version = "0.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "28391a4201ba7eb1984cfeb6862c0b3ea2cfe23332298967c749dddc0d6cd976" + [[package]] name = "tracing" version = "0.1.41" @@ -2018,9 +2025,9 @@ checksum = "589f6da84c646204747d1270a2a5661ea66ed1cced2631d546fdfb155959f9ec" [[package]] name = "winnow" -version = "0.7.6" +version = "0.7.7" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "63d3fcd9bba44b03821e7d699eeee959f3126dcc4aa8e4ae18ec617c2a5cea10" +checksum = "6cb8234a863ea0e8cd7284fcdd4f145233eb00fee02bbdd9861aec44e6477bc5" dependencies = [ "memchr", ] @@ -2072,18 +2079,18 @@ dependencies = [ [[package]] name = "zerocopy" -version = "0.8.24" +version = "0.8.25" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2586fea28e186957ef732a5f8b3be2da217d65c5969d4b1e17f973ebbe876879" +checksum = "a1702d9583232ddb9174e01bb7c15a2ab8fb1bc6f227aa1233858c351a3ba0cb" dependencies = [ "zerocopy-derive", ] [[package]] name = "zerocopy-derive" -version = "0.8.24" +version = "0.8.25" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a996a8f63c5c4448cd959ac1bab0aaa3306ccfd060472f85943ee0750f0169be" +checksum = "28a6e20d751156648aa063f3800b706ee209a32c0b4d9f24be3d980b01be55ef" dependencies = [ "proc-macro2", "quote", diff --git a/src/database.rs b/src/database.rs index 8673729..0a33b9c 100644 --- a/src/database.rs +++ b/src/database.rs @@ -6,7 +6,7 @@ use std::time::{SystemTime, UNIX_EPOCH}; use tracing::{debug, info}; pub async fn fetch_servers(pool: &Pool) -> BoxStream> { - sqlx::query("SELECT address FROM servers ORDER BY lastseen ASC").fetch(pool) + sqlx::query("SELECT address FROM servers ORDER BY lastseen DESC").fetch(pool) } pub async fn fetch_count(pool: &Pool) -> i64 { diff --git a/src/main.rs b/src/main.rs index fd19162..0cbb923 100644 --- a/src/main.rs +++ b/src/main.rs @@ -9,7 +9,6 @@ mod utils; use config::load_config; use indicatif::ProgressStyle; use sqlx::PgPool; -use std::time::Duration; use tracing::error; enum Mode { @@ -17,9 +16,6 @@ enum Mode { Rescan, } -// TODO: add to config file -const TIMEOUT_SECS: Duration = Duration::from_secs(5); - #[tokio::main] async fn main() { tracing_subscriber::fmt::init(); diff --git a/src/masscan.rs b/src/masscan.rs index 3f440ab..9115254 100644 --- a/src/masscan.rs +++ b/src/masscan.rs @@ -69,7 +69,6 @@ pub async fn start(pool: Pool, config: Config, style: ProgressStyle) { } loop { - // TODO: Config changes needed here start_masscan(masscan_config); let servers = match parse_output(masscan_output) { Ok(servers) => servers, diff --git a/src/ping.rs b/src/ping.rs index 1001e31..9bc4a79 100644 --- a/src/ping.rs +++ b/src/ping.rs @@ -26,19 +26,19 @@ pub enum PingServerError { MalformedResponse, } -pub async fn ping_server(host: &(String, u16)) -> Result { +pub async fn ping_server((address, port): &(String, u16)) -> Result { let socket = SocketAddr::V4(SocketAddrV4::new( - Ipv4Addr::from_str(host.0.as_str())?, - host.1, + Ipv4Addr::from_str(address.as_str())?, + *port, )); + // TODO: Rewrite ALL of this below // Connect and create buffer let mut stream = TcpStream::connect(&socket).await?; let mut buffer = [0; 1024]; // Send payload stream.write_all(&PAYLOAD).await?; - // TODO: figure out why a good half of servers are timing out here let total_read = stream.read(&mut buffer).await?; // Decode diff --git a/src/scan.rs b/src/scan.rs index d9fb230..7415965 100644 --- a/src/scan.rs +++ b/src/scan.rs @@ -1,14 +1,14 @@ use crate::config::Config; use crate::utils::scan_results; -use crate::{database, ping, utils}; +use crate::{database, utils}; +use futures_util::future::join_all; use futures_util::TryStreamExt; use indicatif::{ProgressBar, ProgressStyle}; use sqlx::{Pool, Postgres, Row}; use std::sync::Arc; use std::time::{Duration, SystemTime, UNIX_EPOCH}; -use thiserror::Error; use tokio::sync::Mutex; -use tokio::task::JoinSet; +use tokio::task; use tracing::{info, warn}; pub async fn rescan_servers(pool: Pool, config: Config, style: ProgressStyle) { @@ -34,7 +34,7 @@ pub async fn rescan_servers(pool: Pool, config: Config, style: Progres )); let progress_bar = Arc::new(ProgressBar::new(length).with_style(style.clone())); - let mut ping_set = JoinSet::new(); + let mut handles = Vec::new(); let scan_start = SystemTime::now() .duration_since(UNIX_EPOCH) .expect("system time is before the unix epoch") @@ -42,18 +42,22 @@ pub async fn rescan_servers(pool: Pool, config: Config, style: Progres // Streams results from the database while let Some(row) = servers.try_next().await.unwrap() { - let address: String = row.get(0); + let address = row.get::(0); - for port in port_start..=port_end { - ping_set.spawn(utils::run( + (port_start..=port_end).into_iter().for_each(|port| { + handles.push(task::spawn(utils::run( (address.to_owned(), port), transaction.clone(), progress_bar.clone(), - )); - } + ))); + }); } - let results = ping_set.join_all().await; + let results = join_all(handles) + .await + .into_iter() + .filter_map(|r| r.ok()) + .collect::>(); // Print information about scan scan_results(results); @@ -88,25 +92,3 @@ pub async fn rescan_servers(pool: Pool, config: Config, style: Progres } } } - -#[derive(Debug, Error)] -pub enum RunError { - #[error("Error while pinging server")] - PingServer(#[from] ping::PingServerError), - #[error("Error while parsing response")] - ParseResponse(#[from] serde_json::Error), - #[error("Error while updating database")] - DatabaseUpdate(#[from] sqlx::Error), - #[error("Connection timed out")] - TimedOut(#[from] tokio::time::error::Elapsed), -} -impl Into for RunError { - fn into(self) -> usize { - match self { - Self::PingServer(_) => 0, - Self::ParseResponse(_) => 1, - Self::DatabaseUpdate(_) => 2, - Self::TimedOut(_) => 3, - } - } -} diff --git a/src/utils.rs b/src/utils.rs index 7515252..2d0840e 100644 --- a/src/utils.rs +++ b/src/utils.rs @@ -1,12 +1,14 @@ -use crate::scan::RunError; use crate::{database, ping, response}; use indicatif::ProgressBar; use sqlx::PgTransaction; use std::sync::Arc; +use std::time::Duration; +use thiserror::Error; use tokio::sync::{Mutex, Semaphore}; use tracing::{info, warn}; static PERMITS: Semaphore = Semaphore::const_new(2000); +const TIMEOUT_SECS: Duration = Duration::from_secs(5); pub fn scan_results(results: Vec>) { let results_len = results.len(); @@ -37,6 +39,28 @@ pub fn scan_results(results: Vec>) { info!("Commiting {} results to database", results_len - errors_len); } +#[derive(Debug, Error)] +pub enum RunError { + #[error("Error while pinging server")] + PingServer(#[from] ping::PingServerError), + #[error("Error while parsing response")] + ParseResponse(#[from] serde_json::Error), + #[error("Error while updating database")] + DatabaseUpdate(#[from] sqlx::Error), + #[error("Connection timed out")] + TimedOut(#[from] tokio::time::error::Elapsed), +} +impl Into for RunError { + fn into(self) -> usize { + match self { + Self::PingServer(_) => 0, + Self::ParseResponse(_) => 1, + Self::DatabaseUpdate(_) => 2, + Self::TimedOut(_) => 3, + } + } +} + pub async fn run( host: (String, u16), transaction: Arc>>, @@ -47,7 +71,7 @@ pub async fn run( transaction: Arc>>, ) -> Result<(), RunError> { let permit = PERMITS.acquire().await.unwrap(); - let results = tokio::time::timeout(crate::TIMEOUT_SECS, ping::ping_server(&host)).await??; + let results = tokio::time::timeout(TIMEOUT_SECS, ping::ping_server(&host)).await??; drop(permit); let response = response::parse_response(results)?;