app configuration as config file

Author: matthieugouel

Cargo.lock

@@ -9,6 +9,7 @@ caracat = { git = "https://github.com/maxmouchet/caracat" }
 chrono = "0.4.39"
 clap = { version = "4.5.20", features = ["derive"] }
 clap-verbosity-flag = "3.0.0"
+config = "0.15.6"
 csv = "1.3.1"
 env_logger = "0.11.5"
 hyperloglog = "1.0.2"

Cargo.toml

@@ -0,0 +1,69 @@
+use config::Config as AppConfig;
+
+#[allow(dead_code)]
+pub struct ProberConfig {
+    /// Prober ID
+    pub prober_id: u16,
+
+    /// Kafka brokers
+    pub brokers: String,
+
+    /// Kafka consumer topics (comma separated)
+    pub in_topics: String,
+
+    /// Kafka consumer group ID
+    pub in_group_id: String,
+
+    /// Kafka producer topic
+    pub out_topic: String,
+
+    /// Kafka producer Authentication Protocol
+    pub out_auth_protocol: String,
+
+    /// Kafka producer Authentication SASL Username
+    pub out_auth_sasl_username: String,
+
+    /// Kafka producer Authentication SASL Password
+    pub out_auth_sasl_password: String,
+
+    /// Kafka producer Authentication SASL Mechanism
+    pub out_auth_sasl_mechanism: String,
+}
+
+pub fn load_config(config_path: &str) -> AppConfig {
+    AppConfig::builder()
+        .add_source(config::File::with_name(config_path))
+        .add_source(config::Environment::with_prefix("OSIRIS"))
+        .build()
+        .unwrap()
+}
+
+pub fn prober_config(config: AppConfig) -> ProberConfig {
+    ProberConfig {
+        prober_id: config.get_int("prober.id").unwrap_or(0) as u16,
+        brokers: config
+            .get_string("kafka.brokers")
+            .unwrap_or("localhost:9092".to_string()),
+        in_topics: config
+            .get_string("kafka.in_topics")
+            .unwrap_or("osiris".to_string()),
+        in_group_id: config
+            .get_string("kafka.in_group_id")
+            .unwrap_or("osiris".to_string()),
+        out_topic: config
+            .get_string("kafka.out_topic")
+            .unwrap_or("osiris-results".to_string()),
+        out_auth_protocol: config
+            .get_string("kafka.out_auth_protocol")
+            .unwrap_or("PLAINTEXT".to_string()),
+        out_auth_sasl_username: config
+            .get_string("kafka.out_auth_sasl_username")
+            .unwrap_or("osiris".to_string()),
+        out_auth_sasl_password: config
+            .get_string("kafka.out_auth_sasl_password")
+            .unwrap_or("osiris".to_string()),
+        out_auth_sasl_mechanism: config
+            .get_string("kafka.out_auth_sasl_mechanism")
+            .unwrap_or("SCRAM-SHA-512".to_string()),
+    }
+}

src/config.rs

@@ -1,68 +1,13 @@
 use anyhow::Result;
 use caracat::models::Probe;
-use caracat::rate_limiter::RateLimitingMethod;
 use ipnet::IpNet;
 use rand::seq::SliceRandom;
 use rand::thread_rng;
-use std::net::Ipv4Addr;
-use std::net::Ipv6Addr;
-use std::time::Duration;
 use tokio::task;
 
-use crate::prober::probe;
-use crate::producer::{produce, KafkaAuth};
-
-/// Probing configuration.
-#[derive(Debug)]
-pub struct Config {
-    /// Number of probes to send before calling the rate limiter.
-    pub batch_size: u64,
-    /// Identifier encoded in the probes (random by default).
-    pub instance_id: u16,
-    /// Whether to actually send the probes on the network or not.
-    pub dry_run: bool,
-    /// Do not send probes with ttl < min_ttl.
-    pub min_ttl: Option<u8>,
-    /// Do not send probes with ttl > max_ttl.
-    pub max_ttl: Option<u8>,
-    /// Check that replies match valid probes.
-    pub integrity_check: bool,
-    /// Interface from which to send the packets.
-    pub interface: String,
-    /// Source IPv4 address
-    pub src_ipv4_addr: Option<Ipv4Addr>,
-    /// Source IPv6 address
-    pub src_ipv6_addr: Option<Ipv6Addr>,
-    /// Maximum number of probes to send (unlimited by default).
-    pub max_probes: Option<u64>,
-    /// Number of packets to send per probe.
-    pub packets: u64,
-    /// Probing rate in packets per second.
-    pub probing_rate: u64,
-    /// Method to use to limit the packets rate.
-    pub rate_limiting_method: RateLimitingMethod,
-    /// Time in seconds to wait after sending the probes to stop the receiver.
-    pub receiver_wait_time: Duration,
-}
-
-fn create_config() -> Config {
-    Config {
-        batch_size: 128,
-        dry_run: false,
-        min_ttl: None,
-        max_ttl: None,
-        integrity_check: true,
-        instance_id: 0,
-        interface: caracat::utilities::get_default_interface(),
-        src_ipv4_addr: None,
-        src_ipv6_addr: None,
-        max_probes: None,
-        packets: 1,
-        probing_rate: 100,
-        rate_limiting_method: caracat::rate_limiter::RateLimitingMethod::Auto,
-        receiver_wait_time: Duration::new(3, 0),
-    }
-}
+use crate::config::ProberConfig;
+use crate::prober::{load_caracat_config, probe};
+use crate::producer::{produce, KafkaAuth, SaslAuth};
 
 struct Target {
     prefix: IpNet,
@@ -137,26 +82,34 @@ fn generate_probes(target: &Target) -> Result<Vec<Probe>> {
     Ok(probes)
 }
 
-pub async fn handle(
-    prober_id: u16,
-    brokers: &str,
-    _in_topics: &str,
-    _in_group_id: &str,
-    out_topic: &str,
-    out_auth: KafkaAuth,
-    target: &str,
-) -> Result<()> {
+pub async fn handle(config: &ProberConfig, target: &str) -> Result<()> {
+    // Configure Kafka producer authentication
+    let out_auth = match config.out_auth_protocol.as_str() {
+        "PLAINTEXT" => KafkaAuth::PlainText,
+        "SASL_PLAINTEXT" => KafkaAuth::SasalPlainText(SaslAuth {
+            username: config.out_auth_sasl_username.clone(),
+            password: config.out_auth_sasl_password.clone(),
+            mechanism: config.out_auth_sasl_mechanism.clone(),
+        }),
+        _ => {
+            return Err(anyhow::anyhow!(
+                "Invalid Kafka producer authentication protocol"
+            ))
+        }
+    };
+
     // Probe Generation
-    let config = create_config();
+    let caracat_config = load_caracat_config();
     let target = decode_payload(target)?;
     let probes_to_send = generate_probes(&target)?;
 
     // Probing
-    let result = task::spawn_blocking(move || probe(config, probes_to_send.into_iter())).await?;
+    let result =
+        task::spawn_blocking(move || probe(caracat_config, probes_to_send.into_iter())).await?;
     let (_, _, results) = result?;
 
     // Produce the results to Kafka topic
-    produce(brokers, out_topic, prober_id, out_auth, results).await;
+    produce(config, out_auth, results).await;
 
     Ok(())
 }

src/handler.rs

@@ -1,67 +1,51 @@
+mod config;
 mod consumer;
 mod handler;
 mod prober;
 mod producer;
 
 use anyhow::Result;
 use chrono::Local;
-use clap::Parser as CliParser;
+use clap::{Args, Parser, Subcommand};
 use clap_verbosity_flag::{InfoLevel, Verbosity};
 use env_logger::Builder;
 use std::io::Write;
 
 // use crate::consumer::consume;
+use crate::config::{load_config, prober_config};
 use crate::handler::handle;
 
-#[derive(CliParser, Debug, Clone)]
-#[command(version, about, long_about = None)]
-struct Cli {
-    /// Prober ID
-    #[clap(long, default_value = "0")]
-    prober_id: u16,
-
-    /// Kafka brokers
-    #[clap(long, default_value = "localhost:9092")]
-    brokers: String,
-
-    /// Kafka consumer topics (comma separated)
-    #[clap(long, default_value = "osiris")]
-    in_topics: String,
-
-    /// Kafka consumer group ID
-    #[clap(long, default_value = "osiris")]
-    in_group_id: String,
-
-    /// Kafka producer topic
-    #[clap(long, default_value = "osiris-results")]
-    out_topic: String,
-
-    /// Kafka producer Authentication Protocol
-    #[clap(long, default_value = "PLAINTEXT")]
-    out_auth_protocol: String,
-
-    /// Kafka producer Authentication SASL Username
-    #[clap(long, default_value = "osiris")]
-    out_auth_sasl_username: String,
-
-    /// Kafka producer Authentication SASL Password
-    #[clap(long, default_value = "osiris")]
-    out_auth_sasl_password: String,
-
-    /// Kafka producer Authentication SASL Mechanism
-    #[clap(long, default_value = "SCRAM-SHA-512")]
-    out_auth_sasl_mechanism: String,
+#[derive(Debug, Parser)]
+#[clap(name = "Osiris", version)]
+pub struct App {
+    #[clap(flatten)]
+    global_opts: GlobalOpts,
 
-    /// Target (eg., 2606:4700:4700::1111/128,1,32,1)
-    #[arg(index = 1)]
-    target: String,
+    #[clap(subcommand)]
+    command: Command,
+}
+#[derive(Debug, Subcommand)]
+#[command(version, about, long_about = None)]
+enum Command {
+    Prober {
+        /// Configuration file
+        #[arg(short, long)]
+        config: String,
+
+        /// Target (eg., 2606:4700:4700::1111/128,1,32,1)
+        #[arg(index = 1)]
+        target: String,
+    },
+}
 
+#[derive(Debug, Args)]
+struct GlobalOpts {
     /// Verbosity level
     #[clap(flatten)]
     verbose: Verbosity<InfoLevel>,
 }
 
-fn set_logging(cli: &Cli) {
+fn set_logging(cli: &GlobalOpts) {
     Builder::new()
         .format(|buf, record| {
             writeln!(
@@ -78,37 +62,19 @@ fn set_logging(cli: &Cli) {
 
 #[tokio::main]
 async fn main() -> Result<()> {
-    let cli = Cli::parse();
-    set_logging(&cli);
-
-    // Configure Kafka producer authentication
-    let out_auth = match cli.out_auth_protocol.as_str() {
-        "PLAINTEXT" => producer::KafkaAuth::PLAINTEXT,
-        "SASL" => producer::KafkaAuth::SASL(producer::SaslAuth {
-            username: cli.out_auth_sasl_username.clone(),
-            password: cli.out_auth_sasl_password.clone(),
-            mechanism: cli.out_auth_sasl_mechanism.clone(),
-        }),
-        _ => {
-            return Err(anyhow::anyhow!(
-                "Invalid Kafka producer authentication protocol"
-            ))
+    let cli = App::parse();
+    set_logging(&cli.global_opts);
+
+    match cli.command {
+        Command::Prober { config, target } => {
+            let app_config = load_config(&config);
+            let prober_config = prober_config(app_config);
+
+            match handle(&prober_config, &target).await {
+                Ok(_) => (),
+                Err(e) => log::error!("Error: {}", e),
+            }
         }
-    };
-
-    match handle(
-        cli.prober_id,
-        &cli.brokers,
-        &cli.in_topics,
-        &cli.in_group_id,
-        &cli.out_topic,
-        out_auth,
-        &cli.target,
-    )
-    .await
-    {
-        Ok(_) => (),
-        Err(e) => log::error!("Error: {}", e),
     }
 
     Ok(())