353 lines
10 KiB
Rust
353 lines
10 KiB
Rust
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(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_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(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),
|
|
}
|
|
}
|