Added scaffoded project

This commit is contained in:
2026-06-29 04:39:41 -03:00
parent ab82a4b71c
commit 9bd945cc51
17 changed files with 2797 additions and 0 deletions
+352
View File
@@ -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),
}
}
+398
View File
@@ -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(",")
}
+218
View File
@@ -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}"))
}
+103
View File
@@ -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
View File
@@ -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)
}
}
+28
View File
@@ -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
View File
@@ -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
}
+99
View File
@@ -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
View File
@@ -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(",")
}