malwaredb 0.3.6

Service for storing malicious, benign, or unknown files and related metadata and relationships.
// SPDX-License-Identifier: Apache-2.0

use malwaredb_server::State;
#[cfg(feature = "anonymous")]
use malwaredb_server::db::admin::User;

use std::sync::Arc;

use anyhow::Result;
use eframe::egui;
use malwaredb_server::crypto::{EncryptionOption, FileEncryption};
use malwaredb_server::db::MDBConfig;

#[allow(clippy::struct_excessive_bools)]
pub struct OptionsPane {
    /// Malware DB server state
    state: Arc<State>,

    /// Malware DB configuration
    config: MDBConfig,

    /// Error message, if any.
    error: String,

    /// If files are to be compressed.
    compress: bool,

    /// User's change for compression.
    change_compress: bool,

    /// Encryption key for files.
    encryption_key: Option<u32>,

    /// User's change for encryption key.
    change_encryption_key: u32,

    /// User's option to create a new encryption key
    new_encryption_key: EncryptionOption,

    /// User's option to create a new encryption key as a string
    new_encryption_key_str: String,

    /// If encryption is enabled.
    enable_encryption: bool,

    /// If we allow sending samples to Virus Total.
    #[cfg(feature = "vt")]
    send_to_vt: bool,

    /// User's change for VT
    #[cfg(feature = "vt")]
    change_send_to_vt: bool,

    /// User ID for anonymous access
    #[cfg(feature = "anonymous")]
    anonymous_user: Option<u32>,

    /// User's change for anonymous access
    #[cfg(feature = "anonymous")]
    change_anonymous_user: Option<u32>,

    /// List of users
    #[cfg(feature = "anonymous")]
    users: Vec<User>,
}

impl OptionsPane {
    pub async fn new(state: Arc<State>) -> Result<Self> {
        Ok(Self {
            config: state.db_type.get_config().await?,
            error: String::new(),
            compress: state.db_config.compression,
            change_compress: state.db_config.compression,
            encryption_key: state.db_config.get_default_key(),
            change_encryption_key: state.db_config.get_default_key().unwrap_or_default(),
            new_encryption_key: EncryptionOption::NONE,
            new_encryption_key_str: String::new(),
            enable_encryption: state.db_config.get_default_key().is_some(),
            #[cfg(feature = "vt")]
            send_to_vt: state.db_config.send_samples_to_vt,
            #[cfg(feature = "vt")]
            change_send_to_vt: state.db_config.send_samples_to_vt,
            #[cfg(feature = "anonymous")]
            anonymous_user: state.db_config.anonymous_uid,
            #[cfg(feature = "anonymous")]
            change_anonymous_user: state.db_config.anonymous_uid,
            #[cfg(feature = "anonymous")]
            users: state.db_type.list_users().await?,
            state,
        })
    }

    fn changed(&self) -> bool {
        #[cfg(feature = "vt")]
        let vt = self.change_send_to_vt != self.send_to_vt;
        #[cfg(not(feature = "vt"))]
        let vt = false;

        #[cfg(feature = "anonymous")]
        let anon = self.change_anonymous_user != self.anonymous_user;
        #[cfg(not(feature = "anonymous"))]
        let anon = false;

        self.change_compress != self.compress
            //|| self.change_encryption_key != self.encryption_key
            || self.new_encryption_key != EncryptionOption::NONE
            || vt
            || anon
    }

    #[allow(clippy::too_many_lines)]
    pub fn ui(&mut self, ui: &mut egui::Ui) {
        ui.checkbox(&mut self.change_compress, "Compress samples");
        #[cfg(feature = "vt")]
        ui.checkbox(&mut self.change_send_to_vt, "Send samples to VirusTotal");

        let keys = futures::executor::block_on(self.state.encryption_keys()).unwrap_or_default();

        // Not having a common sorting causes the GUI to flash.
        let mut keys_list = keys.iter().collect::<Vec<_>>();
        keys_list.sort_by_key(|(id, _)| *id);

        ui.horizontal(|ui| {
            let key_info = if let Some(key_id) = self.encryption_key
                && let Some(key) = keys.get(&key_id)
            {
                format!("{key}, id: {key_id}")
            } else {
                "No encryption key set".into()
            };

            ui.add_enabled(
                false,
                eframe::egui::Checkbox::new(&mut self.enable_encryption, key_info),
            );
            egui::ComboBox::new("change_encryption_key", "")
                .selected_text(self.change_encryption_key.to_string())
                .show_ui(ui, |ui| {
                    for (available_key_id, algo) in keys_list {
                        ui.selectable_value(
                            &mut self.change_encryption_key,
                            *available_key_id,
                            format!("{available_key_id}: {algo}"),
                        );
                    }
                });
            ui.label("New key");
            egui::ComboBox::new("new_encryption_key", "")
                .selected_text(self.new_encryption_key.to_string())
                .show_ui(ui, |ui| {
                    ui.selectable_value(
                        &mut self.new_encryption_key,
                        EncryptionOption::NONE,
                        "NONE",
                    );
                    ui.selectable_value(&mut self.new_encryption_key, EncryptionOption::Xor, "XOR");
                    ui.selectable_value(&mut self.new_encryption_key, EncryptionOption::RC4, "RC4");
                    ui.selectable_value(
                        &mut self.new_encryption_key,
                        EncryptionOption::AES128,
                        "AES-128",
                    );
                });
            if !self.new_encryption_key_str.is_empty() {
                ui.label(format!("Create key type: {}", self.new_encryption_key_str));
            }
        });

        #[cfg(feature = "anonymous")]
        ui.label(format!("Anonymous access: {:?}", self.anonymous_user));

        if self.changed() {
            ui.horizontal(|ui| {
                if ui.button("Save").clicked() {
                    if self.change_compress {
                        if let Err(e) =
                            futures::executor::block_on(self.state.db_type.enable_compression())
                        {
                            self.error = format!("Error enabling compression: {e}");
                        } else {
                            self.compress = true;
                        }
                    } else if let Err(e) =
                        futures::executor::block_on(self.state.db_type.disable_compression())
                    {
                        self.error = format!("Error disabling compression: {e}");
                    } else {
                        self.compress = false;
                    }

                    if self.new_encryption_key != EncryptionOption::NONE {
                        let key = FileEncryption::from(self.new_encryption_key);
                        match futures::executor::block_on(
                            self.state.db_type.add_file_encryption_key(&key),
                        ) {
                            Ok(new_key_id) => self.change_encryption_key = new_key_id,
                            Err(e) => self.error = format!("Error creating encryption key: {e}"),
                        }
                        match futures::executor::block_on(self.state.db_type.get_config()) {
                            Ok(config) => self.config = config,
                            Err(e) => self.error = format!("Error getting config: {e}"),
                        }
                        self.encryption_key = self.config.get_default_key();
                        self.new_encryption_key = EncryptionOption::NONE;
                    }

                    #[cfg(feature = "vt")]
                    {
                        if self.change_send_to_vt {
                            if let Err(e) =
                                futures::executor::block_on(self.state.db_type.enable_vt_upload())
                            {
                                self.error = format!("Error enabling VT upload: {e}");
                            } else {
                                self.send_to_vt = true;
                            }
                        } else if let Err(e) =
                            futures::executor::block_on(self.state.db_type.disable_vt_upload())
                        {
                            self.error = format!("Error disabling VT upload: {e}");
                        } else {
                            self.send_to_vt = false;
                        }
                    }

                    #[cfg(feature = "anonymous")]
                    {
                        if let Some(user_id) = self.change_anonymous_user {
                            if let Err(e) = futures::executor::block_on(
                                self.state.db_type.set_anonymous_user(user_id),
                            ) {
                                self.error = format!("Error setting anonymous user: {e}");
                            } else {
                                self.anonymous_user = Some(user_id);
                            }
                        } else if let Err(e) =
                            futures::executor::block_on(self.state.db_type.clear_anonymous_user())
                        {
                            self.error = format!("Error clearing anonymous user: {e}");
                        }
                    }
                } // end save button

                if ui.button("Reset").clicked() {
                    if let Ok(config) = futures::executor::block_on(self.state.db_type.get_config())
                    {
                        self.compress = config.compression;
                        #[cfg(feature = "vt")]
                        {
                            self.send_to_vt = config.send_samples_to_vt;
                        }
                    }
                    self.change_compress = self.compress;
                    self.change_encryption_key = self.encryption_key.unwrap_or_default();
                    self.enable_encryption = self.encryption_key.is_some();
                    self.new_encryption_key = EncryptionOption::NONE;
                    self.new_encryption_key_str.clear();

                    #[cfg(feature = "vt")]
                    {
                        self.change_send_to_vt = self.send_to_vt;
                    }
                    #[cfg(feature = "anonymous")]
                    {
                        self.change_anonymous_user = self.anonymous_user;
                        if let Ok(users) =
                            futures::executor::block_on(self.state.db_type.list_users())
                        {
                            self.users = users;
                        }
                    }
                    self.error.clear();
                } // end reset button
            }); // end horizontal
        } // end if changed()

        ui.label(&self.error);
    }
}