Added scaffoded project
This commit is contained in:
+352
@@ -0,0 +1,352 @@
|
||||
use anyhow::{anyhow, Context, Result};
|
||||
use std::{
|
||||
env, fs,
|
||||
net::SocketAddr,
|
||||
path::{Path, PathBuf},
|
||||
};
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum Mode {
|
||||
Local,
|
||||
Coordinator,
|
||||
Worker,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct Config {
|
||||
pub mode: Mode,
|
||||
pub cidrs: Vec<String>,
|
||||
pub port: u16,
|
||||
pub timeout_ms: u64,
|
||||
pub concurrency: usize,
|
||||
pub chunk_size: usize,
|
||||
pub database_url: String,
|
||||
pub coordinator_bind: SocketAddr,
|
||||
pub coordinator_addr: SocketAddr,
|
||||
pub worker_name: String,
|
||||
}
|
||||
|
||||
impl Config {
|
||||
pub fn from_env_and_args() -> Result<Self> {
|
||||
let mut args = env::args().skip(1).collect::<Vec<_>>();
|
||||
let mode = parse_mode(&mut args)?;
|
||||
let cli = CliArgs::from_args(&args)?;
|
||||
|
||||
let cidrs = load_cidrs(cli.cidr_file.as_deref(), cli.cidrs.as_deref())?;
|
||||
Ok(Self {
|
||||
mode,
|
||||
cidrs,
|
||||
port: cli
|
||||
.port
|
||||
.or_else(|| parse_u16("SCAN_PORT", 554).ok())
|
||||
.unwrap_or(554),
|
||||
timeout_ms: cli
|
||||
.timeout_ms
|
||||
.or_else(|| parse_u64("SCAN_TIMEOUT_MS", 1200).ok())
|
||||
.unwrap_or(1200),
|
||||
concurrency: cli
|
||||
.concurrency
|
||||
.or_else(|| parse_usize("SCAN_CONCURRENCY", default_concurrency()).ok())
|
||||
.unwrap_or_else(default_concurrency),
|
||||
chunk_size: cli
|
||||
.chunk_size
|
||||
.or_else(|| parse_usize("SCAN_CHUNK_SIZE", default_chunk_size()).ok())
|
||||
.unwrap_or_else(default_chunk_size),
|
||||
database_url: cli.database_url.unwrap_or_else(database_url_from_env),
|
||||
coordinator_bind: cli
|
||||
.coordinator_bind
|
||||
.or_else(|| parse_socket_addr("COORDINATOR_BIND", "0.0.0.0:666").ok())
|
||||
.unwrap_or_else(|| "0.0.0.0:666".parse().expect("valid default bind address")),
|
||||
coordinator_addr: cli
|
||||
.coordinator_addr
|
||||
.or_else(|| parse_socket_addr("COORDINATOR_ADDR", "127.0.0.1:666").ok())
|
||||
.unwrap_or_else(|| {
|
||||
"127.0.0.1:666"
|
||||
.parse()
|
||||
.expect("valid default coordinator address")
|
||||
}),
|
||||
worker_name: cli.worker_name.unwrap_or_else(default_worker_name),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Default)]
|
||||
struct CliArgs {
|
||||
cidrs: Option<String>,
|
||||
cidr_file: Option<PathBuf>,
|
||||
database_url: Option<String>,
|
||||
coordinator_bind: Option<SocketAddr>,
|
||||
coordinator_addr: Option<SocketAddr>,
|
||||
worker_name: Option<String>,
|
||||
port: Option<u16>,
|
||||
timeout_ms: Option<u64>,
|
||||
concurrency: Option<usize>,
|
||||
chunk_size: Option<usize>,
|
||||
}
|
||||
|
||||
impl CliArgs {
|
||||
fn from_args(args: &[String]) -> Result<Self> {
|
||||
let mut cli = Self::default();
|
||||
let mut index = 0;
|
||||
|
||||
while index < args.len() {
|
||||
let arg = &args[index];
|
||||
match arg.as_str() {
|
||||
"--cidrs" => {
|
||||
cli.cidrs = Some(next_value(args, &mut index, "--cidrs")?);
|
||||
}
|
||||
"--cidr-file" => {
|
||||
cli.cidr_file =
|
||||
Some(PathBuf::from(next_value(args, &mut index, "--cidr-file")?));
|
||||
}
|
||||
"--database" | "--database-url" => {
|
||||
cli.database_url = Some(next_value(args, &mut index, arg)?);
|
||||
}
|
||||
"--bind" => {
|
||||
cli.coordinator_bind = Some(parse_socket_addr_literal(&next_value(
|
||||
args, &mut index, "--bind",
|
||||
)?)?);
|
||||
}
|
||||
"--connect" => {
|
||||
cli.coordinator_addr = Some(parse_socket_addr_literal(&next_value(
|
||||
args,
|
||||
&mut index,
|
||||
"--connect",
|
||||
)?)?);
|
||||
}
|
||||
"--name" => {
|
||||
cli.worker_name = Some(next_value(args, &mut index, "--name")?);
|
||||
}
|
||||
"--port" => {
|
||||
cli.port = Some(parse_u16_literal(&next_value(args, &mut index, "--port")?)?);
|
||||
}
|
||||
"--timeout-ms" => {
|
||||
cli.timeout_ms = Some(parse_u64_literal(&next_value(
|
||||
args,
|
||||
&mut index,
|
||||
"--timeout-ms",
|
||||
)?)?);
|
||||
}
|
||||
"--concurrency" => {
|
||||
cli.concurrency = Some(parse_usize_literal(&next_value(
|
||||
args,
|
||||
&mut index,
|
||||
"--concurrency",
|
||||
)?)?);
|
||||
}
|
||||
"--chunk-size" => {
|
||||
cli.chunk_size = Some(parse_usize_literal(&next_value(
|
||||
args,
|
||||
&mut index,
|
||||
"--chunk-size",
|
||||
)?)?);
|
||||
}
|
||||
_ if arg.starts_with("--") => {
|
||||
return Err(anyhow!("unknown argument: {arg}"));
|
||||
}
|
||||
_ => {
|
||||
return Err(anyhow!("unexpected positional argument: {arg}"));
|
||||
}
|
||||
}
|
||||
|
||||
index += 1;
|
||||
}
|
||||
|
||||
Ok(cli)
|
||||
}
|
||||
}
|
||||
|
||||
fn parse_mode(args: &mut Vec<String>) -> Result<Mode> {
|
||||
if let Some(first) = args.first() {
|
||||
if !first.starts_with('-') {
|
||||
let mode = parse_mode_literal(first)?;
|
||||
args.remove(0);
|
||||
return Ok(mode);
|
||||
}
|
||||
}
|
||||
|
||||
match env::var("SCAN_MODE") {
|
||||
Ok(value) => parse_mode_literal(value.trim()),
|
||||
Err(_) => Ok(Mode::Local),
|
||||
}
|
||||
}
|
||||
|
||||
fn parse_mode_literal(value: &str) -> Result<Mode> {
|
||||
match value.trim().to_ascii_lowercase().as_str() {
|
||||
"local" | "scan" => Ok(Mode::Local),
|
||||
"coordinator" | "coord" | "server" => Ok(Mode::Coordinator),
|
||||
"worker" => Ok(Mode::Worker),
|
||||
other => Err(anyhow!(
|
||||
"modo inválido: {other}. Use local, coordinator ou worker"
|
||||
)),
|
||||
}
|
||||
}
|
||||
|
||||
fn load_cidrs(file: Option<&Path>, inline: Option<&str>) -> Result<Vec<String>> {
|
||||
if let Some(inline) = inline {
|
||||
let cidrs = parse_cidr_list(inline);
|
||||
if !cidrs.is_empty() {
|
||||
return Ok(cidrs);
|
||||
}
|
||||
}
|
||||
|
||||
if let Some(file) = file {
|
||||
let content = fs::read_to_string(file)
|
||||
.with_context(|| format!("failed to read CIDR file {}", file.display()))?;
|
||||
let cidrs = parse_cidr_list(&content);
|
||||
if !cidrs.is_empty() {
|
||||
return Ok(cidrs);
|
||||
}
|
||||
return Err(anyhow!("arquivo de CIDRs vazio: {}", file.display()));
|
||||
}
|
||||
|
||||
if let Ok(value) = env::var("SCAN_CIDRS") {
|
||||
let value = value.trim();
|
||||
if !value.is_empty() {
|
||||
let cidrs = parse_cidr_list(value);
|
||||
if !cidrs.is_empty() {
|
||||
return Ok(cidrs);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if let Ok(path) = env::var("SCAN_CIDR_FILE") {
|
||||
let path = path.trim();
|
||||
if !path.is_empty() {
|
||||
let content = fs::read_to_string(path)
|
||||
.with_context(|| format!("failed to read CIDR file {path}"))?;
|
||||
let cidrs = parse_cidr_list(&content);
|
||||
if !cidrs.is_empty() {
|
||||
return Ok(cidrs);
|
||||
}
|
||||
return Err(anyhow!("arquivo de CIDRs vazio: {path}"));
|
||||
}
|
||||
}
|
||||
|
||||
if let Ok(value) = env::var("SCAN_CIDR") {
|
||||
let value = value.trim();
|
||||
if !value.is_empty() {
|
||||
let cidrs = parse_cidr_list(value);
|
||||
if !cidrs.is_empty() {
|
||||
return Ok(cidrs);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Ok(Vec::new())
|
||||
}
|
||||
|
||||
fn parse_cidr_list(input: &str) -> Vec<String> {
|
||||
input
|
||||
.lines()
|
||||
.flat_map(|line| line.split('#').next())
|
||||
.flat_map(|line| line.split(|c: char| c == ',' || c.is_whitespace()))
|
||||
.map(str::trim)
|
||||
.filter(|entry| !entry.is_empty())
|
||||
.map(ToString::to_string)
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn database_url_from_env() -> String {
|
||||
let value = read_env("DATABASE_URL", "");
|
||||
if value.is_empty() {
|
||||
"postgres://postgres:postgres@127.0.0.1:5432/camfinder".to_string()
|
||||
} else {
|
||||
value
|
||||
}
|
||||
}
|
||||
|
||||
fn default_worker_name() -> String {
|
||||
env::var("WORKER_NAME")
|
||||
.ok()
|
||||
.map(|value| value.trim().to_string())
|
||||
.filter(|value| !value.is_empty())
|
||||
.unwrap_or_else(|| {
|
||||
let host = env::var("HOSTNAME").unwrap_or_else(|_| "worker".to_string());
|
||||
format!("{}-{}", host.trim(), std::process::id())
|
||||
})
|
||||
}
|
||||
|
||||
fn read_env(name: &str, fallback: &str) -> String {
|
||||
env::var(name)
|
||||
.map(|value| value.trim().to_string())
|
||||
.unwrap_or_else(|_| fallback.to_string())
|
||||
}
|
||||
|
||||
fn parse_socket_addr(name: &str, fallback: &str) -> Result<SocketAddr> {
|
||||
parse_socket_addr_literal(&read_env(name, fallback))
|
||||
}
|
||||
|
||||
fn parse_socket_addr_literal(value: &str) -> Result<SocketAddr> {
|
||||
value
|
||||
.trim()
|
||||
.parse()
|
||||
.map_err(|_| anyhow!("invalid socket address: {value}"))
|
||||
}
|
||||
|
||||
fn parse_u16(name: &str, fallback: u16) -> Result<u16> {
|
||||
parse_number(name, fallback)
|
||||
}
|
||||
|
||||
fn parse_u64(name: &str, fallback: u64) -> Result<u64> {
|
||||
parse_number(name, fallback)
|
||||
}
|
||||
|
||||
fn parse_usize(name: &str, fallback: usize) -> Result<usize> {
|
||||
parse_number(name, fallback)
|
||||
}
|
||||
|
||||
fn parse_u16_literal(value: &str) -> Result<u16> {
|
||||
value
|
||||
.trim()
|
||||
.parse()
|
||||
.map_err(|_| anyhow!("invalid u16 value: {value}"))
|
||||
}
|
||||
|
||||
fn parse_u64_literal(value: &str) -> Result<u64> {
|
||||
value
|
||||
.trim()
|
||||
.parse()
|
||||
.map_err(|_| anyhow!("invalid u64 value: {value}"))
|
||||
}
|
||||
|
||||
fn parse_usize_literal(value: &str) -> Result<usize> {
|
||||
value
|
||||
.trim()
|
||||
.parse()
|
||||
.map_err(|_| anyhow!("invalid usize value: {value}"))
|
||||
}
|
||||
|
||||
fn next_value(args: &[String], index: &mut usize, flag: &str) -> Result<String> {
|
||||
let value_index = *index + 1;
|
||||
let value = args
|
||||
.get(value_index)
|
||||
.ok_or_else(|| anyhow!("missing value for {flag}"))?
|
||||
.clone();
|
||||
*index = value_index;
|
||||
Ok(value)
|
||||
}
|
||||
|
||||
fn default_concurrency() -> usize {
|
||||
let cpus = std::thread::available_parallelism()
|
||||
.map(|n| n.get())
|
||||
.unwrap_or(4);
|
||||
(cpus.saturating_mul(2)).clamp(4, 64)
|
||||
}
|
||||
|
||||
fn default_chunk_size() -> usize {
|
||||
default_concurrency().saturating_mul(2).clamp(32, 256)
|
||||
}
|
||||
|
||||
fn parse_number<T>(name: &str, fallback: T) -> Result<T>
|
||||
where
|
||||
T: std::str::FromStr + Copy,
|
||||
{
|
||||
match env::var(name) {
|
||||
Ok(value) => value
|
||||
.trim()
|
||||
.parse::<T>()
|
||||
.map_err(|_| anyhow!("invalid value for {name}: {value}")),
|
||||
Err(_) => Ok(fallback),
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,398 @@
|
||||
use crate::{
|
||||
config::Config,
|
||||
db::Database,
|
||||
models::{ChunkRecord, PendingRange},
|
||||
scan::{build_pending_ranges, load_targets},
|
||||
terminal::{ScreenState, TerminalUi},
|
||||
};
|
||||
use anyhow::{anyhow, Context, Result};
|
||||
use std::{
|
||||
collections::VecDeque,
|
||||
net::Ipv4Addr,
|
||||
sync::{
|
||||
atomic::{AtomicBool, AtomicUsize, Ordering},
|
||||
Arc, Mutex,
|
||||
},
|
||||
time::Instant,
|
||||
};
|
||||
use tokio::{
|
||||
io::{AsyncBufReadExt, AsyncWriteExt, BufReader},
|
||||
net::{TcpListener, TcpStream},
|
||||
task::JoinSet,
|
||||
};
|
||||
|
||||
pub async fn run_coordinator(
|
||||
config: Config,
|
||||
database: Arc<Database>,
|
||||
ui: Arc<TerminalUi>,
|
||||
interrupted: Arc<AtomicBool>,
|
||||
) -> Result<()> {
|
||||
let targets = Arc::new(load_targets(&config.cidrs)?);
|
||||
let target_count = targets.len();
|
||||
let chunk_records = database.load_chunk_records().await?;
|
||||
let (tested_before_run, _open_before_run) = database.summarize_chunks().await?;
|
||||
let pending_ranges = build_pending_ranges(target_count, config.chunk_size, &chunk_records);
|
||||
let pending_count: usize = pending_ranges.iter().map(|range| range.len()).sum();
|
||||
|
||||
if ui.is_enabled() {
|
||||
ui.print_banner();
|
||||
ui.render(&ScreenState {
|
||||
current_ip: if pending_count > 0 {
|
||||
targets[pending_ranges[0].start].to_string()
|
||||
} else {
|
||||
"concluído".to_string()
|
||||
},
|
||||
tested: tested_before_run,
|
||||
total: target_count,
|
||||
});
|
||||
} else {
|
||||
println!(
|
||||
"Coordinator ready on {} for {} hosts across {} CIDRs. Pending {}, already tested {}.",
|
||||
config.coordinator_bind,
|
||||
target_count,
|
||||
config.cidrs.len(),
|
||||
pending_count,
|
||||
tested_before_run
|
||||
);
|
||||
}
|
||||
|
||||
if pending_ranges.is_empty() {
|
||||
if ui.is_enabled() {
|
||||
ui.render(&ScreenState {
|
||||
current_ip: "concluído".to_string(),
|
||||
tested: target_count,
|
||||
total: target_count,
|
||||
});
|
||||
ui.finish();
|
||||
}
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
let listener = TcpListener::bind(config.coordinator_bind)
|
||||
.await
|
||||
.with_context(|| format!("failed to bind {}", config.coordinator_bind))?;
|
||||
let queue = Arc::new(Mutex::new(VecDeque::from(pending_ranges)));
|
||||
let tested_this_run = Arc::new(AtomicUsize::new(0));
|
||||
let open_this_run = Arc::new(AtomicUsize::new(0));
|
||||
let start = Instant::now();
|
||||
let mut handlers = JoinSet::new();
|
||||
|
||||
loop {
|
||||
if interrupted.load(Ordering::SeqCst) {
|
||||
break;
|
||||
}
|
||||
|
||||
tokio::select! {
|
||||
accept_result = listener.accept() => {
|
||||
let (stream, peer) = accept_result?;
|
||||
let queue = Arc::clone(&queue);
|
||||
let targets = Arc::clone(&targets);
|
||||
let database = Arc::clone(&database);
|
||||
let ui = Arc::clone(&ui);
|
||||
let interrupted = Arc::clone(&interrupted);
|
||||
let tested_this_run = Arc::clone(&tested_this_run);
|
||||
let open_this_run = Arc::clone(&open_this_run);
|
||||
let config = config.clone();
|
||||
|
||||
handlers.spawn(async move {
|
||||
handle_worker(
|
||||
stream,
|
||||
peer,
|
||||
queue,
|
||||
targets,
|
||||
database,
|
||||
ui,
|
||||
interrupted,
|
||||
tested_before_run,
|
||||
tested_this_run,
|
||||
open_this_run,
|
||||
config,
|
||||
)
|
||||
.await
|
||||
});
|
||||
}
|
||||
maybe_finished = handlers.join_next(), if !handlers.is_empty() => {
|
||||
if let Some(join_result) = maybe_finished {
|
||||
join_result??;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if queue.lock().expect("queue mutex poisoned").is_empty() && handlers.is_empty() {
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
while let Some(join_result) = handlers.join_next().await {
|
||||
join_result??;
|
||||
}
|
||||
|
||||
let duration = start.elapsed().as_secs_f64();
|
||||
let total_open_count = database.count_open_endpoints(config.port).await?;
|
||||
|
||||
if interrupted.load(Ordering::SeqCst) {
|
||||
if ui.is_enabled() {
|
||||
ui.render(&ScreenState {
|
||||
current_ip: "interrompido".to_string(),
|
||||
tested: tested_before_run + tested_this_run.load(Ordering::SeqCst),
|
||||
total: target_count,
|
||||
});
|
||||
ui.finish();
|
||||
} else {
|
||||
println!("Interrupted after {:.2}s.", duration);
|
||||
println!(
|
||||
"Hosts tested in this run: {}",
|
||||
tested_this_run.load(Ordering::SeqCst)
|
||||
);
|
||||
println!(
|
||||
"Open endpoints found: {}",
|
||||
open_this_run.load(Ordering::SeqCst)
|
||||
);
|
||||
}
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
if ui.is_enabled() {
|
||||
ui.render(&ScreenState {
|
||||
current_ip: "concluído".to_string(),
|
||||
tested: target_count,
|
||||
total: target_count,
|
||||
});
|
||||
ui.finish();
|
||||
} else {
|
||||
println!("Coordinator finished in {:.2}s", duration);
|
||||
println!(
|
||||
"Hosts tested in this run: {}",
|
||||
tested_this_run.load(Ordering::SeqCst)
|
||||
);
|
||||
println!(
|
||||
"Open endpoints found in this run: {}",
|
||||
open_this_run.load(Ordering::SeqCst)
|
||||
);
|
||||
println!("Total open endpoints in DB: {}", total_open_count);
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn handle_worker(
|
||||
stream: TcpStream,
|
||||
peer: std::net::SocketAddr,
|
||||
queue: Arc<Mutex<VecDeque<PendingRange>>>,
|
||||
targets: Arc<Vec<Ipv4Addr>>,
|
||||
database: Arc<Database>,
|
||||
ui: Arc<TerminalUi>,
|
||||
interrupted: Arc<AtomicBool>,
|
||||
tested_before_run: usize,
|
||||
tested_this_run: Arc<AtomicUsize>,
|
||||
open_this_run: Arc<AtomicUsize>,
|
||||
config: Config,
|
||||
) -> Result<()> {
|
||||
let (read_half, mut write_half) = stream.into_split();
|
||||
let mut lines = BufReader::new(read_half).lines();
|
||||
|
||||
let worker_name = match lines.next_line().await? {
|
||||
Some(line) if line.starts_with("HELLO ") => line[6..].trim().to_string(),
|
||||
Some(line) => {
|
||||
return Err(anyhow!("unexpected worker handshake from {peer}: {line}"));
|
||||
}
|
||||
None => return Ok(()),
|
||||
};
|
||||
|
||||
if !worker_name.is_empty() {
|
||||
println!("worker connected from {peer}: {worker_name}");
|
||||
} else {
|
||||
println!("worker connected from {peer}");
|
||||
}
|
||||
|
||||
loop {
|
||||
if interrupted.load(Ordering::SeqCst) {
|
||||
let _ = write_half.write_all(b"SHUTDOWN\n").await;
|
||||
break;
|
||||
}
|
||||
|
||||
let Some(range) = next_range(&queue) else {
|
||||
let _ = write_half.write_all(b"DONE\n").await;
|
||||
break;
|
||||
};
|
||||
|
||||
if let Err(err) = process_range(
|
||||
&mut write_half,
|
||||
&mut lines,
|
||||
range,
|
||||
&targets,
|
||||
&database,
|
||||
&ui,
|
||||
tested_before_run,
|
||||
&tested_this_run,
|
||||
&open_this_run,
|
||||
&config,
|
||||
)
|
||||
.await
|
||||
{
|
||||
requeue_range(&queue, range);
|
||||
return Err(err);
|
||||
}
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn process_range(
|
||||
write_half: &mut tokio::net::tcp::OwnedWriteHalf,
|
||||
lines: &mut tokio::io::Lines<BufReader<tokio::net::tcp::OwnedReadHalf>>,
|
||||
range: PendingRange,
|
||||
targets: &Arc<Vec<Ipv4Addr>>,
|
||||
database: &Arc<Database>,
|
||||
ui: &Arc<TerminalUi>,
|
||||
tested_before_run: usize,
|
||||
tested_this_run: &Arc<AtomicUsize>,
|
||||
open_this_run: &Arc<AtomicUsize>,
|
||||
config: &Config,
|
||||
) -> Result<()> {
|
||||
let ips = targets[range.start..range.end].to_vec();
|
||||
let job_line = format!(
|
||||
"JOB {} {} {} {} {}\n",
|
||||
range.start,
|
||||
range.end,
|
||||
config.port,
|
||||
config.timeout_ms,
|
||||
format_ips(&ips)
|
||||
);
|
||||
write_half.write_all(job_line.as_bytes()).await?;
|
||||
|
||||
let Some(result_line) = lines.next_line().await? else {
|
||||
return Err(anyhow!("worker disconnected before returning job"));
|
||||
};
|
||||
|
||||
let result = parse_result_line(&result_line)?;
|
||||
if result.chunk_start != range.start || result.chunk_end != range.end {
|
||||
return Err(anyhow!(
|
||||
"worker returned mismatched range: expected {}..{}, got {}..{}",
|
||||
range.start,
|
||||
range.end,
|
||||
result.chunk_start,
|
||||
result.chunk_end
|
||||
));
|
||||
}
|
||||
|
||||
for ip in result.open_ips {
|
||||
database.upsert_open_endpoint(ip, config.port).await?;
|
||||
}
|
||||
|
||||
let previous = database.load_chunk_record(range.chunk_start).await?;
|
||||
let tested_before = previous
|
||||
.as_ref()
|
||||
.map(|record| record.tested_count)
|
||||
.unwrap_or(0);
|
||||
let open_before = previous
|
||||
.as_ref()
|
||||
.map(|record| record.open_count)
|
||||
.unwrap_or(0);
|
||||
let tested_count = tested_before + result.tested;
|
||||
let open_count = open_before + result.open_count;
|
||||
database
|
||||
.save_chunk_record(&ChunkRecord {
|
||||
chunk_start: range.chunk_start,
|
||||
chunk_end: range.chunk_end,
|
||||
tested_count,
|
||||
open_count,
|
||||
is_complete: result.tested >= range.len(),
|
||||
})
|
||||
.await?;
|
||||
|
||||
let tested_now = tested_this_run.fetch_add(result.tested, Ordering::SeqCst) + result.tested;
|
||||
open_this_run.fetch_add(result.open_count, Ordering::SeqCst);
|
||||
|
||||
if ui.is_enabled() {
|
||||
let current_ip = ips
|
||||
.last()
|
||||
.map(|ip| ip.to_string())
|
||||
.unwrap_or_else(|| "concluído".to_string());
|
||||
ui.render(&ScreenState {
|
||||
current_ip,
|
||||
tested: tested_before_run + tested_now,
|
||||
total: targets.len(),
|
||||
});
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn requeue_range(queue: &Arc<Mutex<VecDeque<PendingRange>>>, range: PendingRange) {
|
||||
queue
|
||||
.lock()
|
||||
.expect("queue mutex poisoned")
|
||||
.push_front(range);
|
||||
}
|
||||
|
||||
fn next_range(queue: &Arc<Mutex<VecDeque<PendingRange>>>) -> Option<PendingRange> {
|
||||
queue.lock().expect("queue mutex poisoned").pop_front()
|
||||
}
|
||||
|
||||
struct JobResult {
|
||||
chunk_start: usize,
|
||||
chunk_end: usize,
|
||||
tested: usize,
|
||||
open_count: usize,
|
||||
open_ips: Vec<Ipv4Addr>,
|
||||
}
|
||||
|
||||
fn parse_result_line(line: &str) -> Result<JobResult> {
|
||||
let mut parts = line.splitn(6, ' ');
|
||||
let kind = parts.next().unwrap_or("");
|
||||
if kind != "RESULT" {
|
||||
return Err(anyhow!("unexpected worker message: {line}"));
|
||||
}
|
||||
|
||||
let chunk_start = parse_usize_part(parts.next(), "chunk_start", line)?;
|
||||
let chunk_end = parse_usize_part(parts.next(), "chunk_end", line)?;
|
||||
let tested = parse_usize_part(parts.next(), "tested", line)?;
|
||||
let open_count = parse_usize_part(parts.next(), "open_count", line)?;
|
||||
let ips_part = parts.next().unwrap_or("");
|
||||
let open_ips = if ips_part.trim().is_empty() {
|
||||
Vec::new()
|
||||
} else {
|
||||
ips_part
|
||||
.split(',')
|
||||
.map(str::trim)
|
||||
.filter(|entry| !entry.is_empty())
|
||||
.map(|entry| {
|
||||
entry
|
||||
.parse::<Ipv4Addr>()
|
||||
.with_context(|| format!("invalid IP in worker result: {entry}"))
|
||||
})
|
||||
.collect::<Result<Vec<_>>>()?
|
||||
};
|
||||
|
||||
if open_ips.len() != open_count {
|
||||
return Err(anyhow!(
|
||||
"open_count mismatch in worker result: declared {}, received {}",
|
||||
open_count,
|
||||
open_ips.len()
|
||||
));
|
||||
}
|
||||
|
||||
Ok(JobResult {
|
||||
chunk_start,
|
||||
chunk_end,
|
||||
tested,
|
||||
open_count,
|
||||
open_ips,
|
||||
})
|
||||
}
|
||||
|
||||
fn parse_usize_part(part: Option<&str>, name: &str, line: &str) -> Result<usize> {
|
||||
let value = part.ok_or_else(|| anyhow!("missing {name} in worker message: {line}"))?;
|
||||
value
|
||||
.trim()
|
||||
.parse()
|
||||
.map_err(|_| anyhow!("invalid {name} in worker message: {line}"))
|
||||
}
|
||||
|
||||
fn format_ips(ips: &[Ipv4Addr]) -> String {
|
||||
ips.iter()
|
||||
.map(ToString::to_string)
|
||||
.collect::<Vec<_>>()
|
||||
.join(",")
|
||||
}
|
||||
@@ -0,0 +1,218 @@
|
||||
use crate::models::ChunkRecord;
|
||||
use anyhow::{Context, Result};
|
||||
use std::convert::TryFrom;
|
||||
use tokio_postgres::{Client, NoTls};
|
||||
|
||||
pub struct Database {
|
||||
client: Client,
|
||||
}
|
||||
|
||||
impl Database {
|
||||
pub async fn open(database_url: impl AsRef<str>) -> Result<Self> {
|
||||
let database_url = database_url.as_ref().trim();
|
||||
let (client, connection) = tokio_postgres::connect(database_url, NoTls)
|
||||
.await
|
||||
.with_context(|| format!("failed to connect to postgres at {database_url}"))?;
|
||||
|
||||
tokio::spawn(async move {
|
||||
if let Err(err) = connection.await {
|
||||
eprintln!("postgres connection error: {err}");
|
||||
}
|
||||
});
|
||||
|
||||
let db = Self { client };
|
||||
db.init().await?;
|
||||
Ok(db)
|
||||
}
|
||||
|
||||
async fn init(&self) -> Result<()> {
|
||||
self.client
|
||||
.batch_execute(
|
||||
r#"
|
||||
CREATE TABLE IF NOT EXISTS open_rtsp_endpoint (
|
||||
id BIGSERIAL PRIMARY KEY,
|
||||
ip INET NOT NULL UNIQUE,
|
||||
port INTEGER NOT NULL,
|
||||
scanned_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
|
||||
last_seen_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
|
||||
created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
|
||||
updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW()
|
||||
);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS open_rtsp_endpoint_port_idx
|
||||
ON open_rtsp_endpoint (port);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS scan_chunks (
|
||||
chunk_start BIGINT PRIMARY KEY NOT NULL,
|
||||
chunk_end BIGINT NOT NULL,
|
||||
tested_count BIGINT NOT NULL DEFAULT 0,
|
||||
open_count BIGINT NOT NULL DEFAULT 0,
|
||||
is_complete BOOLEAN NOT NULL DEFAULT FALSE,
|
||||
created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
|
||||
updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW()
|
||||
);
|
||||
"#,
|
||||
)
|
||||
.await
|
||||
.context("failed to initialize postgres schema")?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn load_chunk_records(&self) -> Result<Vec<ChunkRecord>> {
|
||||
let rows = self
|
||||
.client
|
||||
.query(
|
||||
r#"
|
||||
SELECT chunk_start, chunk_end, tested_count, open_count, is_complete
|
||||
FROM scan_chunks
|
||||
ORDER BY chunk_start ASC
|
||||
"#,
|
||||
&[],
|
||||
)
|
||||
.await
|
||||
.context("failed to load scan chunks")?;
|
||||
|
||||
rows.into_iter().map(row_to_chunk_record).collect()
|
||||
}
|
||||
|
||||
pub async fn load_chunk_record(&self, chunk_start: usize) -> Result<Option<ChunkRecord>> {
|
||||
let chunk_start = usize_to_i64(chunk_start, "chunk_start")?;
|
||||
let row = self
|
||||
.client
|
||||
.query_opt(
|
||||
r#"
|
||||
SELECT chunk_start, chunk_end, tested_count, open_count, is_complete
|
||||
FROM scan_chunks
|
||||
WHERE chunk_start = $1
|
||||
"#,
|
||||
&[&chunk_start],
|
||||
)
|
||||
.await
|
||||
.context("failed to load scan chunk")?;
|
||||
|
||||
row.map(row_to_chunk_record).transpose()
|
||||
}
|
||||
|
||||
pub async fn save_chunk_record(&self, record: &ChunkRecord) -> Result<()> {
|
||||
let chunk_start = usize_to_i64(record.chunk_start, "chunk_start")?;
|
||||
let chunk_end = usize_to_i64(record.chunk_end, "chunk_end")?;
|
||||
let tested_count = usize_to_i64(record.tested_count, "tested_count")?;
|
||||
let open_count = usize_to_i64(record.open_count, "open_count")?;
|
||||
|
||||
self.client
|
||||
.execute(
|
||||
r#"
|
||||
INSERT INTO scan_chunks (
|
||||
chunk_start,
|
||||
chunk_end,
|
||||
tested_count,
|
||||
open_count,
|
||||
is_complete,
|
||||
created_at,
|
||||
updated_at
|
||||
)
|
||||
VALUES ($1, $2, $3, $4, $5, NOW(), NOW())
|
||||
ON CONFLICT (chunk_start) DO UPDATE SET
|
||||
chunk_end = EXCLUDED.chunk_end,
|
||||
tested_count = EXCLUDED.tested_count,
|
||||
open_count = EXCLUDED.open_count,
|
||||
is_complete = EXCLUDED.is_complete,
|
||||
updated_at = NOW()
|
||||
"#,
|
||||
&[
|
||||
&chunk_start,
|
||||
&chunk_end,
|
||||
&tested_count,
|
||||
&open_count,
|
||||
&record.is_complete,
|
||||
],
|
||||
)
|
||||
.await
|
||||
.context("failed to save scan chunk")?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn upsert_open_endpoint(&self, ip: std::net::Ipv4Addr, port: u16) -> Result<()> {
|
||||
let ip = ip.to_string();
|
||||
let port = i32::from(port);
|
||||
|
||||
self.client
|
||||
.execute(
|
||||
r#"
|
||||
INSERT INTO open_rtsp_endpoint (
|
||||
ip,
|
||||
port,
|
||||
scanned_at,
|
||||
last_seen_at,
|
||||
created_at,
|
||||
updated_at
|
||||
)
|
||||
VALUES ($1::inet, $2, NOW(), NOW(), NOW(), NOW())
|
||||
ON CONFLICT (ip) DO UPDATE SET
|
||||
port = EXCLUDED.port,
|
||||
scanned_at = NOW(),
|
||||
last_seen_at = NOW(),
|
||||
updated_at = NOW()
|
||||
"#,
|
||||
&[&ip, &port],
|
||||
)
|
||||
.await
|
||||
.context("failed to upsert open endpoint")?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn count_open_endpoints(&self, port: u16) -> Result<usize> {
|
||||
let port = i32::from(port);
|
||||
let count: i64 = self
|
||||
.client
|
||||
.query_one(
|
||||
"SELECT COUNT(*)::BIGINT FROM open_rtsp_endpoint WHERE port = $1",
|
||||
&[&port],
|
||||
)
|
||||
.await
|
||||
.context("failed to count open endpoints")?
|
||||
.get(0);
|
||||
|
||||
i64_to_usize(count, "open endpoint count")
|
||||
}
|
||||
|
||||
pub async fn summarize_chunks(&self) -> Result<(usize, usize)> {
|
||||
let row = self
|
||||
.client
|
||||
.query_one(
|
||||
r#"
|
||||
SELECT
|
||||
COALESCE(SUM(tested_count), 0)::BIGINT,
|
||||
COALESCE(SUM(open_count), 0)::BIGINT
|
||||
FROM scan_chunks
|
||||
"#,
|
||||
&[],
|
||||
)
|
||||
.await
|
||||
.context("failed to summarize scan chunks")?;
|
||||
|
||||
let tested = i64_to_usize(row.get::<_, i64>(0), "tested count")?;
|
||||
let open = i64_to_usize(row.get::<_, i64>(1), "open count")?;
|
||||
Ok((tested, open))
|
||||
}
|
||||
}
|
||||
|
||||
fn row_to_chunk_record(row: tokio_postgres::Row) -> Result<ChunkRecord> {
|
||||
Ok(ChunkRecord {
|
||||
chunk_start: i64_to_usize(row.get::<_, i64>(0), "chunk_start")?,
|
||||
chunk_end: i64_to_usize(row.get::<_, i64>(1), "chunk_end")?,
|
||||
tested_count: i64_to_usize(row.get::<_, i64>(2), "tested_count")?,
|
||||
open_count: i64_to_usize(row.get::<_, i64>(3), "open_count")?,
|
||||
is_complete: row.get::<_, bool>(4),
|
||||
})
|
||||
}
|
||||
|
||||
fn i64_to_usize(value: i64, field: &str) -> Result<usize> {
|
||||
usize::try_from(value).with_context(|| format!("{field} out of range: {value}"))
|
||||
}
|
||||
|
||||
fn usize_to_i64(value: usize, field: &str) -> Result<i64> {
|
||||
i64::try_from(value).with_context(|| format!("{field} out of range: {value}"))
|
||||
}
|
||||
@@ -0,0 +1,103 @@
|
||||
use anyhow::{anyhow, Context, Result};
|
||||
use std::net::Ipv4Addr;
|
||||
|
||||
pub fn cidr_to_hosts(cidr: &str) -> Result<Vec<Ipv4Addr>> {
|
||||
let (network, prefix) = cidr
|
||||
.split_once('/')
|
||||
.ok_or_else(|| anyhow!("CIDR inválido: {cidr}"))?;
|
||||
|
||||
let prefix: u8 = prefix
|
||||
.parse()
|
||||
.with_context(|| format!("máscara de sub-rede inválida: {prefix}"))?;
|
||||
if prefix > 32 {
|
||||
return Err(anyhow!("máscara de sub-rede inválida: {prefix}"));
|
||||
}
|
||||
|
||||
let network: Ipv4Addr = network
|
||||
.parse()
|
||||
.with_context(|| format!("endereço IP inválido: {network}"))?;
|
||||
|
||||
let network_u32 = u32::from(network);
|
||||
let mask = if prefix == 0 {
|
||||
0
|
||||
} else {
|
||||
u32::MAX << (32 - prefix)
|
||||
};
|
||||
let network_base = network_u32 & mask;
|
||||
let broadcast = if prefix == 32 {
|
||||
network_base
|
||||
} else {
|
||||
network_base | !mask
|
||||
};
|
||||
|
||||
let start = if prefix >= 31 {
|
||||
network_base
|
||||
} else {
|
||||
network_base + 1
|
||||
};
|
||||
let end = if prefix >= 31 {
|
||||
broadcast
|
||||
} else {
|
||||
broadcast - 1
|
||||
};
|
||||
|
||||
let mut hosts = Vec::new();
|
||||
for ip_u32 in start..=end {
|
||||
hosts.push(Ipv4Addr::from(ip_u32));
|
||||
}
|
||||
|
||||
Ok(hosts)
|
||||
}
|
||||
|
||||
#[allow(dead_code)]
|
||||
pub fn is_routable_ipv4(ip: Ipv4Addr) -> bool {
|
||||
let [a, b, c, _d] = ip.octets();
|
||||
|
||||
if a == 0 {
|
||||
return false;
|
||||
}
|
||||
if a == 10 {
|
||||
return false;
|
||||
}
|
||||
if a == 100 && (64..=127).contains(&b) {
|
||||
return false;
|
||||
}
|
||||
if a == 127 {
|
||||
return false;
|
||||
}
|
||||
if a == 169 && b == 254 {
|
||||
return false;
|
||||
}
|
||||
if a == 172 && (16..=31).contains(&b) {
|
||||
return false;
|
||||
}
|
||||
if a == 192 && b == 0 && c == 0 {
|
||||
return false;
|
||||
}
|
||||
if a == 192 && b == 0 && c == 2 {
|
||||
return false;
|
||||
}
|
||||
if a == 192 && b == 88 && c == 99 {
|
||||
return false;
|
||||
}
|
||||
if a == 192 && b == 168 {
|
||||
return false;
|
||||
}
|
||||
if a == 198 && (b == 18 || b == 19) {
|
||||
return false;
|
||||
}
|
||||
if a == 198 && b == 51 && c == 100 {
|
||||
return false;
|
||||
}
|
||||
if a == 203 && b == 0 && c == 113 {
|
||||
return false;
|
||||
}
|
||||
if (224..=239).contains(&a) {
|
||||
return false;
|
||||
}
|
||||
if a >= 240 {
|
||||
return false;
|
||||
}
|
||||
|
||||
true
|
||||
}
|
||||
+124
@@ -0,0 +1,124 @@
|
||||
mod config;
|
||||
mod coordinator;
|
||||
mod db;
|
||||
mod ip;
|
||||
mod models;
|
||||
mod scan;
|
||||
mod terminal;
|
||||
mod worker;
|
||||
|
||||
use anyhow::Result;
|
||||
use config::{Config, Mode};
|
||||
use db::Database;
|
||||
use std::sync::{
|
||||
atomic::{AtomicBool, Ordering},
|
||||
Arc,
|
||||
};
|
||||
use terminal::{stdout_is_terminal, TerminalUi};
|
||||
|
||||
#[tokio::main]
|
||||
async fn main() -> Result<()> {
|
||||
dotenvy::dotenv().ok();
|
||||
|
||||
let config = Config::from_env_and_args()?;
|
||||
if matches!(config.mode, Mode::Local | Mode::Coordinator) {
|
||||
assert_local_cidrs(&config.cidrs)?;
|
||||
}
|
||||
|
||||
let database = Arc::new(Database::open(&config.database_url).await?);
|
||||
let ui = Arc::new(TerminalUi::new(stdout_is_terminal()));
|
||||
let interrupted = Arc::new(AtomicBool::new(false));
|
||||
|
||||
{
|
||||
let interrupted = Arc::clone(&interrupted);
|
||||
tokio::spawn(async move {
|
||||
let _ = tokio::signal::ctrl_c().await;
|
||||
interrupted.store(true, Ordering::SeqCst);
|
||||
});
|
||||
}
|
||||
|
||||
let result = match config.mode {
|
||||
Mode::Local => scan::run_local_scan(config, database, ui, interrupted).await,
|
||||
Mode::Coordinator => coordinator::run_coordinator(config, database, ui, interrupted).await,
|
||||
Mode::Worker => worker::run_worker(config, ui).await,
|
||||
};
|
||||
|
||||
if let Err(err) = result {
|
||||
eprintln!("error: {err:#}");
|
||||
return Err(err);
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn assert_local_cidrs(cidrs: &[String]) -> Result<()> {
|
||||
if cidrs.is_empty() {
|
||||
return Err(anyhow::anyhow!(
|
||||
"nenhum CIDR foi informado. Use SCAN_CIDRS, SCAN_CIDR_FILE ou --cidrs/--cidr-file"
|
||||
));
|
||||
}
|
||||
|
||||
for cidr in cidrs {
|
||||
assert_local_cidr_inner(cidr)?;
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn assert_local_cidr_inner(cidr: &str) -> Result<()> {
|
||||
let (network, prefix) = cidr
|
||||
.split_once('/')
|
||||
.ok_or_else(|| anyhow::anyhow!("CIDR inválido: {cidr}"))?;
|
||||
|
||||
let prefix: u8 = prefix
|
||||
.parse()
|
||||
.map_err(|_| anyhow::anyhow!("máscara de sub-rede inválida: {prefix}"))?;
|
||||
if prefix > 32 {
|
||||
return Err(anyhow::anyhow!("máscara de sub-rede inválida: {prefix}"));
|
||||
}
|
||||
|
||||
let ip: std::net::Ipv4Addr = network
|
||||
.parse()
|
||||
.map_err(|_| anyhow::anyhow!("endereço IP inválido: {network}"))?;
|
||||
|
||||
let start = u32::from(ip) & cidr_mask(prefix);
|
||||
let end = if prefix == 32 {
|
||||
start
|
||||
} else {
|
||||
start | !cidr_mask(prefix)
|
||||
};
|
||||
|
||||
let private_ranges = [
|
||||
(
|
||||
u32::from(std::net::Ipv4Addr::new(10, 0, 0, 0)),
|
||||
u32::from(std::net::Ipv4Addr::new(10, 255, 255, 255)),
|
||||
),
|
||||
(
|
||||
u32::from(std::net::Ipv4Addr::new(172, 16, 0, 0)),
|
||||
u32::from(std::net::Ipv4Addr::new(172, 31, 255, 255)),
|
||||
),
|
||||
(
|
||||
u32::from(std::net::Ipv4Addr::new(192, 168, 0, 0)),
|
||||
u32::from(std::net::Ipv4Addr::new(192, 168, 255, 255)),
|
||||
),
|
||||
];
|
||||
|
||||
let allowed = private_ranges
|
||||
.iter()
|
||||
.any(|(range_start, range_end)| start >= *range_start && end <= *range_end);
|
||||
if !allowed {
|
||||
return Err(anyhow::anyhow!(
|
||||
"CIDR fora de redes locais permitidas: {cidr}. Use 10.0.0.0/8, 172.16.0.0/12 ou 192.168.0.0/16"
|
||||
));
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn cidr_mask(prefix: u8) -> u32 {
|
||||
if prefix == 0 {
|
||||
0
|
||||
} else {
|
||||
u32::MAX << (32 - prefix)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,28 @@
|
||||
#[derive(Debug, Clone, Copy)]
|
||||
pub struct PendingRange {
|
||||
pub chunk_start: usize,
|
||||
pub chunk_end: usize,
|
||||
pub start: usize,
|
||||
pub end: usize,
|
||||
}
|
||||
|
||||
impl PendingRange {
|
||||
pub fn len(self) -> usize {
|
||||
self.end.saturating_sub(self.start)
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct ChunkRecord {
|
||||
pub chunk_start: usize,
|
||||
pub chunk_end: usize,
|
||||
pub tested_count: usize,
|
||||
pub open_count: usize,
|
||||
pub is_complete: bool,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Default)]
|
||||
pub struct ScanTotals {
|
||||
pub tested: usize,
|
||||
pub open: usize,
|
||||
}
|
||||
+305
@@ -0,0 +1,305 @@
|
||||
use crate::{
|
||||
config::Config,
|
||||
db::Database,
|
||||
ip::cidr_to_hosts,
|
||||
models::{ChunkRecord, PendingRange, ScanTotals},
|
||||
terminal::{ScreenState, TerminalUi},
|
||||
};
|
||||
use anyhow::Result;
|
||||
use std::{
|
||||
net::Ipv4Addr,
|
||||
sync::{
|
||||
atomic::{AtomicBool, AtomicUsize, Ordering},
|
||||
Arc,
|
||||
},
|
||||
time::Instant,
|
||||
};
|
||||
use tokio::{net::TcpStream, time::timeout};
|
||||
|
||||
pub async fn run_local_scan(
|
||||
config: Config,
|
||||
database: Arc<Database>,
|
||||
ui: Arc<TerminalUi>,
|
||||
interrupted: Arc<AtomicBool>,
|
||||
) -> Result<()> {
|
||||
let targets = load_targets(&config.cidrs)?;
|
||||
let target_count = targets.len();
|
||||
let chunk_records = database.load_chunk_records().await?;
|
||||
let (tested_before_run, _open_before_run) = database.summarize_chunks().await?;
|
||||
let pending_ranges = build_pending_ranges(target_count, config.chunk_size, &chunk_records);
|
||||
let pending_count: usize = pending_ranges.iter().map(|range| range.len()).sum();
|
||||
|
||||
if ui.is_enabled() {
|
||||
ui.print_banner();
|
||||
ui.render(&ScreenState {
|
||||
current_ip: if pending_count > 0 {
|
||||
targets[pending_ranges[0].start].to_string()
|
||||
} else {
|
||||
"concluído".to_string()
|
||||
},
|
||||
tested: tested_before_run,
|
||||
total: target_count,
|
||||
});
|
||||
} else {
|
||||
println!(
|
||||
"Scanning {} hosts on {} CIDRs, port {}, concurrency {}, chunk size {}. Pending {}, already tested {}.",
|
||||
target_count,
|
||||
config.cidrs.len(),
|
||||
config.port,
|
||||
config.concurrency,
|
||||
config.chunk_size,
|
||||
pending_count,
|
||||
tested_before_run
|
||||
);
|
||||
}
|
||||
|
||||
if pending_ranges.is_empty() {
|
||||
if ui.is_enabled() {
|
||||
ui.render(&ScreenState {
|
||||
current_ip: "concluído".to_string(),
|
||||
tested: target_count,
|
||||
total: target_count,
|
||||
});
|
||||
ui.finish();
|
||||
} else {
|
||||
println!("Scan concluído em 0.00s");
|
||||
println!("Total de hosts testados: 0");
|
||||
println!("Portas abertas encontradas neste scan: 0");
|
||||
println!(
|
||||
"Total de portas abertas no banco: {}",
|
||||
database.count_open_endpoints(config.port).await?
|
||||
);
|
||||
}
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
let start = Instant::now();
|
||||
let next_range = Arc::new(AtomicUsize::new(0));
|
||||
let tested_this_run = Arc::new(AtomicUsize::new(0));
|
||||
let open_this_run = Arc::new(AtomicUsize::new(0));
|
||||
let mut handles = Vec::new();
|
||||
for _ in 0..config.concurrency.max(1) {
|
||||
let ranges = pending_ranges.clone();
|
||||
let next_range = Arc::clone(&next_range);
|
||||
let tested_this_run = Arc::clone(&tested_this_run);
|
||||
let open_this_run = Arc::clone(&open_this_run);
|
||||
let database = Arc::clone(&database);
|
||||
let ui = Arc::clone(&ui);
|
||||
let targets = targets.clone();
|
||||
let interrupted = Arc::clone(&interrupted);
|
||||
let config = config.clone();
|
||||
|
||||
handles.push(tokio::spawn(async move {
|
||||
loop {
|
||||
if interrupted.load(Ordering::SeqCst) {
|
||||
break;
|
||||
}
|
||||
|
||||
let index = next_range.fetch_add(1, Ordering::SeqCst);
|
||||
let Some(range) = ranges.get(index).copied() else {
|
||||
break;
|
||||
};
|
||||
|
||||
let chunk_stats = scan_range(
|
||||
range,
|
||||
&targets,
|
||||
&config,
|
||||
&database,
|
||||
&ui,
|
||||
&tested_this_run,
|
||||
&open_this_run,
|
||||
&interrupted,
|
||||
tested_before_run,
|
||||
)
|
||||
.await?;
|
||||
|
||||
persist_range(&database, range, chunk_stats).await?;
|
||||
}
|
||||
|
||||
Ok::<(), anyhow::Error>(())
|
||||
}));
|
||||
}
|
||||
|
||||
for handle in handles {
|
||||
handle.await??;
|
||||
}
|
||||
|
||||
let duration = start.elapsed().as_secs_f64();
|
||||
let total_open_count = database.count_open_endpoints(config.port).await?;
|
||||
|
||||
if interrupted.load(Ordering::SeqCst) {
|
||||
if ui.is_enabled() {
|
||||
ui.render(&ScreenState {
|
||||
current_ip: "interrompido".to_string(),
|
||||
tested: tested_before_run + tested_this_run.load(Ordering::SeqCst),
|
||||
total: target_count,
|
||||
});
|
||||
ui.finish();
|
||||
} else {
|
||||
println!("Interrompido após {:.2}s. Progresso salvo.", duration);
|
||||
println!(
|
||||
"Hosts testados nesta execução: {}",
|
||||
tested_this_run.load(Ordering::SeqCst)
|
||||
);
|
||||
println!(
|
||||
"Portas abertas encontradas: {}",
|
||||
open_this_run.load(Ordering::SeqCst)
|
||||
);
|
||||
}
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
if ui.is_enabled() {
|
||||
ui.render(&ScreenState {
|
||||
current_ip: "concluído".to_string(),
|
||||
tested: target_count,
|
||||
total: target_count,
|
||||
});
|
||||
ui.finish();
|
||||
} else {
|
||||
println!("Scan concluído em {:.2}s", duration);
|
||||
println!(
|
||||
"Total de hosts testados: {}",
|
||||
tested_this_run.load(Ordering::SeqCst)
|
||||
);
|
||||
println!(
|
||||
"Portas abertas encontradas neste scan: {}",
|
||||
open_this_run.load(Ordering::SeqCst)
|
||||
);
|
||||
println!("Total de portas abertas no banco: {}", total_open_count);
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn scan_range(
|
||||
range: PendingRange,
|
||||
targets: &[Ipv4Addr],
|
||||
config: &Config,
|
||||
database: &Database,
|
||||
ui: &TerminalUi,
|
||||
tested_this_run: &AtomicUsize,
|
||||
open_this_run: &AtomicUsize,
|
||||
interrupted: &AtomicBool,
|
||||
tested_before_run: usize,
|
||||
) -> Result<ScanTotals> {
|
||||
let mut totals = ScanTotals::default();
|
||||
|
||||
for index in range.start..range.end {
|
||||
if interrupted.load(Ordering::SeqCst) {
|
||||
break;
|
||||
}
|
||||
|
||||
let ip = targets[index];
|
||||
let is_open = check_port_open(ip, config.port, config.timeout_ms).await;
|
||||
totals.tested += 1;
|
||||
let tested_now = tested_this_run.fetch_add(1, Ordering::SeqCst) + 1;
|
||||
|
||||
if is_open {
|
||||
totals.open += 1;
|
||||
open_this_run.fetch_add(1, Ordering::SeqCst);
|
||||
database.upsert_open_endpoint(ip, config.port).await?;
|
||||
}
|
||||
|
||||
if should_render_progress(tested_now, range.end - range.start) || index + 1 == range.end {
|
||||
ui.render(&ScreenState {
|
||||
current_ip: ip.to_string(),
|
||||
tested: tested_before_run + tested_now,
|
||||
total: targets.len(),
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
Ok(totals)
|
||||
}
|
||||
|
||||
async fn persist_range(database: &Database, range: PendingRange, totals: ScanTotals) -> Result<()> {
|
||||
let previous = database.load_chunk_record(range.chunk_start).await?;
|
||||
let tested_before = previous
|
||||
.as_ref()
|
||||
.map(|record| record.tested_count)
|
||||
.unwrap_or(0);
|
||||
let open_before = previous
|
||||
.as_ref()
|
||||
.map(|record| record.open_count)
|
||||
.unwrap_or(0);
|
||||
let tested_count = tested_before + totals.tested;
|
||||
let open_count = open_before + totals.open;
|
||||
let record = ChunkRecord {
|
||||
chunk_start: range.chunk_start,
|
||||
chunk_end: range.chunk_end,
|
||||
tested_count,
|
||||
open_count,
|
||||
is_complete: tested_count >= range.len() + tested_before,
|
||||
};
|
||||
|
||||
database.save_chunk_record(&record).await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn build_pending_ranges(
|
||||
total_targets: usize,
|
||||
chunk_size: usize,
|
||||
records: &[ChunkRecord],
|
||||
) -> Vec<PendingRange> {
|
||||
let mut ranges = Vec::new();
|
||||
let chunk_size = chunk_size.max(1);
|
||||
|
||||
for chunk_start in (0..total_targets).step_by(chunk_size) {
|
||||
let chunk_end = (chunk_start + chunk_size).min(total_targets);
|
||||
let chunk_len = chunk_end.saturating_sub(chunk_start);
|
||||
let record = records
|
||||
.iter()
|
||||
.find(|record| record.chunk_start == chunk_start);
|
||||
let tested_count = record
|
||||
.map(|record| {
|
||||
if record.is_complete {
|
||||
chunk_len
|
||||
} else {
|
||||
record.tested_count.min(chunk_len)
|
||||
}
|
||||
})
|
||||
.unwrap_or(0);
|
||||
|
||||
if tested_count < chunk_len {
|
||||
ranges.push(PendingRange {
|
||||
chunk_start,
|
||||
chunk_end,
|
||||
start: chunk_start + tested_count,
|
||||
end: chunk_end,
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
ranges
|
||||
}
|
||||
|
||||
pub async fn check_port_open(ip: Ipv4Addr, port: u16, timeout_ms: u64) -> bool {
|
||||
let address = std::net::SocketAddr::from((ip, port));
|
||||
matches!(
|
||||
timeout(
|
||||
std::time::Duration::from_millis(timeout_ms),
|
||||
TcpStream::connect(address)
|
||||
)
|
||||
.await,
|
||||
Ok(Ok(_))
|
||||
)
|
||||
}
|
||||
|
||||
pub fn load_targets(cidrs: &[String]) -> Result<Vec<Ipv4Addr>> {
|
||||
let mut targets = Vec::new();
|
||||
|
||||
for cidr in cidrs {
|
||||
targets.extend(cidr_to_hosts(cidr)?);
|
||||
}
|
||||
|
||||
Ok(targets)
|
||||
}
|
||||
|
||||
fn should_render_progress(test_index: usize, chunk_len: usize) -> bool {
|
||||
if chunk_len <= 8 {
|
||||
return true;
|
||||
}
|
||||
|
||||
test_index % 4 == 0
|
||||
}
|
||||
@@ -0,0 +1,99 @@
|
||||
use std::io::{self, IsTerminal, Write};
|
||||
use std::sync::Mutex;
|
||||
|
||||
pub const BANNER: [&str; 9] = [
|
||||
" ░██████ ░██████ ",
|
||||
" ░██ ░██ ░██ ░██ ",
|
||||
"░██ ░██████ ░█████████████░██ ░███████ ░██████ ░████████ ░████████ ░███████ ░██░████ ",
|
||||
"░██ ░██░██ ░██ ░██░████████░██ ░██ ░██ ░██ ░██░██ ░██░██ ░██░███ ",
|
||||
"░██ ░███████░██ ░██ ░██ ░██░██ ░███████ ░██ ░██░██ ░██░█████████░██ ",
|
||||
" ░██ ░██░██ ░██░██ ░██ ░██░██ ░██░██ ░██░██ ░██░██ ░██░██ ░██░██ ░██ ",
|
||||
" ░██████ ░█████░█░██ ░██ ░██ ░██████ ░███████ ░█████░██░██ ░██░██ ░██ ░███████ ░██ ",
|
||||
" ",
|
||||
" ",
|
||||
];
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct ScreenState {
|
||||
pub current_ip: String,
|
||||
pub tested: usize,
|
||||
pub total: usize,
|
||||
}
|
||||
|
||||
pub struct TerminalUi {
|
||||
enabled: bool,
|
||||
lock: Mutex<()>,
|
||||
}
|
||||
|
||||
impl TerminalUi {
|
||||
pub fn new(enabled: bool) -> Self {
|
||||
Self {
|
||||
enabled,
|
||||
lock: Mutex::new(()),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn is_enabled(&self) -> bool {
|
||||
self.enabled
|
||||
}
|
||||
|
||||
pub fn print_banner(&self) {
|
||||
if !self.enabled {
|
||||
return;
|
||||
}
|
||||
|
||||
let _guard = self.lock.lock().expect("terminal mutex poisoned");
|
||||
let mut stdout = io::stdout();
|
||||
let _ = write!(stdout, "\x1b[2J\x1b[H\x1b[?25l");
|
||||
for line in BANNER {
|
||||
let _ = writeln!(stdout, "{line}");
|
||||
}
|
||||
let _ = writeln!(stdout, "IP: aguardando...");
|
||||
let _ = writeln!(
|
||||
stdout,
|
||||
"Progress: [░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░] | 0.0%"
|
||||
);
|
||||
let _ = stdout.flush();
|
||||
}
|
||||
|
||||
pub fn render(&self, state: &ScreenState) {
|
||||
if !self.enabled {
|
||||
return;
|
||||
}
|
||||
|
||||
let _guard = self.lock.lock().expect("terminal mutex poisoned");
|
||||
let mut stdout = io::stdout();
|
||||
let percent = if state.total == 0 {
|
||||
100.0
|
||||
} else {
|
||||
((state.tested as f64 / state.total as f64) * 100.0).min(100.0)
|
||||
};
|
||||
let bar_width = 36usize;
|
||||
let filled = ((percent / 100.0) * bar_width as f64).round() as usize;
|
||||
let empty = bar_width.saturating_sub(filled);
|
||||
let bar = format!("{}{}", "█".repeat(filled), "░".repeat(empty));
|
||||
let ip = format!("{:<15}", state.current_ip);
|
||||
let percent = format!("{:>5.1}%", percent);
|
||||
|
||||
let _ = write!(
|
||||
stdout,
|
||||
"\x1b[2F\x1b[2KIP: {ip}\n\x1b[2KProgress: [{bar}] | {percent}\n"
|
||||
);
|
||||
let _ = stdout.flush();
|
||||
}
|
||||
|
||||
pub fn finish(&self) {
|
||||
if !self.enabled {
|
||||
return;
|
||||
}
|
||||
|
||||
let _guard = self.lock.lock().expect("terminal mutex poisoned");
|
||||
let mut stdout = io::stdout();
|
||||
let _ = write!(stdout, "\x1b[?25h");
|
||||
let _ = stdout.flush();
|
||||
}
|
||||
}
|
||||
|
||||
pub fn stdout_is_terminal() -> bool {
|
||||
io::stdout().is_terminal()
|
||||
}
|
||||
+191
@@ -0,0 +1,191 @@
|
||||
use crate::{
|
||||
config::Config,
|
||||
scan::check_port_open,
|
||||
terminal::{ScreenState, TerminalUi},
|
||||
};
|
||||
use anyhow::{anyhow, Context, Result};
|
||||
use std::{net::Ipv4Addr, sync::Arc};
|
||||
use tokio::{
|
||||
io::{AsyncBufReadExt, AsyncWriteExt, BufReader},
|
||||
net::TcpStream,
|
||||
task::JoinSet,
|
||||
time::sleep,
|
||||
};
|
||||
|
||||
pub async fn run_worker(config: Config, ui: Arc<TerminalUi>) -> Result<()> {
|
||||
if ui.is_enabled() {
|
||||
ui.print_banner();
|
||||
}
|
||||
|
||||
let stream = connect_with_retry(config.coordinator_addr).await?;
|
||||
let (read_half, mut write_half) = stream.into_split();
|
||||
let mut lines = BufReader::new(read_half).lines();
|
||||
write_half
|
||||
.write_all(format!("HELLO {}\n", config.worker_name).as_bytes())
|
||||
.await?;
|
||||
|
||||
loop {
|
||||
let Some(line) = lines.next_line().await? else {
|
||||
break;
|
||||
};
|
||||
|
||||
if line == "DONE" {
|
||||
break;
|
||||
}
|
||||
|
||||
if line == "SHUTDOWN" {
|
||||
break;
|
||||
}
|
||||
|
||||
let job = parse_job_line(&line)?;
|
||||
let result = execute_job(
|
||||
&job.ips,
|
||||
config.port,
|
||||
config.timeout_ms,
|
||||
config.concurrency,
|
||||
&ui,
|
||||
)
|
||||
.await?;
|
||||
|
||||
let result_line = format!(
|
||||
"RESULT {} {} {} {} {}\n",
|
||||
job.chunk_start,
|
||||
job.chunk_end,
|
||||
result.tested,
|
||||
result.open_ips.len(),
|
||||
format_ips(&result.open_ips)
|
||||
);
|
||||
write_half.write_all(result_line.as_bytes()).await?;
|
||||
}
|
||||
|
||||
if ui.is_enabled() {
|
||||
ui.render(&ScreenState {
|
||||
current_ip: "concluído".to_string(),
|
||||
tested: 0,
|
||||
total: 0,
|
||||
});
|
||||
ui.finish();
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn connect_with_retry(addr: std::net::SocketAddr) -> Result<TcpStream> {
|
||||
let mut attempts = 0usize;
|
||||
|
||||
loop {
|
||||
match TcpStream::connect(addr).await {
|
||||
Ok(stream) => return Ok(stream),
|
||||
Err(err) => {
|
||||
attempts += 1;
|
||||
if attempts >= 30 {
|
||||
return Err(err).with_context(|| format!("failed to connect to {addr}"));
|
||||
}
|
||||
|
||||
sleep(std::time::Duration::from_secs(1)).await;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
struct Job {
|
||||
chunk_start: usize,
|
||||
chunk_end: usize,
|
||||
ips: Vec<Ipv4Addr>,
|
||||
}
|
||||
|
||||
struct JobOutcome {
|
||||
tested: usize,
|
||||
open_ips: Vec<Ipv4Addr>,
|
||||
}
|
||||
|
||||
fn parse_job_line(line: &str) -> Result<Job> {
|
||||
let mut parts = line.splitn(6, ' ');
|
||||
let kind = parts.next().unwrap_or("");
|
||||
if kind != "JOB" {
|
||||
return Err(anyhow!("unexpected coordinator message: {line}"));
|
||||
}
|
||||
|
||||
let chunk_start = parse_usize_part(parts.next(), "chunk_start", line)?;
|
||||
let chunk_end = parse_usize_part(parts.next(), "chunk_end", line)?;
|
||||
let _port = parse_usize_part(parts.next(), "port", line)?;
|
||||
let _timeout_ms = parse_usize_part(parts.next(), "timeout_ms", line)?;
|
||||
let ips_part = parts.next().unwrap_or("");
|
||||
let ips = if ips_part.trim().is_empty() {
|
||||
Vec::new()
|
||||
} else {
|
||||
ips_part
|
||||
.split(',')
|
||||
.map(str::trim)
|
||||
.filter(|entry| !entry.is_empty())
|
||||
.map(|entry| {
|
||||
entry
|
||||
.parse::<Ipv4Addr>()
|
||||
.with_context(|| format!("invalid IP in job: {entry}"))
|
||||
})
|
||||
.collect::<Result<Vec<_>>>()?
|
||||
};
|
||||
|
||||
Ok(Job {
|
||||
chunk_start,
|
||||
chunk_end,
|
||||
ips,
|
||||
})
|
||||
}
|
||||
|
||||
async fn execute_job(
|
||||
ips: &[Ipv4Addr],
|
||||
port: u16,
|
||||
timeout_ms: u64,
|
||||
concurrency: usize,
|
||||
ui: &TerminalUi,
|
||||
) -> Result<JobOutcome> {
|
||||
let mut tested = 0usize;
|
||||
let mut open_ips = Vec::new();
|
||||
let concurrency = concurrency.max(1);
|
||||
let mut index = 0usize;
|
||||
|
||||
while index < ips.len() {
|
||||
let end = (index + concurrency).min(ips.len());
|
||||
let mut join_set = JoinSet::new();
|
||||
|
||||
for ip in ips[index..end].iter().copied() {
|
||||
join_set.spawn(async move { (ip, check_port_open(ip, port, timeout_ms).await) });
|
||||
}
|
||||
|
||||
while let Some(join_result) = join_set.join_next().await {
|
||||
let (ip, is_open) = join_result?;
|
||||
tested += 1;
|
||||
if is_open {
|
||||
open_ips.push(ip);
|
||||
}
|
||||
|
||||
if ui.is_enabled() {
|
||||
ui.render(&ScreenState {
|
||||
current_ip: ip.to_string(),
|
||||
tested,
|
||||
total: ips.len(),
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
index = end;
|
||||
}
|
||||
|
||||
Ok(JobOutcome { tested, open_ips })
|
||||
}
|
||||
|
||||
fn parse_usize_part(part: Option<&str>, name: &str, line: &str) -> Result<usize> {
|
||||
let value = part.ok_or_else(|| anyhow!("missing {name} in message: {line}"))?;
|
||||
value
|
||||
.trim()
|
||||
.parse()
|
||||
.map_err(|_| anyhow!("invalid {name} in message: {line}"))
|
||||
}
|
||||
|
||||
fn format_ips(ips: &[Ipv4Addr]) -> String {
|
||||
ips.iter()
|
||||
.map(ToString::to_string)
|
||||
.collect::<Vec<_>>()
|
||||
.join(",")
|
||||
}
|
||||
Reference in New Issue
Block a user