sd-switch 0.6.3

use anyhow::Context;
use i18n_cli::MSG;
use nix::unistd::Uid;
use sd_switch::systemd::dbus::DbusServiceManager;
use sd_switch::systemd::ServiceManager;
use std::{
    env, fmt,
    path::{Path, PathBuf},
    process::exit,
    time::Duration,
};

mod i18n_cli;

const PKG_NAME: Option<&'static str> = option_env!("CARGO_PKG_NAME");
const VERSION: Option<&'static str> = option_env!("CARGO_PKG_VERSION");

const DBUS_SESSION_BUS_ADDRESS: &str = "DBUS_SESSION_BUS_ADDRESS";
const XDG_RUNTIME_DIR: &str = "XDG_RUNTIME_DIR";

#[derive(Debug)]
enum TargetManager {
    System,
    User,
}

#[derive(Debug)]
pub struct Options {
    /// The manager to which we should connect.
    target_manager: TargetManager,

    /// Whether to log, but not perform, actions.
    pub dry_run: bool,

    /// Whether to log all actions taken.
    pub verbose: bool,

    /// How long to wait for units to start, stop, reload, etc. Default timeout
    /// is 2 minutes.
    pub timeout: Duration,

    /// Whether to force use of systemctl, even if D-Bus is available.
    pub force_systemctl: bool,

    /// Path to the directory of old unit files, if any exists.
    pub old_dir: Option<PathBuf>,

    /// Path to the directory of new unit files.
    pub new_dir: PathBuf,
}

/// Parses the given command line arguments.
pub fn parse_args(args: &[String]) -> Options {
    let program = args[0].clone();

    let mut opts = getopts::Options::new();
    opts.optflag("", "system", &MSG.help_system())
        .optflag("", "user", &MSG.help_user())
        .optopt("", "new-units", &MSG.help_new_units(), "DIR")
        .optopt("", "old-units", &MSG.help_old_units(), "DIR")
        .optflag("n", "dry-run", &MSG.help_dry_run())
        .optopt("", "timeout", &MSG.help_timeout(), "MS")
        .optflag("", "force-systemctl", &MSG.help_force_systemctl())
        .optflag("v", "verbose", &MSG.help_verbose())
        .optflag("h", "help", &MSG.help_help())
        .optflag("V", "version", &MSG.help_version());

    let matches = match opts.parse(&args[1..]) {
        Ok(m) => m,
        Err(fail) => {
            use getopts::Fail::*;
            let err = match fail {
                ArgumentMissing(opt) => MSG.err_argument_missing(&opt),
                UnrecognizedOption(opt) => MSG.err_unrecognized_option(&opt),
                OptionMissing(opt) => MSG.err_option_missing(&opt),
                OptionDuplicated(opt) => MSG.err_option_duplicated(&opt),
                UnexpectedArgument(opt) => MSG.err_unexpected_argument(&opt),
            };
            eprintln!("{}\n{}", err, MSG.err_run_help(&program));
            exit(1);
        }
    };

    if matches.opt_present("help") {
        print_usage(&program, &opts);
        exit(0);
    }

    if matches.opt_present("version") {
        println!("{}", package_version());
        exit(0);
    }

    let target_manager = if matches.opt_present("system") && matches.opt_present("user") {
        eprintln!("{}", MSG.err_both_system_and_user());
        exit(1);
    } else if matches.opt_present("system") {
        TargetManager::System
    } else {
        TargetManager::User
    };

    let timeout = matches.opt_str("timeout").unwrap_or(String::from("120000"));
    let timeout = match timeout.parse() {
        Ok(millis) if millis > 0 => Duration::from_millis(millis),
        _ => {
            eprintln!("{}", MSG.err_invalid_timeout());
            exit(1);
        }
    };

    let Some(new_dir) = matches.opt_str("new-units").map(PathBuf::from) else {
        eprintln!(
            "{}\n{}",
            MSG.err_missing_required_option("new-units"),
            MSG.err_run_help(&program)
        );
        exit(1);
    };

    Options {
        target_manager,
        dry_run: matches.opt_present("dry-run"),
        verbose: matches.opt_present("verbose"),
        timeout,
        force_systemctl: matches.opt_present("force-systemctl"),
        old_dir: matches.opt_str("old-units").map(PathBuf::from),
        new_dir,
    }
}

fn package_version() -> String {
    format!(
        "{} {}",
        PKG_NAME.unwrap_or("[unknown package]"),
        VERSION.unwrap_or("[unknown version]")
    )
}

fn print_usage(program: &str, opts: &getopts::Options) {
    print!(
        "{}",
        opts.usage(&MSG.help_brief(package_version().as_str(), program))
    );
}

/// Open connection to the standard D-Bus address.
fn open_std_dbus_connection(
    target_manager: &TargetManager,
) -> anyhow::Result<zbus::blocking::Connection> {
    let connection = match target_manager {
        TargetManager::System => {
            zbus::blocking::Connection::system().with_context(|| MSG.err_open_system_bus())?
        }
        TargetManager::User => {
            zbus::blocking::Connection::session().with_context(|| MSG.err_open_user_bus())?
        }
    };

    let mngr = DbusServiceManager::new(&connection)?;

    if (&mngr).system_status().is_err() {
        anyhow::bail!(MSG.err_bus_failure());
    }

    drop(mngr);

    Ok(connection)
}

/// Open connection to alternative D-Bus session address. Unfortunately, in a
/// few cases the `DBUS_SESSION_BUS_ADDRESS` variable lies and systemd is
/// actually listening on `/run/user/$UID/bus`.
///
/// No alternate system D-Bus address is considered.
fn open_alt_dbus_connection(
    target_manager: &TargetManager,
) -> anyhow::Result<zbus::blocking::Connection> {
    if matches!(target_manager, TargetManager::System) {
        anyhow::bail!("No alternate D-Bus system manager considered");
    }

    let mut bus_path =
        env::var(XDG_RUNTIME_DIR).unwrap_or_else(|_| format!("/run/user/{}", Uid::effective()));
    bus_path.push_str("/bus");

    let orig_dbus_address = env::var(DBUS_SESSION_BUS_ADDRESS)?;

    if Path::new(&bus_path).exists() {
        env::set_var(DBUS_SESSION_BUS_ADDRESS, format!("unix:path={bus_path}"));
    } else {
        anyhow::bail!(MSG.err_no_bus_socket(&bus_path));
    }

    let connection = zbus::blocking::Connection::session()
        .with_context(|| MSG.err_open_user_bus_alt(&bus_path))?;
    let mngr = DbusServiceManager::new(&connection)?;

    if (&mngr).system_status().is_err() {
        anyhow::bail!(MSG.err_bus_failure_alt(&bus_path));
    }

    drop(mngr);

    eprintln!("{}", MSG.warn_alternate_bus(&bus_path));

    // Reset the environment variable. This is necessary since systemctl may depend on it.
    env::set_var(DBUS_SESSION_BUS_ADDRESS, orig_dbus_address);

    Ok(connection)
}

fn run_switch(
    system_manager: impl sd_switch::systemd::ServiceManager,
    options: &Options,
) -> anyhow::Result<()> {
    sd_switch::switch(
        &system_manager,
        options.old_dir.as_deref(),
        &options.new_dir,
        options.dry_run,
        options.verbose,
        options.timeout,
    )
    .with_context(|| MSG.err_switch_failed())
}

fn log_error<A, E: fmt::Display>(verbose: bool, r: Result<A, E>) -> Result<A, E> {
    match r {
        Err(ref e) if verbose => {
            println!("{e}");
            r
        }
        _ => r,
    }
}

fn main() -> anyhow::Result<()> {
    let options = parse_args(&env::args().collect::<Vec<_>>());

    if options.verbose {
        println!("{}", package_version());
        println!("Options are {options:#?}");
    }

    if options.dry_run {
        println!("{}", MSG.performing_dry_run());
    }

    if options.verbose {
        println!("{}", MSG.enable_verbose_output());
    }

    if !options.new_dir.is_dir() {
        eprintln!("{}", MSG.err_not_a_directory(options.new_dir.as_path()));
        exit(1);
    }

    if let Some(ref old_dir) = options.old_dir {
        if !old_dir.is_dir() {
            eprintln!("{}", MSG.err_not_a_directory(old_dir));
            exit(1);
        }
    }

    let verbose = options.verbose;

    if !options.force_systemctl {
        // Try connecting to systemd using D-Bus. We first try the D-Bus address
        // indicated by the system environment. If that doesn't work, then we
        // forcibly attempt to use the standard socket path that systemd tend to
        // use.
        for open_dbus in [open_std_dbus_connection, open_alt_dbus_connection] {
            if let Ok(conn) = log_error(verbose, open_dbus(&options.target_manager)) {
                if let Ok(mngr) = log_error(verbose, DbusServiceManager::new(&conn)) {
                    return run_switch(&mngr, &options);
                }
            }
        }
    }

    let mngr = sd_switch::systemd::systemctl::SystemctlServiceManager::new(matches!(
        options.target_manager,
        TargetManager::System
    ))?;
    if log_error(verbose, (&mngr).system_status()).is_ok() {
        return run_switch(&mngr, &options);
    }

    anyhow::bail!(MSG.err_no_systemd_conn())
}