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
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126

use std::collections::{HashMap, HashSet};

use strum::IntoEnumIterator;
use strum_macros::EnumIter;

use crate::extensions::ExtensionCategory::*;

#[derive(Debug)]
pub struct ExtensionsCount {
    pub categories_set: HashMap<ExtensionCategory, HashSet<String>>,
    pub extensionlist: ExtensionList,
}

#[derive(Debug, Eq, PartialEq, Hash, Copy, Clone, EnumIter)]
pub enum ExtensionCategory {
    Docs,
    Config,
    Archives,
    Database,
    Code,
    Exe,
    Email,
    PasswordVault,
    Event,
    Others,
}

#[derive(Debug)]
pub struct ExtensionList {
    pub categories: HashMap<ExtensionCategory, Vec<&'static str>>,
}

impl ExtensionList {
    pub fn new() -> ExtensionList {
        let doc = vec![
            "doc", "docx", "docb", "docm", "pdf", "djvu", "odt", "xls", "xlsx", "csv", "tsv",
            "ppt", "pptx", "pst", "ost", "msg", "eml", "vsd", "vsdx", "txt", "rtf", "jpeg", "jpg",
            "png", "gif", "tiff", "tif", "bmp", "raw", "psd", "svg", "mp3", "flac", "alac", "wav",
            "aac", "ogg", "wma", "mp4", "mkv", "wmv", "flv", "mpg", "avi",
        ];
        let config = vec!["ini", "inf", "json", "yaml", "yml", "xml"];
        let archive = vec!["zip", "rar", "7z", "gz", "tgz", "tar", "gzip", "dat"];
        let database = vec![
            "db", "sql", "sqlitedb", "sqlite", "sqlite3", "dbf", "fdb", "mdb", "mde", "ora",
            "abbdb", "abbde", "odb", "sdf", "mdf", "ldf", "ndf", "kdbx",
        ];
        let code = vec![
            "iso", "jar", "c", "h", "hpp", "cpp", "cxx", "hxx", "java", "class", "php", "js",
            "html", "sh", "asp", "sh", "jar", "rb", "jsp", "cs", "vb", "pl", "py", "rst",
        ];
        let exe = vec!["exe", "dll"];
        let email = vec!["eml", "email"];
        let password_vault = vec![
            "1pux",
            "opvault",
            "agilekeychain",
            "kdb",
            "kdbx",
            "pwrep",
            "pgpf",
            "psw",
            "passwordwallet4",
            "pswx",
        ];
        let event = vec!["evtx"];
        let others = vec![];

        let mut categories = HashMap::new();
        categories.insert(Docs, doc);
        categories.insert(Config, config);
        categories.insert(Archives, archive);
        categories.insert(Database, database);
        categories.insert(Code, code);
        categories.insert(Exe, exe);
        categories.insert(Email, email);
        categories.insert(PasswordVault, password_vault);
        categories.insert(Event, event);
        categories.insert(Others, others);

        ExtensionList {
            categories: categories,
        }
    }

    pub fn get_extension_category(&self, extension: &str) -> ExtensionCategory {
        let extension_low = &extension.to_lowercase();
        for (k, v) in &self.categories {
            if v.contains(&&**extension_low) {
                return *k;
            }
        }
        return Others;
    }
}

impl ExtensionsCount {
    pub fn new() -> ExtensionsCount {
        let mut cats_entries: HashMap<ExtensionCategory, HashSet<String>> = HashMap::new();
        for cat in ExtensionCategory::iter() {
            cats_entries.insert(cat, HashSet::new());
        }
        ExtensionsCount {
            categories_set: cats_entries,
            extensionlist: ExtensionList::new(),
        }
    }

    pub fn count_all(&self) -> usize {
        ExtensionCategory::iter()
            .map(|c| self.categories_set[&c].len())
            .sum()
    }

    pub fn count_category(&self, cat: ExtensionCategory) -> usize {
        self.categories_set[&cat].len()
    }

    pub fn add_cat_extension(&mut self, extension: &str) {
        let extension = extension.trim_matches(char::from(0));
        if !extension.is_empty() {
            let extension_category = self.extensionlist.get_extension_category(extension);
            let val = self.categories_set.get_mut(&extension_category).unwrap();
            val.insert(String::from(extension));
        }
    }
}