This commit is contained in:
funtimes909
2025-04-29 12:11:09 +12:00
parent 5a48dbb66d
commit 8d663d35ef
7 changed files with 71 additions and 63 deletions
Generated
+26 -19
View File
@@ -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",
+1 -1
View File
@@ -6,7 +6,7 @@ use std::time::{SystemTime, UNIX_EPOCH};
use tracing::{debug, info};
pub async fn fetch_servers(pool: &Pool<Postgres>) -> BoxStream<Result<PgRow, Error>> {
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<Postgres>) -> i64 {
-4
View File
@@ -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();
-1
View File
@@ -69,7 +69,6 @@ pub async fn start(pool: Pool<Postgres>, config: Config, style: ProgressStyle) {
}
loop {
// TODO: Config changes needed here
start_masscan(masscan_config);
let servers = match parse_output(masscan_output) {
Ok(servers) => servers,
+4 -4
View File
@@ -26,19 +26,19 @@ pub enum PingServerError {
MalformedResponse,
}
pub async fn ping_server(host: &(String, u16)) -> Result<String, PingServerError> {
pub async fn ping_server((address, port): &(String, u16)) -> Result<String, PingServerError> {
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
+14 -32
View File
@@ -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<Postgres>, config: Config, style: ProgressStyle) {
@@ -34,7 +34,7 @@ pub async fn rescan_servers(pool: Pool<Postgres>, 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<Postgres>, 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::<String, _>(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::<Vec<_>>();
// Print information about scan
scan_results(results);
@@ -88,25 +92,3 @@ pub async fn rescan_servers(pool: Pool<Postgres>, 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<usize> for RunError {
fn into(self) -> usize {
match self {
Self::PingServer(_) => 0,
Self::ParseResponse(_) => 1,
Self::DatabaseUpdate(_) => 2,
Self::TimedOut(_) => 3,
}
}
}
+26 -2
View File
@@ -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<Result<(), RunError>>) {
let results_len = results.len();
@@ -37,6 +39,28 @@ pub fn scan_results(results: Vec<Result<(), RunError>>) {
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<usize> 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<Mutex<PgTransaction<'_>>>,
@@ -47,7 +71,7 @@ pub async fn run(
transaction: Arc<Mutex<PgTransaction<'_>>>,
) -> 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)?;