1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86

use std::collections::HashMap;
use std::ops::Index;

use registry::*;
use strum::IntoEnumIterator;
use strum_macros::EnumIter;

use crate::extensions::ExtensionList;

#[derive(Debug, EnumIter, PartialEq, Eq, Hash, Clone)]
pub enum Param {
    DebugPath,
    ConfigPath,
    NumVersion,
    UtilsPath,
    AppId,
    KillPolicy,
}

#[derive(PartialEq)]
pub enum KillPolicy {
    Suspend,
    Kill,
    DoNothing,
}

impl Param {
    fn convert_to_str(param: &Param) -> &str {
        match param {
            Param::ConfigPath => "CONFIG_PATH", // incidents reports, exclusions list
            Param::NumVersion => "NUM_VERSION",
            Param::DebugPath => "DEBUG_PATH", // dir with prediction.csv (used for debug)
            Param::UtilsPath => "UTILS_PATH", // toast.exe
            Param::AppId => "APP_ID",         // AppUserModelID for toast notifications
            Param::KillPolicy => "KILL_POLICY", // SUSPEND / KILL
        }
    }
}

#[derive(Debug)]
pub struct Config {
    params: HashMap<Param, String>,
    pub extensions_list: ExtensionList,
    pub threshold_drivermsgs: usize,
    pub threshold_prediction: f32,
    pub timesteps_stride: usize,
}

impl Config {
    pub fn new() -> Config {
        let mut params: HashMap<Param, String> = HashMap::new();
        for param in Param::iter() {
            let regkey = Hive::LocalMachine
                .open(r"SOFTWARE\Owlyshield", Security::Read)
                .expect("Cannot open registry hive");
            let val = regkey
                .value(Param::convert_to_str(&param))
                .expect(&format!("Cannot open registry key {:?}", param))
                .to_string();
            params.insert(param, val);
        }
        Config {
            params,
            extensions_list: ExtensionList::new(),
            threshold_drivermsgs: 100,
            threshold_prediction: 0.65,
            timesteps_stride: 20,
        }
    }

    pub fn get_kill_policy(&self) -> KillPolicy {
        match self[Param::KillPolicy].as_str() {
            "KILL" => KillPolicy::Kill,
            "SUSPEND" => KillPolicy::Suspend,
            &_ => KillPolicy::DoNothing,
        }
    }
}

impl Index<Param> for Config {
    type Output = String;

    fn index(&self, index: Param) -> &Self::Output {
        &self.params[&index]
    }
}