amphetamine 0.1.2

//! The one privileged capability Amphetamine uses, and how it is granted.
//!
//! macOS lets any user *raise* a process's nice value without privileges, but
//! only root can lower it again. That asymmetry is the whole problem: without
//! help, deprioritising an app during a focus session would be a one-way door.
//!
//! The grant is therefore scoped to exactly one command shape —
//! `renice 0 -p <pid>` — which can only ever return a process to *normal*
//! priority. It cannot deprioritise anything (that needs no privileges
//! anyway), cannot run another program, and cannot open a shell. Amphetamine
//! only ever demotes processes already sitting at nice 0, so restoring to 0 is
//! exact rather than approximate, and the grant never needs to be widened.

use anyhow::{Context, Result, bail};
use std::os::unix::fs::DirBuilderExt;
use std::path::{Path, PathBuf};
use std::process::Command;

pub const SUDOERS_PATH: &str = "/etc/sudoers.d/amphetamine";
pub const RENICE: &str = "/usr/bin/renice";

pub fn sudoers_rule(user: &str) -> String {
    format!(
        "# Amphetamine — installed by `amph setup`.\n\
         #\n\
         # Grants exactly one capability: returning a process to normal scheduling\n\
         # priority. The `0` is literal, so this rule cannot deprioritise anything,\n\
         # cannot run any other command, and cannot open a shell.\n\
         #\n\
         # It exists because macOS lets any user raise a process's nice value\n\
         # unprivileged, but only root can lower it back. Without this line, a focus\n\
         # session could deprioritise an app and never undo it.\n\
         #\n\
         # Remove with:  sudo rm {SUDOERS_PATH}\n\
         #\n\
         {user} ALL=(root) NOPASSWD: {RENICE} 0 -p [0-9]*\n"
    )
}

pub fn current_user() -> String {
    std::env::var("USER").unwrap_or_else(|_| {
        String::from_utf8_lossy(
            &Command::new("id")
                .arg("-un")
                .output()
                .map(|o| o.stdout)
                .unwrap_or_default(),
        )
        .trim()
        .to_owned()
    })
}

pub fn is_installed() -> bool {
    Path::new(SUDOERS_PATH).exists()
}

/// Whether a demotion could actually be undone right now.
///
/// Both halves matter. A bare `sudo -n` probe can succeed purely because the
/// user ran `sudo` in the last few minutes and the timestamp is still warm —
/// which would green-light a demotion that becomes unrestorable once the
/// timestamp expires. Requiring the drop-in to exist as well removes that trap.
pub fn can_restore() -> bool {
    if !is_installed() {
        return false;
    }
    // Reniceing ourselves to 0 when we are already at 0 is a genuine no-op, so
    // this probes the grant without changing anything.
    Command::new("sudo")
        .args(["-n", RENICE, "0", "-p", &std::process::id().to_string()])
        .stdout(std::process::Stdio::null())
        .stderr(std::process::Stdio::null())
        .status()
        .is_ok_and(|s| s.success())
}

/// Returns one process to nice 0.
pub fn renice_to_zero(pid: i32) -> Result<()> {
    let out = Command::new("sudo")
        .args(["-n", RENICE, "0", "-p", &pid.to_string()])
        .output()
        .context("running sudo renice")?;
    if !out.status.success() {
        bail!(
            "renice {pid}: {}",
            String::from_utf8_lossy(&out.stderr).trim()
        );
    }
    Ok(())
}

/// Installs the drop-in, validating it before it can take effect.
///
/// A malformed file in `/etc/sudoers.d` can lock the machine out of `sudo`
/// entirely, so the content is checked with `visudo -c` while it is still in a
/// private directory, and only then moved into place.
pub fn install() -> Result<()> {
    let user = current_user();
    if user.is_empty() {
        bail!("could not determine the current username");
    }
    let rule = sudoers_rule(&user);

    // A 0700 staging directory: nothing else can swap the file out between
    // validation and installation.
    let dir = staging_dir()?;
    let staged = dir.join("amphetamine");
    std::fs::write(&staged, &rule).with_context(|| format!("writing {}", staged.display()))?;

    let check = Command::new("sudo")
        .args(["/usr/sbin/visudo", "-c", "-f"])
        .arg(&staged)
        .output()
        .context("running visudo to validate the rule")?;
    if !check.status.success() {
        std::fs::remove_dir_all(&dir).ok();
        bail!(
            "refusing to install: visudo rejected the rule\n{}",
            String::from_utf8_lossy(&check.stderr).trim()
        );
    }

    let install = Command::new("sudo")
        .args(["install", "-m", "0440", "-o", "root", "-g", "wheel"])
        .arg(&staged)
        .arg(SUDOERS_PATH)
        .status()
        .context("installing the sudoers drop-in")?;
    std::fs::remove_dir_all(&dir).ok();
    if !install.success() {
        bail!("could not install {SUDOERS_PATH}");
    }
    if !can_restore() {
        bail!("{SUDOERS_PATH} was installed but the grant does not work");
    }
    Ok(())
}

pub fn uninstall() -> Result<()> {
    if !is_installed() {
        return Ok(());
    }
    let ok = Command::new("sudo")
        .args(["rm", "-f", SUDOERS_PATH])
        .status()
        .context("removing the sudoers drop-in")?;
    if !ok.success() {
        bail!("could not remove {SUDOERS_PATH}");
    }
    Ok(())
}

fn staging_dir() -> Result<PathBuf> {
    let dir = std::env::temp_dir().join(format!("amph-setup-{}", std::process::id()));
    std::fs::DirBuilder::new()
        .recursive(true)
        .mode(0o700)
        .create(&dir)
        .with_context(|| format!("creating {}", dir.display()))?;
    Ok(dir)
}

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

    #[test]
    fn rule_grants_only_renice_to_zero() {
        let rule = sudoers_rule("someone");
        let spec = rule
            .lines()
            .find(|l| !l.starts_with('#') && !l.trim().is_empty())
            .expect("rule must have exactly one directive");

        assert_eq!(
            spec,
            "someone ALL=(root) NOPASSWD: /usr/bin/renice 0 -p [0-9]*"
        );
        // The literal 0 is what makes this safe: no wildcard in the nice slot.
        assert!(!spec.contains("renice *"));
        assert!(!spec.contains("ALL:ALL"));
        assert!(!spec.contains("(ALL)"));
        // One directive only — no second line smuggling in a broader grant.
        assert_eq!(
            rule.lines()
                .filter(|l| !l.starts_with('#') && !l.trim().is_empty())
                .count(),
            1
        );
    }

    #[test]
    fn rule_is_accepted_by_visudo() {
        // Validating the exact text we would install catches a syntax error
        // here rather than after it is already in /etc/sudoers.d.
        let dir = staging_dir().unwrap();
        let f = dir.join("candidate");
        std::fs::write(&f, sudoers_rule(&current_user())).unwrap();
        let out = Command::new("/usr/sbin/visudo")
            .args(["-c", "-f"])
            .arg(&f)
            .output();
        std::fs::remove_dir_all(&dir).ok();

        // visudo may need root to run at all; only assert when it actually ran.
        if let Ok(o) = out
            && !String::from_utf8_lossy(&o.stderr).contains("Permission denied")
        {
            assert!(
                o.status.success(),
                "visudo rejected our rule: {}{}",
                String::from_utf8_lossy(&o.stdout),
                String::from_utf8_lossy(&o.stderr)
            );
        }
    }

    #[test]
    fn can_restore_is_false_without_the_dropin() {
        // Whatever the sudo timestamp says, no drop-in means no guarantee.
        if !is_installed() {
            assert!(!can_restore());
        }
    }
}