malarkey/src/main.rs

use chrono::{Date, Local, TimeZone};
use std::collections::BTreeMap;
use std::path::{Path, PathBuf};
use sysinfo::{ProcessExt, SystemExt};
use winreg::{RegValue, types::FromRegValue};
use winrt_notification::{Duration, Sound, Toast};

fn make_holidays() -> BTreeMap<Date<Local>, &'static str> {
    let mut holidays = BTreeMap::new();

    holidays.insert(Local.ymd(2021, 1, 1), "New Year's Day");
    holidays.insert(Local.ymd(2021, 2, 12), "Chinese New Year");
    holidays.insert(Local.ymd(2021, 2, 13), "Chinese New Year");
    holidays.insert(Local.ymd(2021, 4, 2), "Good Friday");
    holidays.insert(Local.ymd(2021, 5, 1), "Labour Day");
    holidays.insert(Local.ymd(2021, 5, 13), "Hari Raya Puasa");
    holidays.insert(Local.ymd(2021, 7, 20), "Vesak Day");
    holidays.insert(Local.ymd(2021, 8, 9), "National Day");
    holidays.insert(Local.ymd(2021, 11, 4), "Deepavali");
    holidays.insert(Local.ymd(2021, 12, 25), "Christmas");

    holidays.insert(Local.ymd(2022, 1, 1), "New Year's Day");
    holidays.insert(Local.ymd(2022, 2, 1), "Chinese New Year");
    holidays.insert(Local.ymd(2022, 2, 2), "Chinese New Year");
    holidays.insert(Local.ymd(2022, 4, 15), "Good Friday");
    holidays.insert(Local.ymd(2022, 5, 1), "Labour Day");
    holidays.insert(Local.ymd(2022, 5, 2), "Hari Raya Puasa");
    holidays.insert(Local.ymd(2022, 5, 3), "Labour Day (Observed)");
    holidays.insert(Local.ymd(2022, 5, 15), "Vesak Day");
    holidays.insert(Local.ymd(2022, 7, 9), "Hari Raya Haji");
    holidays.insert(Local.ymd(2022, 8, 9), "National Day");
    holidays.insert(Local.ymd(2022, 10, 24), "Deepavali");
    holidays.insert(Local.ymd(2022, 12, 25), "Christmas");
    holidays
}

fn process_is_running(name: &str) -> bool {
    // First we update all information of our system struct.
    let mut system = sysinfo::System::new_all();
    system.refresh_all();
    for (_pid, proc_) in system.get_processes() {
        if proc_.name() == name {
            return true;
        }
    }
    false
}

fn try_terminating(name: &str) -> Option<std::path::PathBuf> {
    let mut process_path = None;

    // First we update all information of our system struct.
    let mut system = sysinfo::System::new_all();
    system.refresh_all();

    // Now let's print every process' id and name:
    for (pid, proc_) in system.get_processes() {
        // println!("{}:{} => status: {:?}", pid, proc_.name(), proc_.status());
        if proc_.name() == name {
            println!("LARK FOUND!!!  PID: {}, status: {:?}", pid, proc_.status());
            println!("Original exe: {}", proc_.exe().display());
            process_path = Some(proc_.exe().to_owned());
            proc_.kill(sysinfo::Signal::Term);
        }
    }
    process_path
}

fn is_working_hours<T: chrono::Datelike + chrono::Timelike>(date_time: &T) -> bool {
    let holidays = make_holidays();

    if let Some(holiday) = holidays.get(&chrono::Local::today()) {
        println!("Happy {}! No work today.", holiday);
        return false;
    }

    match date_time.weekday() {
        chrono::Weekday::Sun | chrono::Weekday::Sat => return false,
        _ => (),
    }

    date_time.hour() >= 9 && date_time.hour() <= 18
}

fn get_lark_path() -> Option<PathBuf> {
    use winreg::enums::*;
    use winreg::RegKey;
    let hkcu = RegKey::predef(HKEY_CURRENT_USER);
    let lark_cfg = match hkcu.open_subkey("SOFTWARE\\Ion\\Larkinator") {
        Ok(v) => Some(v),
        Err(_) => None,
    }?;
    let v: std::io::Result<RegValue> = lark_cfg.get_raw_value("Lark Path");
    let mut buf = v.unwrap();

    // Ensure the buffer is big enough.
    if buf.bytes.len() <= 2 {
        return None;
    }

    // Sometimes it adds a NULL at the end. Strip off NULL termination.
    if buf.bytes[buf.bytes.len()-2] == b'\0' {
        buf.bytes.pop();
        buf.bytes.pop();
    }

    let osstr = std::ffi::OsString::from_reg_value(&buf).unwrap();
    Some(Path::new(&osstr).to_owned())
}

fn set_lark_path(new_path: &Path) -> std::io::Result<()> {
    use winreg::enums::*;
    use winreg::RegKey;
    let hkcu = RegKey::predef(HKEY_CURRENT_USER);
    let (key, _disp) = hkcu.create_subkey("SOFTWARE\\Ion\\Larkinator")?;
    key.set_value("Lark Path", &new_path.as_os_str())?;
    Ok(())
}

fn main() {
    let process_name = "Lark.exe";

    loop {
        let local = Local::now(); // e.g. `2014-11-28T21:45:59.324310806+09:00`
        let should_be_running = is_working_hours(&local);
        let is_running = process_is_running(process_name);
        // println!("Should? {}  Running? {}", should_be_running, is_running);
        if should_be_running && !is_running {
            println!("Running Lark!");
            if let Some(p) = get_lark_path() {
                println!("Lark path: {}", p.display());

                std::process::Command::new(p)
                    .spawn()
                    .expect("Couldn't run Lark!");

                Toast::new(Toast::POWERSHELL_APP_ID)
                    .title("Started Lark")
                    .text1("Running Lark because it's time to work")
                    .sound(Some(Sound::SMS))
                    .duration(Duration::Short)
                    .show()
                    .expect("unable to toast");
            } else {
                println!("Would run Lark, but no path found");
            }
        } else if !should_be_running && is_running {
            println!("Terminating Lark, since it should't be running");
            if let Some(p) = try_terminating(process_name) {
                set_lark_path(&p).ok();
                Toast::new(Toast::POWERSHELL_APP_ID)
                    .title("Terminated Lark")
                    .text1("Lark was terminated because it is outside working hours")
                    .sound(Some(Sound::SMS))
                    .duration(Duration::Short)
                    .show()
                    .expect("unable to toast");
            }
        }
        std::thread::sleep(std::time::Duration::from_secs(3));
    }
}