Added registry functionality to allow notification to have a logo

Cargo.toml

@@ -25,3 +25,6 @@ keyring = { version = "3", features = ["apple-native", "windows-native", "sync-s
 uuid = { version = "1.11.0", features = ["v4", "fast-rng", "macro-diagnostics"] }
 async-trait = "0.1.83"
 mockall = "0.13.1"
+
+[target.'cfg(windows)'.dependencies]
+winreg = "0.10"

src/main.rs

@@ -2,8 +2,9 @@ mod api;
 mod controller;
 mod creds;
 mod models;
-
 use crate::controller::Controller;
+#[cfg(target_os = "windows")]
+mod registry;
 use ::time::UtcOffset;
 use api::Api;
 use clap::{Parser, Subcommand};
@@ -25,6 +26,10 @@ use std::process::exit;
 use std::process::Command;
 #[cfg(target_os = "windows")]
 use std::process::Stdio;
+#[cfg(target_os = "windows")]
+use winreg::enums::*;
+#[cfg(target_os = "windows")]
+use winreg::RegKey;
 
 #[derive(Serialize, Deserialize)]
 struct Pid {
@@ -71,6 +76,7 @@ const CONFIG_DIR: &str = ".tuenroll";
 const PID_FILE: &str = "process.json";
 const LAST_CHECK_FILE: &str = "last_check.json";
 const LOG_FILE: &str = "tuenroll.log";
+const LOGO: &str = "logo.png";
 
 #[allow(clippy::zombie_processes)]
 #[tokio::main]
@@ -84,6 +90,11 @@ async fn main() {
             "{}",
             "Automate your TU Delft exam registrations. Let's get you set up!".bright_cyan()
         );
+
+        download_logo().await;
+        // Sets up the registry values to be able to display notifications with a logo
+        #[cfg(target_os = "windows")]
+        setup_registry();
     }
 
     set_up_logging();
@@ -443,13 +454,17 @@ fn process_is_running() -> bool {
 }
 
 fn show_notification(body: &str) {
-    if let Err(e) = Notification::new()
-        .summary("TUEnroll")
-        .body(body)
-        .icon("info")
-        .timeout(5 * 1000) // 5 seconds
-        .show()
-    {
+    let mut notification = Notification::new();
+
+    let notification = notification.body(body).timeout(5 * 1000); // 5 seconds
+
+    #[cfg(target_os = "windows")]
+    let notification = notification.app_id(APP_NAME);
+
+    #[cfg(target_os = "linux")]
+    let notification = notification.image_path(get_config_path(CONFIG_DIR, LOGO).to_str().unwrap());
+
+    if let Err(e) = notification.show() {
         error!("Failed to show notification: {}", e);
     }
 }
@@ -539,6 +554,60 @@ fn run_on_boot_linux(interval: &u32) -> Result<(), Box<dyn std::error::Error>> {
     Ok(())
 }
 
+#[cfg(target_os = "windows")]
+fn setup_registry() {
+    use registry::RegistryHandler;
+
+    if registry::registry(
+        get_config_path(CONFIG_DIR, LOGO).to_str().unwrap(),
+        APP_NAME,
+        &RegistryHandler::new(RegKey::predef(HKEY_CURRENT_USER)),
+    )
+    .is_ok()
+    {
+        info!("Registry succesfully setup.");
+    } else {
+        error!("Registry setup was unsuccesful");
+    }
+}
+
+async fn download_logo() {
+    let logo_path = get_config_path(CONFIG_DIR, LOGO);
+
+    if logo_path.exists() {
+        return;
+    }
+
+    let logo_url = "https://raw.githubusercontent.com/dhruvan2006/tuenroll/main/logo.png";
+
+    if let Some(parent) = logo_path.parent() {
+        std::fs::create_dir_all(parent).expect("Failed to create log directory");
+    }
+    let _ = std::fs::OpenOptions::new()
+        .write(true)
+        .create(true)
+        .truncate(true)
+        .open(&logo_path)
+        .unwrap();
+
+    let response = reqwest::get(logo_url)
+        .await
+        .expect("Request did not succeed. Try again later.");
+
+    info!("Downloading logo.");
+
+    std::fs::write(
+        &logo_path,
+        response
+            .bytes()
+            .await
+            .expect("Error occured while downloading image bytes"),
+    )
+    .expect("Could not write to file.");
+
+    info!("Writing logo to file");
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;

src/registry.rs

@@ -0,0 +1,128 @@
+use winreg::enums::*;
+use winreg::RegKey;
+
+pub trait RegistryHandlerTrait {
+    fn create_subkey(&self, path: &str) -> Result<(), Box<dyn std::error::Error>>;
+    fn set_value(
+        &self,
+        path: &str,
+        key: &str,
+        value: &str,
+    ) -> Result<(), Box<dyn std::error::Error>>;
+}
+
+pub struct RegistryHandler {
+    hkcu: RegKey,
+}
+
+impl RegistryHandler {
+    pub fn new(hkcu: RegKey) -> Self {
+        Self { hkcu }
+    }
+}
+
+impl RegistryHandlerTrait for RegistryHandler {
+    fn create_subkey(&self, path: &str) -> Result<(), Box<dyn std::error::Error>> {
+        self.hkcu.create_subkey(path)?;
+        Ok(())
+    }
+
+    fn set_value(
+        &self,
+        path: &str,
+        key: &str,
+        value: &str,
+    ) -> Result<(), Box<dyn std::error::Error>> {
+        let subkey = self.hkcu.open_subkey_with_flags(path, KEY_WRITE)?;
+        subkey.set_value(key, &value)?;
+        Ok(())
+    }
+}
+
+pub fn registry(
+    logo_path: &str,
+    app_name: &str,
+    handler: &impl RegistryHandlerTrait,
+) -> Result<(), Box<dyn std::error::Error>> {
+    // Define variables for the registry entry
+    let aumid = app_name;
+    let display_name = app_name;
+    let icon_uri = logo_path;
+    let path = format!("Software\\Classes\\AppUserModelId\\{}", aumid);
+
+    // Open the registry key (or create it if it doesn't exist)
+    handler
+        .create_subkey(&path)
+        .expect("Could not write to registry");
+
+    // Set registry values
+    handler
+        .set_value(&path, "DisplayName", display_name)
+        .expect("Could not write DisplayName to registry");
+    handler
+        .set_value(&path, "IconUri", icon_uri)
+        .expect("Could not write IconUri to registry");
+
+    Ok(())
+}
+
+#[allow(dead_code)]
+mod tests {
+    use std::collections::HashMap;
+    use std::sync::Mutex;
+
+    use super::*;
+
+    struct MockRegistryHandler {
+        data: Mutex<HashMap<String, HashMap<String, String>>>,
+    }
+
+    impl MockRegistryHandler {
+        pub fn new() -> Self {
+            Self {
+                data: Mutex::new(HashMap::new()),
+            }
+        }
+    }
+
+    impl RegistryHandlerTrait for MockRegistryHandler {
+        fn create_subkey(&self, path: &str) -> Result<(), Box<dyn std::error::Error>> {
+            let mut data = self.data.lock().unwrap();
+            data.entry(path.to_string()).or_default();
+            Ok(())
+        }
+
+        fn set_value(
+            &self,
+            path: &str,
+            key: &str,
+            value: &str,
+        ) -> Result<(), Box<dyn std::error::Error>> {
+            let mut data = self.data.lock().unwrap();
+            //let mut data = data.entry(path)
+            let map = data.get_mut(path).unwrap();
+
+            map.insert(key.to_string(), value.to_string());
+            Ok(())
+        }
+    }
+
+    #[test]
+    fn test_registry_with_mock() {
+        let mock_handler = MockRegistryHandler::new();
+
+        let logo_path = "C:\\Path\\To\\Logo.ico";
+        let app_name = "TestApp";
+
+        let result = registry(logo_path, app_name, &mock_handler);
+        assert!(result.is_ok());
+
+        let data = mock_handler.data.lock().unwrap();
+        let subkey = data.get("Software\\Classes\\AppUserModelId\\TestApp");
+        assert!(subkey.is_some());
+
+        let values = subkey.unwrap();
+        assert_eq!(values.get("DisplayName"), Some(&app_name.to_string()));
+        assert_eq!(values.get("IconUri"), Some(&logo_path.to_string()));
+    }
+}