honey-biscuit-workshop/workshop-engine/src/pm/apt.rs

use crate::pm::PackageManager;
use os_info::{Info, Version};
use serde::Deserialize;
use serde_yaml::Value;
use tokio::process::Command;

#[derive(Clone)]
pub struct Apt;

/// Repository configuration for apt.
#[derive(Deserialize, Clone)]
struct AptRepo {
    /// The deb line URL (e.g., `https://repo.example.com/ubuntu`)
    url: String,
    /// Distribution codename override (e.g., `focal`, `bookworm`).
    /// If not set, extracted from os_info at runtime.
    #[serde(default)]
    distro: Option<String>,
    /// Components (e.g., `main`, `universe`, `contrib`).
    /// Defaults to `["main"]` if not specified.
    #[serde(default = "default_components")]
    components: Vec<String>,
    /// GPG key URL to download and install.
    #[serde(default)]
    key_url: Option<String>,
    /// GPG key ID to fetch from keyserver.
    #[serde(default)]
    key_id: Option<String>,
    /// Path to save the GPG keyring file.
    #[serde(default = "default_keyring_dir")]
    keyring_path: Option<String>,
    /// Architecture constraint (e.g., `amd64`, `arm64`).
    #[serde(default)]
    arch: Option<String>,
    /// Whether this is a source repo (`deb-src`).
    #[serde(default)]
    source: bool,
}

fn default_components() -> Vec<String> {
    vec!["main".to_string()]
}

fn default_keyring_dir() -> Option<String> {
    Some("/usr/share/keyrings".to_string())
}

impl PackageManager for Apt {
    fn name(&self) -> &'static str {
        "apt"
    }

    fn binary(&self) -> &'static str {
        "apt-get"
    }

    fn update(&self) -> Vec<Command> {
        let mut cmd = Command::new("apt-get");
        cmd.arg("update");
        vec![cmd]
    }

    fn install(&self, package: &[String]) -> Vec<Command> {
        let mut cmd = Command::new("apt-get");
        cmd.arg("install");
        cmd.arg("-y");
        cmd.args(package);
        vec![cmd]
    }

    /// Changes the primary apt mirror by replacing the URL in /etc/apt/sources.list.
    ///
    /// Uses sed to replace the existing mirror URL with the user-specified repo URL.
    fn change_mirror(&self, repo: &str, _osinfo: Option<&Info>) -> Vec<Command> {
        let script = format!(
            "sed -i 's|^deb http[^ ]*|deb {}|' /etc/apt/sources.list",
            repo
        );
        let mut cmd = Command::new("sh");
        cmd.arg("-c").arg(script);
        vec![cmd]
    }

    /// Adds a custom apt repository by creating a .list file in /etc/apt/sources.list.d/
    /// and optionally installing a GPG key.
    fn add_repository(&self, repo: &Value, osinfo: Option<&Info>) -> Vec<Command> {
        let repos: Vec<AptRepo> = match serde_yaml::from_value(repo.clone()) {
            Ok(r) => r,
            Err(e) => {
                log::error!("Failed to parse apt repository: {}", e);
                return Vec::new();
            }
        };

        let codename = repos
            .first()
            .and_then(|r| r.distro.clone())
            .or_else(|| extract_codename(osinfo));

        let codename = match codename {
            Some(c) => c,
            None => {
                log::error!(
                    "Cannot determine distribution codename for apt repository, \
                     set 'distro' field in config or provide os_info"
                );
                return Vec::new();
            }
        };

        let mut cmds: Vec<Command> = vec![];

        for repo_cfg in &repos {
            // Install GPG key if configured
            if let Some(key_url) = &repo_cfg.key_url {
                let keyring = repo_cfg
                    .keyring_path
                    .clone()
                    .unwrap_or_else(|| "/usr/share/keyrings".to_string());

                let keyring_name = format!("{}-keyring.gpg", sanitize_repo_name(&repo_cfg.url));
                let keyring_path = format!("{}/{}", keyring, keyring_name);

                let mut cmd = Command::new("sh");
                cmd.arg("-c").arg(format!(
                    "curl -fsSL '{}' | gpg --dearmor -o '{}' && chmod 644 '{}'",
                    key_url, keyring_path, keyring_path
                ));
                cmds.push(cmd);
            }

            if let Some(key_id) = &repo_cfg.key_id {
                let mut cmd = Command::new("sh");
                cmd.arg("-c").arg(format!(
                    "gpg --keyserver keyserver.ubuntu.com --recv-keys {} && \
                     gpg --export {} | gpg --dearmor > /usr/share/keyrings/{}-keyring.gpg",
                    key_id,
                    key_id,
                    sanitize_repo_name(&repo_cfg.url)
                ));
                cmds.push(cmd);
            }

            // Build the deb line
            let arch_part = repo_cfg
                .arch
                .as_ref()
                .map(|a| format!("[arch={}] ", a))
                .unwrap_or_default();

            let keyring_name = if repo_cfg.key_url.is_some() || repo_cfg.key_id.is_some() {
                format!(
                    " [signed-by=/usr/share/keyrings/{}-keyring.gpg]",
                    sanitize_repo_name(&repo_cfg.url)
                )
            } else {
                String::new()
            };

            let deb_type = if repo_cfg.source { "deb-src" } else { "deb" };
            let components = repo_cfg.components.join(" ");
            let deb_line = format!(
                "{deb_type} {arch_part}{keyring_name} {url} {codename} {components}\n",
                deb_type = deb_type,
                arch_part = arch_part,
                keyring_name = keyring_name,
                url = repo_cfg.url,
                codename = codename,
                components = components,
            );

            // Write the .list file
            let list_name = format!(
                "/etc/apt/sources.list.d/{}.list",
                sanitize_repo_name(&repo_cfg.url)
            );
            let mut cmd = Command::new("sh");
            cmd.arg("-c").arg(format!(
                "cat > '{}' << 'APT_EOF'\n{}APT_EOF",
                list_name, deb_line
            ));
            cmds.push(cmd);
        }

        cmds
    }
}

/// Sanitizes a URL to a safe filename component.
fn sanitize_repo_name(url: &str) -> String {
    url.trim_start_matches("https://")
        .trim_start_matches("http://")
        .trim_start_matches("ppa:")
        .replace(|c: char| !c.is_alphanumeric() && c != '.' && c != '-', "_")
        .trim_matches('_')
        .to_string()
}

/// Extracts the distribution codename from os_info, falling back to version-based heuristics.
fn extract_codename(info: Option<&Info>) -> Option<String> {
    let info = info?;

    // 1. Direct codename from os_info (supported in newer versions)
    if let Some(codename) = info.codename()
        && !codename.is_empty()
    {
        return Some(codename.to_string());
    }

    // 2. Try to infer from version string
    match info.version() {
        Version::Custom(v) if !v.chars().next().map(|c| c.is_numeric()).unwrap_or(true) => {
            // Some systems put codename in Custom field
            return Some(v.clone());
        }
        Version::Semantic(major, minor, _) => {
            // For Ubuntu/Debian, fall back to version number
            return Some(format!("{}.{}", major, minor));
        }
        _ => {}
    }

    None
}

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

    #[test]
    fn test_name() {
        assert_eq!(Apt.name(), "apt");
    }

    #[test]
    fn test_binary() {
        assert_eq!(Apt.binary(), "apt-get");
    }

    #[test]
    fn test_update_command_contents() {
        let cmds = Apt.update();
        let cmd = cmds[0].as_std();
        assert_eq!(cmd.get_program().to_string_lossy(), "apt-get");
        let args: Vec<String> = cmd
            .get_args()
            .map(|a| a.to_string_lossy().into_owned())
            .collect();
        assert_eq!(args, vec!["update"]);
    }

    #[test]
    fn test_install_command_contents() {
        let cmds = Apt.install(&["curl".to_string()]);
        let cmd = cmds[0].as_std();
        assert_eq!(cmd.get_program().to_string_lossy(), "apt-get");
        let args: Vec<String> = cmd
            .get_args()
            .map(|a| a.to_string_lossy().into_owned())
            .collect();
        assert_eq!(args, vec!["install", "-y", "curl"]);
    }

    #[test]
    fn test_install_multiple_packages() {
        let cmds = Apt.install(&["curl".to_string(), "git".to_string()]);
        let cmd = cmds[0].as_std();
        let args: Vec<String> = cmd
            .get_args()
            .map(|a| a.to_string_lossy().into_owned())
            .collect();
        assert_eq!(args, vec!["install", "-y", "curl", "git"]);
    }

    #[test]
    fn test_change_mirror_command() {
        let cmds = Apt.change_mirror("http://mirror.tuna.tsinghua.edu.cn/ubuntu", None);
        let cmd = cmds[0].as_std();
        assert_eq!(cmd.get_program().to_string_lossy(), "sh");
        let args: Vec<String> = cmd
            .get_args()
            .map(|a| a.to_string_lossy().into_owned())
            .collect();
        let full_arg = args.join(" ");
        assert!(full_arg.contains("sed -i"));
        assert!(full_arg.contains("/etc/apt/sources.list"));
        assert!(full_arg.contains("mirror.tuna.tsinghua.edu.cn"));
    }

    #[test]
    fn test_add_repository_simple_deb_line() {
        let repo_yaml = serde_yaml::from_value(serde_yaml::Value::Sequence(vec![
            serde_yaml::Value::Mapping(
                [
                    (
                        serde_yaml::Value::String("url".to_string()),
                        serde_yaml::Value::String("https://repo.example.com/ubuntu".to_string()),
                    ),
                    (
                        serde_yaml::Value::String("distro".to_string()),
                        serde_yaml::Value::String("focal".to_string()),
                    ),
                    (
                        serde_yaml::Value::String("components".to_string()),
                        serde_yaml::Value::Sequence(vec![
                            serde_yaml::Value::String("main".to_string()),
                            serde_yaml::Value::String("universe".to_string()),
                        ]),
                    ),
                ]
                .into_iter()
                .collect(),
            ),
        ]))
        .unwrap();

        let cmds = Apt.add_repository(&repo_yaml, None);
        assert!(!cmds.is_empty());
        let last_cmd = cmds.last().unwrap().as_std();
        let args: Vec<String> = last_cmd
            .get_args()
            .map(|a| a.to_string_lossy().into_owned())
            .collect();
        let full_arg = args.join(" ");
        assert!(full_arg.contains("cat >"));
        assert!(full_arg.contains("sources.list.d"));
        assert!(full_arg.contains("repo.example.com"));
        assert!(full_arg.contains("focal"));
        assert!(full_arg.contains("main universe"));
    }

    #[test]
    fn test_add_repository_with_key_url() {
        let repo_yaml = serde_yaml::from_value(serde_yaml::Value::Sequence(vec![
            serde_yaml::Value::Mapping(
                [
                    (
                        serde_yaml::Value::String("url".to_string()),
                        serde_yaml::Value::String("https://repo.example.com/ubuntu".to_string()),
                    ),
                    (
                        serde_yaml::Value::String("distro".to_string()),
                        serde_yaml::Value::String("focal".to_string()),
                    ),
                    (
                        serde_yaml::Value::String("key_url".to_string()),
                        serde_yaml::Value::String("https://repo.example.com/key.gpg".to_string()),
                    ),
                ]
                .into_iter()
                .collect(),
            ),
        ]))
        .unwrap();

        let cmds = Apt.add_repository(&repo_yaml, None);
        // First command: key download, last command: list file
        assert!(cmds.len() >= 2);
        let first_cmd = cmds[0].as_std();
        let args: Vec<String> = first_cmd
            .get_args()
            .map(|a| a.to_string_lossy().into_owned())
            .collect();
        let full_arg = args.join(" ");
        assert!(full_arg.contains("curl"));
        assert!(full_arg.contains("key.gpg"));
        assert!(full_arg.contains("gpg --dearmor"));
    }

    #[test]
    fn test_sanitize_repo_name() {
        assert_eq!(
            sanitize_repo_name("https://example.com/repo"),
            "example.com_repo"
        );
        assert_eq!(
            sanitize_repo_name("http://mirror.tuna.tsinghua.edu.cn/ubuntu"),
            "mirror.tuna.tsinghua.edu.cn_ubuntu"
        );
        assert_eq!(sanitize_repo_name("ppa:user/name"), "user_name");
    }

    #[test]
    fn test_extract_codename_from_codename_field() {
        let info = Info::with_type(os_info::Type::Ubuntu);
        assert!(extract_codename(Some(&info)).is_none());
    }

    #[test]
    fn test_add_repository_invalid_yaml_returns_empty() {
        let invalid_yaml = serde_yaml::Value::String("invalid".to_string());
        let cmds = Apt.add_repository(&invalid_yaml, None);
        assert!(cmds.is_empty());
    }

    #[test]
    fn test_add_repository_missing_codename_returns_empty() {
        let repo_yaml = serde_yaml::from_value(serde_yaml::Value::Sequence(vec![
            serde_yaml::Value::Mapping(
                [(
                    serde_yaml::Value::String("url".to_string()),
                    serde_yaml::Value::String("https://repo.example.com/ubuntu".to_string()),
                )]
                .into_iter()
                .collect(),
            ),
        ]))
        .unwrap();

        let cmds = Apt.add_repository(&repo_yaml, None);
        assert!(cmds.is_empty());
    }
}