orpheus/src/inventory.rs

use serde::de::{self, Deserializer, MapAccess, Visitor};
use serde::{Deserialize, Serialize};
use serde_yaml_bw::Value as YamlValue;
use std::cmp;
use std::collections::HashMap;
use std::fmt;

#[derive(Serialize, Deserialize, Debug, Clone)]
pub(crate) struct Inventory {
    pub webserver: String,
    pub hosts: HashMap<String, Host>,
}

#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Eq)]
pub(crate) struct Host {
    pub(crate) name: String,
    pub(crate) version: String,
    pub(crate) r#type: String,
    pub(crate) cmdline: Vec<CmdParams>,
}

#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Eq)]
pub(crate) struct BootParam {
    kernel: String,
    initrd: Vec<String>,
    cmdline: String,
    message: String,
}

#[derive(Serialize, Debug, Clone, PartialEq, Eq)]
pub(crate) enum CmdParams {
    DataSource(String),
    Apkovl(String),
    Modloop,
    String(String),
}

impl Inventory {
    fn build_cmdline(&self, host: &str) -> String {
        let host = self.hosts.get(host).unwrap();
        let mut result_string = String::new();
        for param in &host.cmdline {
            match param {
                CmdParams::DataSource(path) => result_string.push_str(&format!(
                    "ds=nocloud;s={}/tiny-cloud/{}",
                    self.webserver, path
                )),
                CmdParams::Apkovl(path) => {
                    result_string.push_str(&format!("apkovl={}/apkovl/{}", self.webserver, path))
                }
                CmdParams::Modloop => result_string.push_str(&format!(
                    "modloop={}/{}/modloop-{}",
                    self.webserver, host.version, host.r#type
                )),
                CmdParams::String(cmdstring) => result_string.push_str(&cmdstring),
            };
            result_string.push_str(" ");
        }
        // remove last space
        result_string = result_string.trim().into();
        return result_string;
    }

    pub fn get(&self, mac_addr: &str) -> Result<BootParam, ()> {
        let host = self.hosts.get(mac_addr).unwrap();
        let boot_param = BootParam {
            kernel: format!(
                "{}/{}/vmlinuz-{}",
                self.webserver, host.version, host.r#type
            )
            .into(),
            initrd: vec![format!(
                "{}/{}/initramfs-{}",
                self.webserver, host.version, host.r#type
            )
            .into()],
            cmdline: self.build_cmdline(&mac_addr),
            message: format!("Booting host: {}", &host.name).into(),
        };
        return Ok(boot_param);
    }
}

impl<'de> Deserialize<'de> for CmdParams {
    fn deserialize<D>(deserializer: D) -> Result<CmdParams, D::Error>
    where
        D: Deserializer<'de>,
    {
        struct CmdParamsVisitor;

        impl<'de> Visitor<'de> for CmdParamsVisitor {
            type Value = CmdParams;

            fn expecting(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
                formatter.write_str("a cmdline parameter in the form of a string or map")
            }

            fn visit_str<E>(self, value: &str) -> Result<CmdParams, E>
            where
                E: de::Error,
            {
                if value == "modloop" {
                    return Ok(CmdParams::Modloop);
                } else {
                    return Ok(CmdParams::String(value.to_string()));
                }
            }

            fn visit_map<M>(self, mut access: M) -> Result<CmdParams, M::Error>
            where
                M: MapAccess<'de>,
            {
                if let Some((key, value)) = access.next_entry::<String, Option<String>>()? {
                    match key.as_str() {
                        "ds" => {
                            let value =
                                value.ok_or_else(|| de::Error::missing_field("ds value"))?;
                            Ok(CmdParams::DataSource(value))
                        }
                        "apkovl" => {
                            let value =
                                value.ok_or_else(|| de::Error::missing_field("apkovl value"))?;
                            Ok(CmdParams::Apkovl(value))
                        }
                        "modloop" => Ok(CmdParams::Modloop),
                        _ => Err(de::Error::unknown_field(&key, &["ds", "apkovl", "modloop"])),
                    }
                } else {
                    Err(de::Error::custom("Expected a key-value pair"))
                }
            }
        }

        deserializer.deserialize_any(CmdParamsVisitor)
    }
}

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

    #[test]
    fn cmdparams_deserialization_string() {
        let yaml_data = "---\nconsole=tty0";
        let cmd_param: CmdParams = serde_yaml::from_str(yaml_data).unwrap();
        assert_eq!(cmd_param, CmdParams::String("console=tty0".to_string()));
    }

    #[test]
    fn cmdparams_deserialization_modloop() {
        let yaml_data = "---\nmodloop";
        let cmd_param: CmdParams = serde_yaml::from_str(yaml_data).unwrap();
        assert_eq!(cmd_param, CmdParams::Modloop);
    }

    #[test]
    fn cmdparams_deserialization_ds() {
        let yaml_data = "---\nds: k3s-server";
        let cmd_param: CmdParams = serde_yaml::from_str(yaml_data).unwrap();
        assert_eq!(cmd_param, CmdParams::DataSource("k3s-server".to_string()));
    }

    #[test]
    fn cmdparams_deserialization_apkovl() {
        let yaml_data = "---\napkovl: nginx.apkov.tar.gz";
        let cmd_param: CmdParams = serde_yaml::from_str(yaml_data).unwrap();
        assert_eq!(
            cmd_param,
            CmdParams::Apkovl("nginx.apkov.tar.gz".to_string())
        );
    }

    #[test]
    fn cmdparams_deserialization_invalid() {
        let yaml_data = "---\nunknown: value";
        let result: Result<CmdParams, _> = serde_yaml::from_str(yaml_data);
        assert!(result.is_err());
    }

    // test Host deserialization
    #[test]
    fn host_deserialization_valid() {
        let yaml_data = r#"---
        name: hostname
        version: alpine-v3.21
        type: test
        cmdline:
          - console=tty0
        "#;
        let host: Host = serde_yaml::from_str(yaml_data).unwrap();
        assert_eq!(
            host,
            Host {
                name: "hostname".to_string(),
                version: "alpine-v3.21".to_string(),
                r#type: "test".to_string(),
                cmdline: vec!(CmdParams::String("console=tty0".to_string()))
            }
        );

        let yaml_data = r#"---
        name: hostname
        version: alpine-v3.21
        type: test
        cmdline: []
        "#;
        let host: Host = serde_yaml::from_str(yaml_data).unwrap();
        assert_eq!(
            host,
            Host {
                name: "hostname".to_string(),
                version: "alpine-v3.21".to_string(),
                r#type: "test".to_string(),
                cmdline: vec!()
            }
        );
    }

    #[test]
    fn host_deserialization_invalid() {
        let yaml_data = r#"---
        version: alpine-v3.21
        type: test
        cmdline: []
        "#;
        let host: Result<Host, _> = serde_yaml::from_str(yaml_data);
        assert!(host.is_err());
    }

    #[test]
    fn build_cmdline() {
        use serde_json::json;
        let host = Host {
            name: "test".into(),
            version: "test".into(),
            r#type: "test".into(),
            cmdline: vec![
                CmdParams::String("test".into()),
                CmdParams::Apkovl("apkovl.tar.gz".into()),
                CmdParams::Modloop,
                CmdParams::DataSource("test-ds-config".into()),
            ],
        };
        let mut hosts: HashMap<String, Host> = HashMap::new();
        hosts.insert("a:b:c:d:e:f".into(), host.clone());

        let inventory = Inventory {
            webserver: "http://example.net".into(),
            hosts: hosts,
        };

        let cmd_string = inventory.build_cmdline("a:b:c:d:e:f");

        assert_eq!(
            cmd_string,
            "test apkovl=http://example.net/apkovl/apkovl.tar.gz modloop=http://example.net/test/modloop-test ds=nocloud;s=http://example.net/tiny-cloud/test-ds-config"
        );

        assert_eq!(
            json!(inventory.get("a:b:c:d:e:f").unwrap()),
            json!({
            "kernel": "http://example.net/test/vmlinuz-test",
            "initrd": ["http://example.net/test/initramfs-test"],
            "cmdline": "test apkovl=http://example.net/apkovl/apkovl.tar.gz modloop=http://example.net/test/modloop-test ds=nocloud;s=http://example.net/tiny-cloud/test-ds-config",
            "message": "Booting host: test"})
        );
    }
}