603 lines
19 KiB
Rust
603 lines
19 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,
|
|
Rtsp,
|
|
RedisCoordinator,
|
|
RedisWorker,
|
|
}
|
|
|
|
#[derive(Debug, Clone)]
|
|
pub struct Config {
|
|
pub mode: Mode,
|
|
pub cidrs: Vec<String>,
|
|
pub rtsp_targets: Vec<String>,
|
|
pub rtsp_paths: Vec<String>,
|
|
pub redis_url: String,
|
|
pub redis_queue_name: String,
|
|
pub redis_input_csv: String,
|
|
pub redis_worker_count: usize,
|
|
pub snapshot_dir: PathBuf,
|
|
pub sqlite_path: PathBuf,
|
|
pub capture_timeout_ms: u64,
|
|
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: String,
|
|
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())?;
|
|
let rtsp_targets = load_csv_targets(cli.targets_file.as_deref(), cli.targets.as_deref())?;
|
|
let rtsp_paths = load_rtsp_paths(cli.paths_file.as_deref(), cli.paths.as_deref())?;
|
|
Ok(Self {
|
|
mode,
|
|
cidrs,
|
|
rtsp_targets,
|
|
rtsp_paths,
|
|
redis_url: cli
|
|
.redis_url
|
|
.unwrap_or_else(|| read_env("REDIS_URL", "redis://127.0.0.1:6379/0")),
|
|
redis_queue_name: cli
|
|
.redis_queue_name
|
|
.unwrap_or_else(|| read_env("RTSP_QUEUE_NAME", "camfinder:rtsp:queue")),
|
|
redis_input_csv: cli
|
|
.redis_input_csv
|
|
.unwrap_or_else(|| read_env("RTSP_INPUT_CSV", "camfinder-open-rtsp.csv")),
|
|
redis_worker_count: cli
|
|
.redis_worker_count
|
|
.or_else(|| parse_usize("RTSP_WORKER_COUNT", 2).ok())
|
|
.unwrap_or(2),
|
|
snapshot_dir: cli
|
|
.snapshot_dir
|
|
.unwrap_or_else(|| PathBuf::from(read_env("RTSP_SNAPSHOT_DIR", "dump/feeds"))),
|
|
sqlite_path: cli
|
|
.sqlite_path
|
|
.unwrap_or_else(|| PathBuf::from(read_env("RTSP_SQLITE_PATH", "dump/rtsp-results.sqlite"))),
|
|
capture_timeout_ms: cli
|
|
.capture_timeout_ms
|
|
.or_else(|| parse_u64("RTSP_CAPTURE_TIMEOUT_MS", 15000).ok())
|
|
.unwrap_or(15000),
|
|
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
|
|
.unwrap_or_else(|| read_env("COORDINATOR_ADDR", "127.0.0.1:666")),
|
|
worker_name: cli.worker_name.unwrap_or_else(default_worker_name),
|
|
})
|
|
}
|
|
}
|
|
|
|
#[derive(Debug, Default)]
|
|
struct CliArgs {
|
|
cidrs: Option<String>,
|
|
cidr_file: Option<PathBuf>,
|
|
targets: Option<String>,
|
|
targets_file: Option<PathBuf>,
|
|
paths: Option<String>,
|
|
paths_file: Option<PathBuf>,
|
|
redis_url: Option<String>,
|
|
redis_queue_name: Option<String>,
|
|
redis_input_csv: Option<String>,
|
|
redis_worker_count: Option<usize>,
|
|
snapshot_dir: Option<PathBuf>,
|
|
sqlite_path: Option<PathBuf>,
|
|
capture_timeout_ms: Option<u64>,
|
|
database_url: Option<String>,
|
|
coordinator_bind: Option<SocketAddr>,
|
|
coordinator_addr: Option<String>,
|
|
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")?));
|
|
}
|
|
"--targets" => {
|
|
cli.targets = Some(next_value(args, &mut index, "--targets")?);
|
|
}
|
|
"--targets-file" => {
|
|
cli.targets_file = Some(PathBuf::from(next_value(
|
|
args,
|
|
&mut index,
|
|
"--targets-file",
|
|
)?));
|
|
}
|
|
"--paths" => {
|
|
cli.paths = Some(next_value(args, &mut index, "--paths")?);
|
|
}
|
|
"--paths-file" => {
|
|
cli.paths_file =
|
|
Some(PathBuf::from(next_value(args, &mut index, "--paths-file")?));
|
|
}
|
|
"--redis-url" => {
|
|
cli.redis_url = Some(next_value(args, &mut index, "--redis-url")?);
|
|
}
|
|
"--queue" | "--queue-name" => {
|
|
cli.redis_queue_name = Some(next_value(args, &mut index, arg)?);
|
|
}
|
|
"--input-csv" | "--csv-file" => {
|
|
cli.redis_input_csv = Some(next_value(args, &mut index, arg)?);
|
|
}
|
|
"--worker-count" => {
|
|
cli.redis_worker_count = Some(parse_usize_literal(&next_value(
|
|
args,
|
|
&mut index,
|
|
"--worker-count",
|
|
)?)?);
|
|
}
|
|
"--snapshot-dir" => {
|
|
cli.snapshot_dir =
|
|
Some(PathBuf::from(next_value(args, &mut index, "--snapshot-dir")?));
|
|
}
|
|
"--sqlite-path" => {
|
|
cli.sqlite_path =
|
|
Some(PathBuf::from(next_value(args, &mut index, "--sqlite-path")?));
|
|
}
|
|
"--capture-timeout-ms" => {
|
|
cli.capture_timeout_ms = Some(parse_u64_literal(&next_value(
|
|
args,
|
|
&mut index,
|
|
"--capture-timeout-ms",
|
|
)?)?);
|
|
}
|
|
"--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(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),
|
|
"rtsp" => Ok(Mode::Rtsp),
|
|
"redis-coordinator" | "redis-coord" | "redis-server" => Ok(Mode::RedisCoordinator),
|
|
"redis-worker" | "redis-client" => Ok(Mode::RedisWorker),
|
|
other => Err(anyhow!(
|
|
"modo inválido: {other}. Use local, coordinator, worker, rtsp, redis-coordinator ou redis-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 load_csv_targets(file: Option<&Path>, inline: Option<&str>) -> Result<Vec<String>> {
|
|
if let Some(inline) = inline {
|
|
let targets = parse_target_list(inline);
|
|
if !targets.is_empty() {
|
|
return Ok(targets);
|
|
}
|
|
}
|
|
|
|
if let Some(file) = file {
|
|
let content = fs::read_to_string(file)
|
|
.with_context(|| format!("failed to read CSV file {}", file.display()))?;
|
|
let targets = parse_csv_targets(&content);
|
|
if !targets.is_empty() {
|
|
return Ok(targets);
|
|
}
|
|
return Err(anyhow!("arquivo CSV vazio: {}", file.display()));
|
|
}
|
|
|
|
if let Ok(path) = env::var("RTSP_TARGETS_FILE") {
|
|
let path = path.trim();
|
|
if !path.is_empty() {
|
|
let content = fs::read_to_string(path)
|
|
.with_context(|| format!("failed to read CSV file {path}"))?;
|
|
let targets = parse_csv_targets(&content);
|
|
if !targets.is_empty() {
|
|
return Ok(targets);
|
|
}
|
|
return Err(anyhow!("arquivo CSV vazio: {path}"));
|
|
}
|
|
}
|
|
|
|
if let Ok(value) = env::var("RTSP_TARGETS") {
|
|
let value = value.trim();
|
|
if !value.is_empty() {
|
|
let targets = parse_target_list(value);
|
|
if !targets.is_empty() {
|
|
return Ok(targets);
|
|
}
|
|
}
|
|
}
|
|
|
|
Ok(Vec::new())
|
|
}
|
|
|
|
fn parse_csv_targets(input: &str) -> Vec<String> {
|
|
input
|
|
.lines()
|
|
.map(str::trim)
|
|
.filter(|line| !line.is_empty() && !line.starts_with('#'))
|
|
.filter_map(|line| line.split(',').next())
|
|
.map(str::trim)
|
|
.map(strip_quotes)
|
|
.filter(|entry| !entry.is_empty())
|
|
.filter(|entry| !is_csv_header(entry))
|
|
.map(ToString::to_string)
|
|
.collect()
|
|
}
|
|
|
|
fn parse_target_list(input: &str) -> Vec<String> {
|
|
input
|
|
.lines()
|
|
.flat_map(|line| line.split(|c: char| c == ',' || c.is_whitespace()))
|
|
.map(str::trim)
|
|
.map(strip_quotes)
|
|
.filter(|entry| !entry.is_empty())
|
|
.filter(|entry| !is_csv_header(entry))
|
|
.map(ToString::to_string)
|
|
.collect()
|
|
}
|
|
|
|
fn load_rtsp_paths(file: Option<&Path>, inline: Option<&str>) -> Result<Vec<String>> {
|
|
if let Some(inline) = inline {
|
|
let paths = parse_rtsp_paths(inline);
|
|
if !paths.is_empty() {
|
|
return Ok(paths);
|
|
}
|
|
}
|
|
|
|
if let Some(file) = file {
|
|
let content = fs::read_to_string(file)
|
|
.with_context(|| format!("failed to read RTSP path file {}", file.display()))?;
|
|
let paths = parse_rtsp_paths(&content);
|
|
if !paths.is_empty() {
|
|
return Ok(paths);
|
|
}
|
|
return Err(anyhow!("arquivo de paths RTSP vazio: {}", file.display()));
|
|
}
|
|
|
|
if let Ok(path) = env::var("RTSP_PATHS_FILE") {
|
|
let path = path.trim();
|
|
if !path.is_empty() {
|
|
return load_rtsp_paths_from_file(Path::new(path));
|
|
}
|
|
}
|
|
|
|
if let Ok(value) = env::var("RTSP_PATHS") {
|
|
let value = value.trim();
|
|
if !value.is_empty() {
|
|
let paths = parse_rtsp_paths(value);
|
|
if !paths.is_empty() {
|
|
return Ok(paths);
|
|
}
|
|
}
|
|
}
|
|
|
|
load_rtsp_paths_from_file(Path::new("rtsp_paths.txt"))
|
|
}
|
|
|
|
fn load_rtsp_paths_from_file(file: &Path) -> Result<Vec<String>> {
|
|
let content = fs::read_to_string(file)
|
|
.with_context(|| format!("failed to read RTSP path file {}", file.display()))?;
|
|
let paths = parse_rtsp_paths(&content);
|
|
if !paths.is_empty() {
|
|
return Ok(paths);
|
|
}
|
|
Err(anyhow!("arquivo de paths RTSP vazio: {}", file.display()))
|
|
}
|
|
|
|
fn parse_rtsp_paths(input: &str) -> Vec<String> {
|
|
let mut seen = std::collections::HashSet::new();
|
|
input
|
|
.lines()
|
|
.flat_map(|line| line.split('#').next())
|
|
.flat_map(|line| line.split(|c: char| c == ',' || c.is_whitespace()))
|
|
.map(str::trim)
|
|
.map(strip_quotes)
|
|
.filter(|entry| !entry.is_empty())
|
|
.filter(|entry| entry.starts_with('/'))
|
|
.filter_map(|entry| {
|
|
let entry = entry.to_string();
|
|
if seen.insert(entry.clone()) {
|
|
Some(entry)
|
|
} else {
|
|
None
|
|
}
|
|
})
|
|
.collect()
|
|
}
|
|
|
|
fn strip_quotes(value: &str) -> &str {
|
|
value
|
|
.strip_prefix('"')
|
|
.and_then(|value| value.strip_suffix('"'))
|
|
.or_else(|| {
|
|
value
|
|
.strip_prefix('\'')
|
|
.and_then(|value| value.strip_suffix('\''))
|
|
})
|
|
.unwrap_or(value)
|
|
}
|
|
|
|
fn is_csv_header(value: &str) -> bool {
|
|
matches!(
|
|
value.to_ascii_lowercase().as_str(),
|
|
"ip" | "host" | "hostname" | "address" | "target" | "camera"
|
|
)
|
|
}
|
|
|
|
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),
|
|
}
|
|
}
|