Improve The Security Model

swhkd/src/daemon.rs

@@ -2,22 +2,19 @@ use crate::config::Value;
 use clap::Parser;
 use config::Hotkey;
 use evdev::{AttributeSet, Device, InputEventKind, Key};
-use nix::{
-    sys::stat::{umask, Mode},
-    unistd::{Group, Uid},
-};
+use nix::sys::stat::{umask, Mode};
 use signal_hook::consts::signal::*;
 use signal_hook_tokio::Signals;
 use std::{
     collections::{HashMap, HashSet},
     env,
     error::Error,
-    fs,
-    fs::Permissions,
-    io::prelude::*,
-    os::unix::{fs::PermissionsExt, net::UnixStream},
+    fs::{self, OpenOptions, Permissions},
+    io::Read,
+    os::unix::{fs::PermissionsExt, net::UnixListener},
     path::{Path, PathBuf},
-    process::{exit, id},
+    process::{exit, id, Command, Stdio},
+    time::{SystemTime, UNIX_EPOCH},
 };
 use sysinfo::{ProcessExt, System, SystemExt};
 use tokio::select;
@@ -31,9 +28,6 @@ mod environ;
 mod perms;
 mod uinput;
 
-#[cfg(test)]
-mod tests;
-
 struct KeyboardState {
     state_modifiers: HashSet<config::Modifier>,
     state_keysyms: AttributeSet<evdev::Key>,
@@ -61,9 +55,17 @@ struct Args {
     #[arg(short, long)]
     debug: bool,
 
+    /// Set Server Refresh Time in milliseconds
+    #[arg(short, long)]
+    refresh: Option<u64>,
+
     /// Take a list of devices from the user
     #[arg(short = 'D', long, num_args = 0.., value_delimiter = ' ')]
     device: Vec<String>,
+
+    /// Set a custom log file. (Defaults to ${XDG_DATA_HOME:-$HOME/.local/share}/swhks-current_unix_time.log)
+    #[arg(short, long, value_name = "FILE")]
+    log: Option<PathBuf>,
 }
 
 #[tokio::main]
@@ -78,36 +80,58 @@ async fn main() -> Result<(), Box<dyn Error>> {
 
     env_logger::init();
     log::trace!("Logger initialized.");
+    perms::raise_privileges();
 
-    let env = environ::Env::construct();
+    let invoking_uid = get_uid()?;
+    let uname = get_uname_from_uid(invoking_uid)?;
+
+    let mut env = refresh_env(&uname, invoking_uid).unwrap();
     log::trace!("Environment Aquired");
 
-    let invoking_uid = env.pkexec_id;
+    let log_file_name = if let Some(val) = args.log {
+        val
+    } else {
+        let time = match SystemTime::now().duration_since(UNIX_EPOCH) {
+            Ok(n) => n.as_secs().to_string(),
+            Err(_) => {
+                log::error!("SystemTime before UnixEpoch!");
+                exit(1);
+            }
+        };
 
-    setup_swhkd(invoking_uid, env.xdg_runtime_dir.clone().to_string_lossy().to_string());
+        format!("{}/swhkd/swhkd-{}.log", env.fetch_xdg_data_path().to_string_lossy(), time).into()
+    };
 
-    let load_config = || {
-        // Drop privileges to the invoking user.
-        perms::drop_privileges(invoking_uid);
+    let log_path = Path::new(&log_file_name);
+    if let Some(p) = log_path.parent() {
+        if !p.exists() {
+            if let Err(e) = fs::create_dir_all(p) {
+                log::error!("Failed to create log dir: {}", e);
+            }
+        }
+    }
 
-        let config_file_path: PathBuf =
-            args.config.as_ref().map_or_else(|| env.fetch_xdg_config_path(), |file| file.clone());
+    setup_swhkd(invoking_uid, env.xdg_runtime_dir(invoking_uid));
 
+    let config_file_path: PathBuf =
+        args.config.as_ref().map_or_else(|| env.fetch_xdg_config_path(), |file| file.clone());
+    let load_config = || {
         log::debug!("Using config file path: {:#?}", config_file_path);
 
         match config::load(&config_file_path) {
             Err(e) => {
                 log::error!("Config Error: {}", e);
                 exit(1)
             }
-            Ok(out) => {
-                // Escalate back to the root user after reading the config file.
-                perms::raise_privileges();
-                out
-            }
+            Ok(out) => out,
         }
     };
 
+    // The server cool down is set to 650ms by default
+    // which is calculated based on the default repeat cooldown
+    // along with it, an additional 120ms is added to it, just to be safe.
+    let server_cooldown = args.refresh.unwrap_or(650 + 120);
+
     let mut modes = load_config();
     let mut mode_stack: Vec<usize> = vec![0];
     let arg_devices: Vec<String> = args.device;
@@ -167,8 +191,8 @@ async fn main() -> Result<(), Box<dyn Error>> {
     let repeat_cooldown_duration: u64 = args.cooldown.unwrap_or(default_cooldown);
 
     let mut signals = Signals::new([
-        SIGUSR1, SIGUSR2, SIGHUP, SIGABRT, SIGBUS, SIGCHLD, SIGCONT, SIGINT, SIGPIPE, SIGQUIT,
-        SIGSYS, SIGTERM, SIGTRAP, SIGTSTP, SIGVTALRM, SIGXCPU, SIGXFSZ,
+        SIGUSR1, SIGUSR2, SIGHUP, SIGABRT, SIGBUS, SIGCONT, SIGINT, SIGPIPE, SIGQUIT, SIGSYS,
+        SIGTERM, SIGTRAP, SIGTSTP, SIGVTALRM, SIGXCPU, SIGXFSZ,
     ])?;
 
     let mut execution_is_paused = false;
@@ -191,21 +215,28 @@ async fn main() -> Result<(), Box<dyn Error>> {
 
     // The initial sleep duration is never read because last_hotkey is initialized to None
     let hotkey_repeat_timer = sleep(Duration::from_millis(0));
+    let next_env = sleep(Duration::from_millis(server_cooldown));
     tokio::pin!(hotkey_repeat_timer);
+    tokio::pin!(next_env);
 
-    // The socket we're sending the commands to.
-    let socket_file_path = env.fetch_xdg_runtime_socket_path();
     loop {
         select! {
             _ = &mut hotkey_repeat_timer, if &last_hotkey.is_some() => {
                 let hotkey = last_hotkey.clone().unwrap();
                 if hotkey.keybinding.on_release {
                     continue;
                 }
-                send_command(hotkey.clone(), &socket_file_path, &modes, &mut mode_stack);
+                send_command(hotkey.clone(), &modes, &mut mode_stack, &uname, &env, log_path);
                 hotkey_repeat_timer.as_mut().reset(Instant::now() + Duration::from_millis(repeat_cooldown_duration));
             }
 
+            // Refresh the environment
+            _ = &mut next_env => {
+                log::info!("Refreshing Env");
+                env = refresh_env(&uname, invoking_uid).unwrap();
+                next_env.as_mut().reset(Instant::now() + Duration::from_secs(server_cooldown));
+            }
+
             Some(signal) = signals.next() => {
                 match signal {
                     SIGUSR1 => {
@@ -322,7 +353,7 @@ async fn main() -> Result<(), Box<dyn Error>> {
                     0 => {
                         if last_hotkey.is_some() && pending_release {
                             pending_release = false;
-                            send_command(last_hotkey.clone().unwrap(), &socket_file_path, &modes, &mut mode_stack);
+                            send_command(last_hotkey.clone().unwrap(), &modes, &mut mode_stack, &uname, &env, log_path);
                             last_hotkey = None;
                         }
                         if let Some(modifier) = modifiers_map.get(&key) {
@@ -385,7 +416,7 @@ async fn main() -> Result<(), Box<dyn Error>> {
                             pending_release = true;
                             break;
                         }
-                        send_command(hotkey.clone(), &socket_file_path, &modes, &mut mode_stack);
+                        send_command(hotkey.clone(), &modes, &mut mode_stack, &uname, &env, log_path);
                         hotkey_repeat_timer.as_mut().reset(Instant::now() + Duration::from_millis(repeat_cooldown_duration));
                         continue;
                     }
@@ -395,32 +426,6 @@ async fn main() -> Result<(), Box<dyn Error>> {
     }
 }
 
-fn socket_write(command: &str, socket_path: PathBuf) -> Result<(), Box<dyn Error>> {
-    let mut stream = UnixStream::connect(socket_path)?;
-    stream.write_all(command.as_bytes())?;
-    Ok(())
-}
-
-pub fn check_input_group() -> Result<(), Box<dyn Error>> {
-    if !Uid::current().is_root() {
-        let groups = nix::unistd::getgroups();
-        for groups in groups.iter() {
-            for group in groups {
-                let group = Group::from_gid(*group);
-                if group.unwrap().unwrap().name == "input" {
-                    log::error!("Note: INVOKING USER IS IN INPUT GROUP!!!!");
-                    log::error!("THIS IS A HUGE SECURITY RISK!!!!");
-                }
-            }
-        }
-        log::error!("Consider using `pkexec swhkd ...`");
-        exit(1);
-    } else {
-        log::warn!("Running swhkd as root!");
-        Ok(())
-    }
-}
-
 pub fn check_device_is_keyboard(device: &Device) -> bool {
     if device.supported_keys().map_or(false, |keys| keys.contains(Key::KEY_ENTER)) {
         if device.name() == Some("swhkd virtual output") {
@@ -434,7 +439,7 @@ pub fn check_device_is_keyboard(device: &Device) -> bool {
     }
 }
 
-pub fn setup_swhkd(invoking_uid: u32, runtime_path: String) {
+pub fn setup_swhkd(invoking_uid: u32, runtime_path: PathBuf) {
     // Set a sane process umask.
     log::trace!("Setting process umask.");
     umask(Mode::S_IWGRP | Mode::S_IWOTH);
@@ -454,7 +459,7 @@ pub fn setup_swhkd(invoking_uid: u32, runtime_path: String) {
     }
 
     // Get the PID file path for instance tracking.
-    let pidfile: String = format!("{}swhkd_{}.pid", runtime_path, invoking_uid);
+    let pidfile: String = format!("{}/swhkd_{}.pid", runtime_path.to_string_lossy(), invoking_uid);
     if Path::new(&pidfile).exists() {
         log::trace!("Reading {} file and checking for running instances.", pidfile);
         let swhkd_pid = match fs::read_to_string(&pidfile) {
@@ -487,18 +492,15 @@ pub fn setup_swhkd(invoking_uid: u32, runtime_path: String) {
             exit(1);
         }
     }
-
-    // Check if the user is in input group.
-    if check_input_group().is_err() {
-        exit(1);
-    }
 }
 
 pub fn send_command(
     hotkey: Hotkey,
-    socket_path: &Path,
     modes: &[config::Mode],
     mode_stack: &mut Vec<usize>,
+    uname: &str,
+    env: &environ::Env,
+    log_path: &Path,
 ) {
     log::info!("Hotkey pressed: {:#?}", hotkey);
     let command = hotkey.command;
@@ -533,9 +535,106 @@ pub fn send_command(
     if commands_to_send.ends_with(" &&") {
         commands_to_send = commands_to_send.strip_suffix(" &&").unwrap().to_string();
     }
-    if let Err(e) = socket_write(&commands_to_send, socket_path.to_path_buf()) {
-        log::error!("Failed to send command to swhks through IPC.");
-        log::error!("Please make sure that swhks is running.");
-        log::error!("Err: {:#?}", e)
-    };
+
+    launch(&commands_to_send, uname, env, log_path);
+}
+
+/// Launch Commands
+fn launch(command: &str, uname: &str, env: &environ::Env, log_path: &Path) {
+    let mut cmd = Command::new("su");
+    cmd.arg(uname)
+        .arg("-c")
+        .arg("-l")
+        .arg(command)
+        .stdin(Stdio::null())
+        .stdout(match OpenOptions::new().append(true).create(true).open(log_path) {
+            Ok(file) => file,
+            Err(e) => {
+                _ = Command::new("notify-send").arg(format!("ERROR {}", e)).spawn();
+                exit(1);
+            }
+        })
+        .stderr(match OpenOptions::new().append(true).create(true).open(log_path) {
+            Ok(file) => file,
+            Err(e) => {
+                _ = Command::new("notify-send").arg(format!("ERROR {}", e)).spawn();
+                exit(1);
+            }
+        });
+
+    for (key, value) in &env.pairs {
+        cmd.env(key, value);
+    }
+
+    match cmd.spawn() {
+        Ok(_) => log::info!("Command executed successfully."),
+        Err(e) => log::error!("Failed to execute command: {}", e),
+    }
+}
+
+/// Get the UID of the user that is not a system user
+fn get_uid() -> Result<u32, Box<dyn Error>> {
+    let status_content = fs::read_to_string(format!("/proc/{}/loginuid", std::process::id()))?;
+    let uid = status_content.trim().parse::<u32>()?;
+    Ok(uid)
+}
+
+fn get_uname_from_uid(uid: u32) -> Result<String, Box<dyn Error>> {
+    let passwd = fs::read_to_string("/etc/passwd").unwrap();
+    let lines: Vec<&str> = passwd.split('\n').collect();
+    for line in lines {
+        let parts: Vec<&str> = line.split(':').collect();
+        if parts.len() > 2 {
+            let Ok(user_id) = parts[2].parse::<u32>() else {
+                continue;
+            };
+            if user_id == uid {
+                return Ok(parts[0].to_string());
+            }
+        }
+    }
+    Err("User not found".into())
+}
+
+fn get_file_paths(runtime_dir: &str) -> (String, String) {
+    let pid_file_path = format!("{}/swhks.pid", runtime_dir);
+    let sock_file_path = format!("{}/swhkd.sock", runtime_dir);
+
+    (pid_file_path, sock_file_path)
+}
+
+fn refresh_env(uname: &str, invoking_uid: u32) -> Result<environ::Env, Box<dyn Error>> {
+    let env = environ::Env::construct(uname, None);
+
+    let (_pid_path, sock_path) =
+        get_file_paths(env.xdg_runtime_dir(invoking_uid).to_str().unwrap());
+
+    if Path::new(&sock_path).exists() {
+        fs::remove_file(&sock_path)?;
+    }
+
+    let mut result: String = String::new();
+    let listener = UnixListener::bind(&sock_path)?;
+    fs::set_permissions(sock_path, fs::Permissions::from_mode(0o666))?;
+    // TODO: Implement Wait mechanism
+    loop {
+        log::warn!("Waiting for Server...");
+        match listener.accept() {
+            Ok((mut socket, _addr)) => {
+                let mut buf = String::new();
+                socket.read_to_string(&mut buf)?;
+                if buf.is_empty() {
+                    continue;
+                }
+                log::info!("Server Instance found!");
+                result.push_str(&buf);
+                break;
+            }
+            Err(e) => {
+                log::info!("Sock Err: {}", e);
+            }
+        }
+    }
+    log::trace!("Environment Refreshed");
+    Ok(environ::Env::construct(uname, Some(&result)))
 }