From eb5fbc8e39f8f1f9a90128d85ec815349555c11d Mon Sep 17 00:00:00 2001
From: Vicki Pfau <vi@endrift.com>
Date: Wed, 24 Apr 2024 18:06:31 -0700
Subject: [PATCH] systemd: Add methods for enabling, disabling, masking, and
 unmasking

---
 src/hardware.rs |  4 +--
 src/systemd.rs  | 80 ++++++++++++++++++++++++++++++++++++++++++++++++-
 src/wifi.rs     |  2 +-
 3 files changed, 82 insertions(+), 4 deletions(-)

diff --git a/src/hardware.rs b/src/hardware.rs
index 948bb9f..03a0a24 100644
--- a/src/hardware.rs
+++ b/src/hardware.rs
@@ -126,7 +126,7 @@ impl FanControl {
 
     pub async fn get_state(&self) -> Result<FanControlState> {
         let jupiter_fan_control =
-            SystemdUnit::new(self.connection.clone(), "jupiter_2dfan_2dcontrol_2eservice").await?;
+            SystemdUnit::new(self.connection.clone(), "jupiter-fan-control.service").await?;
         let active = jupiter_fan_control.active().await?;
         Ok(match active {
             true => FanControlState::Os,
@@ -137,7 +137,7 @@ impl FanControl {
     pub async fn set_state(&self, state: FanControlState) -> Result<()> {
         // Run what steamos-polkit-helpers/jupiter-fan-control does
         let jupiter_fan_control =
-            SystemdUnit::new(self.connection.clone(), "jupiter_2dfan_2dcontrol_2eservice").await?;
+            SystemdUnit::new(self.connection.clone(), "jupiter-fan-control.service").await?;
         match state {
             FanControlState::Os => jupiter_fan_control.start().await,
             FanControlState::Bios => jupiter_fan_control.stop().await,
diff --git a/src/systemd.rs b/src/systemd.rs
index 8f9f3e6..481828a 100644
--- a/src/systemd.rs
+++ b/src/systemd.rs
@@ -4,6 +4,7 @@
  *
  * SPDX-License-Identifier: MIT
  */
+#![allow(dead_code)]
 
 use anyhow::{anyhow, Result};
 use std::path::PathBuf;
@@ -29,11 +30,38 @@ trait SystemdUnit {
     default_path = "/org/freedesktop/systemd1"
 )]
 trait SystemdManager {
+    async fn enable_unit_files(
+        &self,
+        files: &[&str],
+        runtime: bool,
+        force: bool,
+    ) -> Result<(bool, Vec<(String, String, String)>)>;
+
+    async fn disable_unit_files(
+        &self,
+        files: &[&str],
+        runtime: bool,
+    ) -> Result<Vec<(String, String, String)>>;
+
+    async fn mask_unit_files(
+        &self,
+        files: &[&str],
+        runtime: bool,
+    ) -> Result<Vec<(String, String, String)>>;
+
+    async fn unmask_unit_files(
+        &self,
+        files: &[&str],
+        runtime: bool,
+    ) -> Result<Vec<(String, String, String)>>;
+
     async fn reload(&self) -> Result<()>;
 }
 
 pub struct SystemdUnit<'dbus> {
+    connection: Connection,
     proxy: SystemdUnitProxy<'dbus>,
+    name: String,
 }
 
 pub async fn daemon_reload(connection: &Connection) -> Result<()> {
@@ -44,7 +72,7 @@ pub async fn daemon_reload(connection: &Connection) -> Result<()> {
 
 impl<'dbus> SystemdUnit<'dbus> {
     pub async fn new(connection: Connection, name: &str) -> Result<SystemdUnit<'dbus>> {
-        let path = PathBuf::from("/org/freedesktop/systemd1/unit").join(name);
+        let path = PathBuf::from("/org/freedesktop/systemd1/unit").join(escape(name));
         let path = String::from(path.to_str().ok_or(anyhow!("Unit name {name} invalid"))?);
         Ok(SystemdUnit {
             proxy: SystemdUnitProxy::builder(&connection)
@@ -52,6 +80,8 @@ impl<'dbus> SystemdUnit<'dbus> {
                 .path(path)?
                 .build()
                 .await?,
+            connection,
+            name: String::from(name),
         })
     }
 
@@ -70,7 +100,55 @@ impl<'dbus> SystemdUnit<'dbus> {
         Ok(())
     }
 
+    pub async fn enable(&self) -> Result<bool> {
+        let manager = SystemdManagerProxy::new(&self.connection).await?;
+        let (_, res) = manager.enable_unit_files(&[self.name.as_str()], false, false).await?;
+        Ok(res.len() > 0)
+    }
+
+    pub async fn disable(&self) -> Result<bool> {
+        let manager = SystemdManagerProxy::new(&self.connection).await?;
+        let res = manager.disable_unit_files(&[self.name.as_str()], false).await?;
+        Ok(res.len() > 0)
+    }
+
+    pub async fn mask(&self) -> Result<bool> {
+        let manager = SystemdManagerProxy::new(&self.connection).await?;
+        let res = manager.mask_unit_files(&[self.name.as_str()], false).await?;
+        Ok(res.len() > 0)
+    }
+
+    pub async fn unmask(&self) -> Result<bool> {
+        let manager = SystemdManagerProxy::new(&self.connection).await?;
+        let res = manager.unmask_unit_files(&[self.name.as_str()], false).await?;
+        Ok(res.len() > 0)
+    }
+
     pub async fn active(&self) -> Result<bool> {
         Ok(self.proxy.active_state().await? == "active")
     }
 }
+
+pub fn escape(name: &str) -> String {
+    let mut parts = String::new();
+    for c in name.chars() {
+        if c.is_ascii_alphanumeric() {
+            parts.push(c);
+        } else {
+            let escaped = format!("_{:02x}", u32::from(c));
+            parts.push_str(escaped.as_str());
+        }
+    }
+    parts
+}
+
+#[cfg(test)]
+mod test {
+    use super::*;
+
+    #[test]
+    fn test_escape() {
+        assert_eq!(escape("systemd"), "systemd");
+        assert_eq!(escape("system d"), "system_20d");
+    }
+}
diff --git a/src/wifi.rs b/src/wifi.rs
index 6a461ad..2daf3ae 100644
--- a/src/wifi.rs
+++ b/src/wifi.rs
@@ -151,7 +151,7 @@ async fn restart_iwd(connection: Connection) -> Result<()> {
         .inspect_err(|message| error!("restart_iwd: reload systemd got an error: {message}"))?;
 
     // worked, now restart iwd
-    let unit = SystemdUnit::new(connection, "iwd_2eservice").await?;
+    let unit = SystemdUnit::new(connection, "iwd.service").await?;
     unit.restart()
         .await
         .inspect_err(|message| error!("restart_iwd: restart unit got an error: {message}"))