sd-switch 0.6.3

use crate::systemd::UnitManager;
use crate::systemd::UnitStatus;

use crate::error::Error;
use std::collections::HashMap;
use std::process;
use std::str::FromStr;
use std::sync::mpsc;
use std::thread;
use std::{result::Result, time::Duration};

use super::SystemStatus;

pub struct SystemctlServiceManager {
    system: bool,
}

pub struct SystemctlUnitManager {
    status: SystemctlUnitStatus,
}

#[derive(Clone, Debug, PartialEq)]
pub struct SystemctlUnitStatus {
    name: String,
    description: String,
    active_state: String,
    address: String,
    refuse_manual_start: bool,
    refuse_manual_stop: bool,
}

impl UnitStatus for SystemctlUnitStatus {
    fn name(&self) -> &str {
        &self.name
    }

    fn description(&self) -> &str {
        &self.description
    }

    fn active_state(&self) -> &str {
        &self.active_state
    }
}

struct Job {
    unit_name: String,
    child: process::Child,
}

pub struct SystemctlJobSet<'a> {
    manager: &'a SystemctlServiceManager,
    jobs: Vec<Job>,
}

impl SystemctlJobSet<'_> {
    fn new(manager: &SystemctlServiceManager) -> SystemctlJobSet<'_> {
        SystemctlJobSet {
            manager,
            jobs: Vec::new(),
        }
    }
}

impl super::JobSet for SystemctlJobSet<'_> {
    fn reload_unit(&mut self, unit_name: &str) -> Result<(), Error> {
        let child = self
            .manager
            .command()
            .arg("--job-mode=replace")
            .arg("reload")
            .arg(unit_name)
            .spawn()?;
        self.jobs.push(Job {
            unit_name: unit_name.to_string(),
            child,
        });
        Ok(())
    }

    fn restart_unit(&mut self, unit_name: &str) -> Result<(), Error> {
        let child = self
            .manager
            .command()
            .arg("--job-mode=replace")
            .arg("restart")
            .arg(unit_name)
            .spawn()?;
        self.jobs.push(Job {
            unit_name: unit_name.to_string(),
            child,
        });
        Ok(())
    }

    fn start_unit(&mut self, unit_name: &str) -> Result<(), Error> {
        let child = self
            .manager
            .command()
            .arg("--job-mode=replace")
            .arg("start")
            .arg(unit_name)
            .spawn()?;
        self.jobs.push(Job {
            unit_name: unit_name.to_string(),
            child,
        });
        Ok(())
    }

    fn stop_unit(&mut self, unit_name: &str) -> Result<(), Error> {
        let child = self
            .manager
            .command()
            .arg("--job-mode=replace")
            .arg("stop")
            .arg(unit_name)
            .spawn()?;
        self.jobs.push(Job {
            unit_name: unit_name.to_string(),
            child,
        });
        Ok(())
    }

    fn wait_for_all<F>(&mut self, job_handler: F, timeout: Duration) -> Result<(), Error>
    where
        F: Fn(&str, &str) + Send + 'static,
    {
        if self.jobs.is_empty() {
            return Ok(());
        }

        let (tx, rx) = mpsc::channel();

        let jobs: Vec<Job> = self.jobs.drain(..).collect();
        let _ = thread::spawn(move || {
            for mut job in jobs {
                let result = if job.child.wait().is_ok_and(|r| r.success()) {
                    "done"
                } else {
                    "failed"
                };
                job_handler(&job.unit_name, result);
            }

            let _ = tx.send(());
        });

        rx.recv_timeout(timeout)
            .map_err(|e| Error::SdSwitch(e.to_string()))
    }
}

impl SystemctlServiceManager {
    /// Creates a new systemctl service manager instance. The given Boolean
    /// indicates whether we should connect to the system or user service
    /// manager.
    pub fn new(system: bool) -> Result<SystemctlServiceManager, Error> {
        Ok(SystemctlServiceManager { system })
    }

    /// Start a systemctl command call, connecting to the system or user
    /// manager.
    fn command(&self) -> process::Command {
        let mut cmd = process::Command::new("systemctl");

        cmd.arg(if self.system { "--system" } else { "--user" });

        cmd
    }
}

impl<'a> super::ServiceManager for &'a SystemctlServiceManager {
    type UnitManager = SystemctlUnitManager;
    type UnitStatus = SystemctlUnitStatus;
    type JobSet = SystemctlJobSet<'a>;

    fn system_status(&self) -> Result<SystemStatus, Error> {
        let result = self.command().arg("is-system-running").output()?;

        let output = String::from_utf8_lossy(&result.stdout);

        SystemStatus::from_str(output.trim_end())
    }

    /// Performs a systemd daemon reload, blocking until complete.
    fn daemon_reload(&self) -> Result<(), Error> {
        let result = self.command().arg("daemon-reload").status()?;

        if result.success() {
            Ok(())
        } else {
            Err(Error::SdSwitch(String::from(
                "Error performing daemon reload",
            )))
        }
    }

    fn reset_failed(&self) -> Result<(), Error> {
        let result = self.command().arg("reset-failed").status()?;

        if result.success() {
            Ok(())
        } else {
            Err(Error::SdSwitch(String::from(
                "Error resetting failed units",
            )))
        }
    }

    /// Builds a unit manager for the unit with the given status.
    fn unit_manager(&self, status: &SystemctlUnitStatus) -> Result<SystemctlUnitManager, Error> {
        Ok(SystemctlUnitManager {
            status: status.clone(),
        })
    }

    fn new_job_set(&self) -> Result<SystemctlJobSet<'a>, Error> {
        Ok(SystemctlJobSet::new(self))
    }

    fn list_units_by_states(&self, states: &[&str]) -> Result<Vec<SystemctlUnitStatus>, Error> {
        let mut command = self.command();

        command.arg("show").arg("*");

        if !states.is_empty() {
            command.arg("--state").arg(states.join(","));
        }

        let output = command.output()?;
        let output = String::from_utf8_lossy(&output.stdout);

        read_show_units(&output)
    }
}

/// Reads the output of `systemctl show` for one unit. Stops at empty line of EOF.
fn read_show_unit(kvs: &str) -> Result<SystemctlUnitStatus, Error> {
    let new_error =
        |msg| move || Error::SdSwitch(format!("Unexpected output from systemctl show: {msg}"));
    let missing_field = |field| move || new_error(format!("Missing '{field}' field"))();
    let unit = kvs
        .lines()
        .map(|line| {
            line.split_once('=')
                .ok_or_else(new_error(format!("Invalid unit line: {line}")))
        })
        .collect::<Result<HashMap<&str, &str>, Error>>()?;

    let name = unit.get("Id").ok_or_else(missing_field("Id"))?;
    let description = unit
        .get("Description")
        .ok_or_else(missing_field("Description"))?;
    let active_state = unit
        .get("ActiveState")
        .ok_or_else(missing_field("ActiveState"))?;
    let address = unit.get("Id").ok_or_else(missing_field("Id"))?;
    let refuse_manual_start = unit
        .get("RefuseManualStart")
        .ok_or_else(missing_field("RefuseManualStart"))?;
    let refuse_manual_stop = unit
        .get("RefuseManualStop")
        .ok_or_else(missing_field("RefuseManualStop"))?;

    Ok(SystemctlUnitStatus {
        name: (*name).to_string(),
        description: (*description).to_string(),
        active_state: (*active_state).to_string(),
        address: (*address).to_string(),
        refuse_manual_start: *refuse_manual_start == "yes",
        refuse_manual_stop: *refuse_manual_stop == "yes",
    })
}

/// Reads the complete output of `systemctl show`, keeping only the ones with the given states.
fn read_show_units(handle: &str) -> Result<Vec<SystemctlUnitStatus>, Error> {
    handle.split("\n\n").map(read_show_unit).collect()
}

impl UnitManager for SystemctlUnitManager {
    fn refuse_manual_start(&self) -> Result<bool, Error> {
        Ok(self.status.refuse_manual_start)
    }

    fn refuse_manual_stop(&self) -> Result<bool, Error> {
        Ok(self.status.refuse_manual_stop)
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn can_read_show_units() -> Result<(), Error> {
        let raw = r"Id=service1.service
Description=Service 1
ActiveState=active
Type=simple
RefuseManualStart=yes
RefuseManualStop=no

Id=service2.service
Description=Service 2
ActiveState=active
Type=simple
RefuseManualStart=no
RefuseManualStop=yes
";

        let statuses = read_show_units(raw)?;

        assert_eq!(
            statuses,
            vec![
                SystemctlUnitStatus {
                    name: "service1.service".to_string(),
                    description: "Service 1".to_string(),
                    active_state: "active".to_string(),
                    address: "service1.service".to_string(),
                    refuse_manual_start: true,
                    refuse_manual_stop: false
                },
                SystemctlUnitStatus {
                    name: "service2.service".to_string(),
                    description: "Service 2".to_string(),
                    active_state: "active".to_string(),
                    address: "service2.service".to_string(),
                    refuse_manual_start: false,
                    refuse_manual_stop: true
                }
            ]
        );

        Ok(())
    }
}