From a58de4cc41d7133c8a9df225e9342e914a9f3e26 Mon Sep 17 00:00:00 2001
From: funtimes909 <funtimes90969@proton.me>
Date: Sat, 26 Apr 2025 01:01:23 +1200
Subject: [PATCH] Implement requested changes

---
 Cargo.lock     |  7 +++++++
 Cargo.toml     |  1 +
 src/main.rs    | 11 +++++++----
 src/masscan.rs | 50 ++++++++++++++++++++++++--------------------------
 4 files changed, 39 insertions(+), 30 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index db86202..c491b8d 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -6,6 +6,7 @@ version = 4
 name = "ServerSeekerV2-Rust"
 version = "0.1.0"
 dependencies = [
+ "anyhow",
  "futures-core",
  "futures-util",
  "indicatif",
@@ -40,6 +41,12 @@ version = "0.2.21"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "683d7910e743518b0e34f1186f92494becacb047c7b6bf616c96772180fef923"
 
+[[package]]
+name = "anyhow"
+version = "1.0.98"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "e16d2d3311acee920a9eb8d33b8cbc1787ce4a264e85f964c2404b969bdcd487"
+
 [[package]]
 name = "atoi"
 version = "2.0.0"
diff --git a/Cargo.toml b/Cargo.toml
index c615260..2e111eb 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -17,6 +17,7 @@ sqlx = { version = "0.8", features = ["postgres", "runtime-tokio"] }
 indicatif = { version = "0.17" }
 futures-core = "0.3"
 futures-util = "0.3"
+anyhow = "1.0.98"
 
 [profile.release]
 strip = true
diff --git a/src/main.rs b/src/main.rs
index e194a20..fd19162 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -45,11 +45,14 @@ async fn main() {
 		}
 	};
 
-	let style = ProgressStyle::with_template("[{elapsed}] [{bar:40.white/blue}] {pos:>7}/{len:7}")
-		.unwrap()
-		.progress_chars("=>-");
+	let style = ProgressStyle::with_template(
+		"[{elapsed}] [{bar:40.white/blue}] {pos:>7}/{len:7} ETA {eta}",
+	)
+	.unwrap()
+	.progress_chars("=>-");
 
-	let mode = Mode::Discovery;
+	// TODO: CLI arguments for this
+	let mode = Mode::Rescan;
 
 	match mode {
 		Mode::Discovery => masscan::start(pool, config, style).await,
diff --git a/src/masscan.rs b/src/masscan.rs
index a30c869..3f440ab 100644
--- a/src/masscan.rs
+++ b/src/masscan.rs
@@ -22,20 +22,16 @@ pub struct Port {
 	port: u16,
 }
 
-pub fn start_masscan(masscan_config: String) {
+pub fn start_masscan(masscan_config: &str) {
 	info!("Starting masscan...");
 	if cfg!(target_os = "windows") {
 		Command::new("cmd.exe")
-			.arg("/c")
-			.arg(format!("masscan -c {masscan_config}"))
+			.args(["masscan", "-c", masscan_config])
 			.spawn()
 			.expect("failed to execute process")
 	} else {
-		Command::new("/bin/sh")
-			.arg("-c")
-			// Masscan needs to be run as root, nothing I can do about it
-			// TODO! Change this to be a toggle in the config for security
-			.arg(format!("sudo masscan -c {masscan_config}"))
+		Command::new("sudo")
+			.args(["masscan", "-c", masscan_config])
 			.spawn()
 			.expect("failed to execute process")
 	}
@@ -45,29 +41,24 @@ pub fn start_masscan(masscan_config: String) {
 	info!("Masscan completed");
 }
 
-pub fn parse_output(masscan_output: String) -> Vec<(String, u16)> {
-	let file = std::fs::read_to_string(&masscan_output).unwrap_or_else(|_| {
-		error!("Masscan output not found");
-		std::process::exit(1)
-	});
+pub fn parse_output(masscan_output: &str) -> anyhow::Result<Vec<(String, u16)>> {
+	let file = std::fs::read_to_string(&masscan_output)?;
+	let output = serde_json::from_str::<Vec<Masscan>>(&file)?;
 
-	let output = serde_json::from_str::<Vec<Masscan>>(&file).unwrap_or_else(|_| {
-		error!("Failed to read masscan output");
-		std::process::exit(1)
-	});
-
-	output
+	Ok(output
 		.into_iter()
 		.flat_map(|entry| {
-			let ip = entry.ip;
-			entry.ports.into_iter().map(move |p| (ip.clone(), p.port))
+			entry
+				.ports
+				.into_iter()
+				.map(move |p| (entry.ip.clone(), p.port))
 		})
-		.collect()
+		.collect())
 }
 
 pub async fn start(pool: Pool<Postgres>, config: Config, style: ProgressStyle) {
-	let masscan_config = config.masscan.config_file;
-	let masscan_output = config.masscan.output_file;
+	let masscan_config = config.masscan.config_file.as_str();
+	let masscan_output = config.masscan.output_file.as_str();
 
 	let transaction = Arc::new(Mutex::new(
 		pool.begin().await.expect("failed to create transaction"),
@@ -79,8 +70,15 @@ pub async fn start(pool: Pool<Postgres>, config: Config, style: ProgressStyle) {
 
 	loop {
 		// TODO: Config changes needed here
-		start_masscan(masscan_config.clone());
-		let servers = parse_output(masscan_output.clone());
+		start_masscan(masscan_config);
+		let servers = match parse_output(masscan_output) {
+			Ok(servers) => servers,
+			Err(e) => {
+				error!("Failed to parse masscan output: {e:?}");
+				std::process::exit(1);
+			}
+		};
+
 		info!("{} servers found", servers.len());
 
 		let progress_bar =