add subcommand

Author: sofiworker

Cargo.lock

@@ -19,4 +19,7 @@ libc = "0.2"
 ctrlc = { version = "3.4", features = ["termination"] }
 rustix = "0.38.30"
 ipnet = "2.9.0"
-lazy_static = "1.4"
+lazy_static = "1.4"
+config = "0.14.0"
+notify = "6.0"
+clap = { version = "4.5.8", features = ["derive"] }

Cargo.toml

@@ -12,7 +12,7 @@ use std::{borrow::Borrow, collections::HashMap, io::Error, sync::Mutex};
 use tokio::task;
 
 pub struct Capture {
-    pub capture_pool: Arc<Mutex<HashMap<String, task::JoinHandle<u32>>>>,
+    capture_pool: Arc<Mutex<HashMap<String, task::JoinHandle<u32>>>>,
     capture_interface_list: Vec<NetworkInterface>,
 }
 

src/capture/capture.rs

@@ -1,28 +1,20 @@
 use pnet::datalink;
-use std::sync::Mutex;
 
 pub struct InterfaceHandler {
-    interface_list: Mutex<Vec<datalink::NetworkInterface>>,
+    interface_list: Vec<String>,
 }
 
 impl InterfaceHandler {
     pub fn new() -> Self {
         Self {
-            interface_list: Mutex::new(Vec::new()),
+            interface_list: vec![],
         }
     }
 
     pub fn listen_interface(&self) {
-        let listen_interfaces_thread = tokio::spawn(async {
-            // loop {
-            //     datalink::interfaces()
-            // }
-        });
     }
 
     pub fn get_interfaces(&self) -> Vec<datalink::NetworkInterface> {
-        // let guard = self.interface_list.lock().unwrap();
-        // guard.clone()
-        datalink::interfaces()
+        return vec![];
     }
 }

src/capture/interface.rs

@@ -1 +1,33 @@
+use std::process;
+use clap::{command, value_parser, Arg, ArgAction, ArgGroup, ArgMatches, Command};
+use tokio::signal;
+use crate::capture::capture::Capture;
 
+pub fn capture() -> Command {
+    Command::new("capture")
+}
+
+
+async fn start_capture() {
+    // here to start handle interfaces and data capture
+    tokio::spawn(async {
+        let capture = Capture::new();
+        match capture.run() {
+            Ok(_) => {}
+            Err(e) => {
+                panic!("{}", e);
+            }
+        }
+    });
+
+
+    match signal::ctrl_c().await {
+        Ok(()) => {
+            println!("bye!!!!");
+            process::exit(0);
+        }
+        Err(err) => {
+            eprintln!("Unable to listen for shutdown signal: {}", err);
+        }
+    }
+}

src/cmd/capture.rs

@@ -0,0 +1,10 @@
+use clap::{Arg, Command};
+
+pub fn config() -> Command {
+    Command::new("config").args(
+        [
+            Arg::new("test"),
+            Arg::new("")
+        ]
+    )
+}

src/cmd/config.rs

@@ -1,3 +1,5 @@
-mod capture;
-mod reload;
-mod stop;
+pub mod capture;
+pub mod reload;
+pub mod stop;
+pub mod version;
+pub mod config;

src/cmd/mod.rs

@@ -1 +1,5 @@
+use clap::Command;
 
+pub fn reload() -> Command {
+    Command::new("reload")
+}

src/cmd/reload.rs

@@ -0,0 +1,5 @@
+use clap::Command;
+
+pub fn stop() -> Command {
+    Command::new("stop")
+}

src/cmd/stop.rs

@@ -0,0 +1,5 @@
+use clap::Command;
+
+pub fn version() -> Command {
+    Command::new("version")
+}

src/cmd/version.rs

@@ -1,26 +1,36 @@
-use lazy_static::lazy_static;
+#![allow(deprecated)]
+use std::collections::HashMap;
 use serde_derive::{Deserialize, Serialize};
-use std::fs::File;
-use std::io::Read;
+use config::{Config, File};
+use notify::{Event, RecommendedWatcher, RecursiveMode, Watcher};
+use std::path::Path;
+use std::sync::mpsc::channel;
+use std::sync::RwLock;
+use std::thread;
+use std::time::Duration;
+
+const CONFIG_PATH :&str = "src/config.toml";
 
 #[derive(Serialize, Deserialize, Debug)]
-pub struct Config {
-    pub net: NetConfig,
-    pub service: ServiceConfig,
-    pub log: LogConfig,
+pub struct ArkimeConfig {
+    net: NetConfig,
+    service: ServiceConfig,
+    log: LogConfig,
 }
 
 #[derive(Serialize, Deserialize, Debug)]
 pub struct NetConfig {
-    pub dynamic_interfaces: bool,
-    pub bpf_filter: String,
-    pub link_name: Vec<String>,
+    dynamic_interfaces: bool,
+    bpf_filter: String,
+    link_name: Vec<String>,
 }
 
 #[derive(Serialize, Deserialize, Debug)]
 pub struct ServiceConfig {
+    hostname: String,
     host: String,
     port: u32,
+    socket_path: String,
 }
 
 #[derive(Serialize, Deserialize, Debug)]
@@ -29,30 +39,72 @@ pub struct LogConfig {
     output: String,
 }
 
-impl Default for Config {
-    fn default() -> Self {
-        let file_path = "src/config.toml";
+lazy_static::lazy_static! {
+    static ref SETTINGS: RwLock<Config> = RwLock::new({
+        let settings = Config::builder()
+        .add_source(config::Environment::with_prefix("ARKIME-RUST").separator("_"))
+        .add_source(File::with_name(CONFIG_PATH));
+        settings.build().unwrap()
+    });
+}
 
-        let mut file = match File::open(file_path) {
-            Ok(f) => f,
-            Err(e) => panic!("error is op en conf {e}"),
-        };
+fn watch() {
+    thread::spawn(|| {
+        let (tx, rx) = channel();
 
-        let mut str = String::new();
-        match file.read_to_string(&mut str) {
-            Ok(s) => s,
-            Err(e) => panic!("error read str {}", e),
-        };
+        let mut watcher: RecommendedWatcher = Watcher::new(
+            tx,
+            notify::Config::default().with_poll_interval(Duration::from_secs(1000)),
+        )
+            .unwrap();
 
-        toml::from_str(&str).expect("parse config file failed")
-    }
-}
+        watcher
+            .watch(
+                Path::new(CONFIG_PATH),
+                RecursiveMode::NonRecursive,
+            )
+            .unwrap();
+        loop {
+            match rx.recv() {
+                Ok(Ok(Event {
+                          kind: notify::event::EventKind::Modify(_),
+                          ..
+                      })) => {
+                    println!(" * Settings.toml written; refreshing configuration ...");
+                    SETTINGS.write().unwrap().refresh().unwrap();
+                    show();
+                }
+
+                Err(e) => println!("watch error: {:?}", e),
 
-impl Config {
-    pub fn get<'a>() -> &'a Self {
-        lazy_static! {
-            static ref CACHE: Config = Config::default();
+                _ => {
+                    // Ignore event
+                }
+            }
         }
-        &CACHE
+    });
+}
+
+impl ArkimeConfig {
+
+    pub fn init() {
+        SETTINGS.read().unwrap();
+        watch();
+    }
+
+    pub fn get_interface(self) -> Vec<String> {
+        return self.net.link_name;
     }
 }
+
+fn show() {
+    println!(
+        " * Settings :: \n\x1b[31m{:?}\x1b[0m",
+        SETTINGS
+            .read()
+            .unwrap()
+            .clone()
+            .try_deserialize::<HashMap<String, String>>()
+            .unwrap()
+    );
+}

src/conf/conf.rs

@@ -1,3 +1,3 @@
 pub mod conf;
 
-pub use conf::Config;
+pub use conf::ArkimeConfig;

src/conf/mod.rs

@@ -1,15 +1,14 @@
 [net]
-dynamic_interfaces = true
 bpf_filter = ""
-link_name = ["lo"]
+dynamic_interfaces = true
+link_name = ["lo", "ens*"]
 
 [service]
+hostname = ""
 host = "127.0.0.1"
 port = 8080
-
-[core]
-hostname = ""
+socket_path = ""
 
 [log]
 level = "info"
-output = "arkime.log"
+output = "arkime-rust.log"

src/config.toml

@@ -1,7 +1,4 @@
-use std::process;
-
-use capture::capture::Capture;
-use tokio::signal;
+use clap::Command;
 
 mod capture;
 mod cmd;
@@ -17,29 +14,19 @@ mod layer;
 
 #[tokio::main]
 async fn main() {
-    conf::Config::get();
-
-    // here to start handle interfaces and data capture
-    tokio::spawn(async {
-        init_capture();
-    });
-    match signal::ctrl_c().await {
-        Ok(()) => {
-            println!("bye!!!!");
-            process::exit(0);
-        }
-        Err(err) => {
-            eprintln!("Unable to listen for shutdown signal: {}", err);
-        }
-    }
-}
+    let root_cmd = Command::new("arkime-rust").subcommands([
+        cmd::capture::capture(),
+        cmd::config::config(),
+        cmd::reload::reload(),
+        cmd::stop::stop(),
+        cmd::version::version()
+    ]);
+    let mut help_cmd = root_cmd.clone();
+    let matches = root_cmd.get_matches();
 
-fn init_capture() {
-    let capture = Capture::new();
-    match capture.run() {
-        Ok(_) => {}
-        Err(e) => {
-            panic!("{}", e);
-        }
+    match matches.subcommand() {
+        Some(("capture", capture_cmd)) => {},
+        _ => help_cmd.print_help().unwrap()
     }
+    conf::ArkimeConfig::init();
 }