watchman/src/watchman.rs

use crate::monitors::devices::USBMon;
use crate::monitors::filechanges::FileChanges;
// use crate::monitors::network::NETMon;
use crate::monitors::ssh_burn_file::SSHBurnMon;
use crate::monitors::clients::NewClientMon;
use config::Config;
use std::path::Path;
use std::process::exit;
use std::sync::{Arc, Mutex};
use std::time::{Duration, SystemTime};

use crate::arguments::Arguments;
use crate::config::setup_service;
use crate::util::{get_cache_dir, get_config_dir, write_default_blacklist, write_default_config};

pub trait EventMonitor {
    async fn check(&mut self);
}

#[derive(Debug)]
pub enum Monitors {
    USBMon(USBMon),
    // NetMon(NETMon),
    SSHBurnMon(SSHBurnMon),
    FileChanges(FileChanges),
    NewClients(NewClientMon),
}
pub struct AppState {
    pub mon_usb: bool,
    pub detection_triggered: bool,
    pub monitors: Arc<Mutex<Vec<Monitors>>>,
    pub settings_map: Config,
}

impl AppState {
    pub fn new() -> Self {
        AppState {
            mon_usb: true,
            detection_triggered: false,
            monitors: Arc::new(Mutex::new(vec![])),
            settings_map: Config::default(),
        }
    }

    pub fn config(&mut self, args: Arguments) {
        // check for config file if it doesn't exist write default config
        println!("Configuring..");
        let cache_dir = get_cache_dir();
        let config_dir = get_config_dir();

        let mut settings_file_path = format!("{}config/settings.toml", config_dir);

        let mut webhook = String::new();
        if let Some(wh) = args.webhook_url {
            webhook = wh
        }

        if let Some(service_command) = args.service {
            match service_command.as_str() {
                "install" => {
                    println!("Installing..");
                    setup_service();
                    if !Path::new(settings_file_path.as_str()).exists() {
                        println!("Settings.toml does not exist");
                        crate::util::write_default_config(settings_file_path.clone(), webhook.clone());
                        exit(0);
                    }
                }
                _ => {}
            }
        }

        if let Some(config_path) = args.config_path {
            settings_file_path = config_path.to_string();
        }

        if !Path::new(settings_file_path.as_str()).exists() {
            println!("Settings.toml does not exist");
            write_default_config(settings_file_path.clone(), webhook);
        }

        let blacklist_file_path = format!("{}config/file_mon_blacklist", config_dir);
        if !Path::new(blacklist_file_path.as_str()).exists() {
            println!("file_mon_blacklist does not exist");
            write_default_blacklist(blacklist_file_path.clone());
        }
        let config = Config::builder();
        let settings = config
            .add_source(config::File::with_name(settings_file_path.as_str()))
            .build()
            .expect("Configuration error: unable to parse settings");

        // initialize monitoring modules
        let mut monitors: Vec<Monitors> = vec![];
        if settings
            .get::<String>("usb_mon_enabled")
            .unwrap()
            .eq_ignore_ascii_case("true")
        {
            monitors.push(Monitors::USBMon(USBMon::new(settings.clone())));
        }
        // if settings
        //     .get::<String>("net_mon_enabled")
        //     .unwrap()
        //     .eq_ignore_ascii_case("true")
        // {
        //     monitors.push(Monitors::NetMon(NETMon::new(settings.clone())));
        // }
        if settings
            .get::<String>("burn_file_mon_enabled")
            .unwrap()
            .eq_ignore_ascii_case("true")
        {
            monitors.push(Monitors::SSHBurnMon(SSHBurnMon::new(settings.clone())));
        }
        if settings
            .get::<String>("fs_mon_enabled")
            .unwrap()
            .eq_ignore_ascii_case("true")
        {
            monitors.push(Monitors::FileChanges(FileChanges::new(
                settings.clone(),
                blacklist_file_path,
                cache_dir,
            )));
        }
        // Enable new client scanning monitor if configured
        if settings
            .get::<String>("client_scan_enabled")
            .unwrap_or("false".to_string())
            .eq_ignore_ascii_case("true")
        {
            monitors.push(Monitors::NewClients(NewClientMon::new(settings.clone())));
        }
        self.monitors.clone_from(&Arc::new(Mutex::new(monitors)));
        self.settings_map.clone_from(&settings)
    }

    pub async fn run(&mut self) {
        let fs_check_tick = self
            .settings_map
            .get::<String>("fs_tick_delay_seconds")
            .expect("tick_delay_seconds not found in Settings.toml");
        let n_fs_check_tick = fs_check_tick
            .parse::<u64>()
            .expect("unable to parse fs_tick_delay_seconds");
        let tick = self
            .settings_map
            .get::<String>("tick_delay_seconds")
            .expect("tick_delay_seconds not found in Settings.toml");
        let n_tick = tick
            .parse::<u64>()
            .expect("unable to parse tick_delay_seconds");
        let mut last = SystemTime::now();
        loop {
            let mut binding = self.monitors.lock();
            let bind = binding.as_mut().unwrap();
            for i in bind.iter_mut() {
                match i {
                    Monitors::USBMon(e) => {
                        // println!("{:#?}", e);
                        e.check().await;
                    }
                    // Monitors::NetMon(e) => {
                    //     // println!("{:#?}", e);
                    //     e.check().await;
                    // }
                    Monitors::SSHBurnMon(e) => {
                        // println!("{:#?}", e);
                        e.check().await;
                    }
                    Monitors::FileChanges(e) => {
                        // println!("{:#?}", e);
                        let now = SystemTime::now();
                        let dur_since = now.duration_since(last).unwrap();
                        if dur_since.as_secs() > n_fs_check_tick {
                            // println!("fs_tick");
                            last = SystemTime::now();
                            e.check().await;
                        }
                    }
                    Monitors::NewClients(e) => {
                        // println!("{:#?}", e);
                        e.check().await;
                    }
                }
                tokio::time::sleep(Duration::new(n_tick, 0)).await;
                // println!("tick");
            }
        }
    }
}