keygen-rs 0.11.1

//! Machine activation and management.
//!
//! This module provides functionality for activating machines against licenses,
//! managing machine heartbeats, and checking out machine files for offline use.

use crate::certificate::CertificateFileResponse;
use crate::client::{Client, ClientOptions, Response};
use crate::config::get_config;
use crate::config::KeygenConfig;
use crate::errors::Error;
use crate::insert_optional;
use crate::machine_file::MachineFile;
use crate::KeygenResponseData;
use chrono::{DateTime, Utc};
#[cfg(not(target_arch = "wasm32"))]
use futures::future::{BoxFuture, FutureExt};
use serde::{Deserialize, Serialize};
use serde_json::{json, Value};
use std::collections::HashMap;
use std::sync::Arc;
#[cfg(not(target_arch = "wasm32"))]
use std::time::Duration;
#[cfg(not(target_arch = "wasm32"))]
use tokio::sync::mpsc;

/// Heartbeat status as returned by the Keygen API
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "SCREAMING_SNAKE_CASE")]
pub enum HeartbeatStatus {
    Alive,
    Dead,
    NotStarted,
    Resurrected,
}

impl HeartbeatStatus {
    /// Parses a HeartbeatStatus from a string, returning None for unknown values
    pub fn parse(s: &str) -> Option<Self> {
        match s.to_uppercase().as_str() {
            "ALIVE" => Some(Self::Alive),
            "DEAD" => Some(Self::Dead),
            "NOT_STARTED" => Some(Self::NotStarted),
            "RESURRECTED" => Some(Self::Resurrected),
            _ => None,
        }
    }
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub(crate) struct MachineAttributes {
    pub fingerprint: String,
    pub name: Option<String>,
    pub platform: Option<String>,
    pub hostname: Option<String>,
    pub ip: Option<String>,
    pub cores: Option<i32>,
    pub metadata: Option<HashMap<String, Value>>,
    #[serde(rename = "requireHeartbeat")]
    pub require_heartbeat: bool,
    #[serde(rename = "heartbeatStatus")]
    pub heartbeat_status: String,
    #[serde(rename = "heartbeatDuration")]
    pub heartbeat_duration: Option<i32>,
    pub created: DateTime<Utc>,
    pub updated: DateTime<Utc>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub(crate) struct MachineResponse {
    pub data: KeygenResponseData<MachineAttributes>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub(crate) struct MachinesResponse {
    pub data: Vec<KeygenResponseData<MachineAttributes>>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Machine {
    pub id: String,
    pub fingerprint: String,
    pub name: Option<String>,
    pub platform: Option<String>,
    pub hostname: Option<String>,
    pub ip: Option<String>,
    pub cores: Option<i32>,
    pub metadata: Option<HashMap<String, Value>>,
    #[serde(rename = "requireHeartbeat")]
    pub require_heartbeat: bool,
    #[serde(rename = "heartbeatStatus")]
    pub heartbeat_status: String,
    #[serde(rename = "heartbeatDuration")]
    pub heartbeat_duration: Option<i32>,
    pub created: DateTime<Utc>,
    pub updated: DateTime<Utc>,
    pub account_id: Option<String>,
    pub environment_id: Option<String>,
    pub product_id: Option<String>,
    pub license_id: Option<String>,
    pub owner_id: Option<String>,
    pub group_id: Option<String>,
    #[serde(skip)]
    pub config: Option<Arc<KeygenConfig>>,
}

#[derive(Debug, Clone, Default)]
pub struct MachineCheckoutOpts {
    pub ttl: Option<i64>,
    pub include: Option<Vec<String>>,
}

impl MachineCheckoutOpts {
    pub fn new() -> Self {
        Self::default()
    }

    pub fn with_ttl(ttl: i64) -> Self {
        Self {
            ttl: Some(ttl),
            ..Self::default()
        }
    }

    pub fn with_include(include: Vec<String>) -> Self {
        Self {
            include: Some(include),
            ..Self::default()
        }
    }
}

#[derive(Debug, Clone, Serialize, Deserialize, Default)]
pub struct MachineListFilters {
    pub license: Option<String>,
    pub user: Option<String>,
    pub platform: Option<String>,
    pub name: Option<String>,
    pub fingerprint: Option<String>,
    pub ip: Option<String>,
    pub hostname: Option<String>,
    pub product: Option<String>,
    pub owner: Option<String>,
    pub group: Option<String>,
    pub policy: Option<String>,
    pub key: Option<String>,
    pub metadata: Option<HashMap<String, Value>>,
    pub page_number: Option<i32>,
    pub page_size: Option<i32>,
    pub limit: Option<i32>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MachineCreateRequest {
    pub fingerprint: String,
    pub name: Option<String>,
    pub platform: Option<String>,
    pub hostname: Option<String>,
    pub ip: Option<String>,
    pub cores: Option<i32>,
    pub metadata: Option<HashMap<String, Value>>,
    pub license_id: String,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MachineUpdateRequest {
    pub name: Option<String>,
    pub platform: Option<String>,
    pub hostname: Option<String>,
    pub ip: Option<String>,
    pub cores: Option<i32>,
    pub metadata: Option<HashMap<String, Value>>,
}

impl Machine {
    pub(crate) fn from(data: KeygenResponseData<MachineAttributes>) -> Machine {
        Machine {
            id: data.id,
            fingerprint: data.attributes.fingerprint,
            name: data.attributes.name,
            platform: data.attributes.platform,
            hostname: data.attributes.hostname,
            ip: data.attributes.ip,
            cores: data.attributes.cores,
            metadata: data.attributes.metadata,
            require_heartbeat: data.attributes.require_heartbeat,
            heartbeat_status: data.attributes.heartbeat_status,
            heartbeat_duration: data.attributes.heartbeat_duration,
            created: data.attributes.created,
            updated: data.attributes.updated,
            account_id: data.relationships.account_id(),
            environment_id: data.relationships.environment_id(),
            product_id: data.relationships.product_id(),
            license_id: data.relationships.license_id(),
            owner_id: data.relationships.owner_id(),
            group_id: data.relationships.group_id(),
            config: None,
        }
    }

    /// Associates a configuration with this Machine
    pub fn with_config(mut self, config: KeygenConfig) -> Self {
        self.config = Some(Arc::new(config));
        self
    }

    /// Gets a client for this machine, using the associated config or global config
    fn get_client(&self) -> Result<Client, Error> {
        let config = if let Some(ref cfg) = self.config {
            cfg.as_ref().clone()
        } else {
            get_config()?
        };
        Client::new(ClientOptions::from(config))
    }

    pub async fn deactivate(&self) -> Result<(), Error> {
        let client = self.get_client()?;
        let _response = client
            .delete::<(), serde_json::Value>(&format!("machines/{}", self.id), None::<&()>)
            .await?;
        Ok(())
    }

    pub async fn checkout(&self, options: &MachineCheckoutOpts) -> Result<MachineFile, Error> {
        let mut query = json!({
            "encrypt": 1
        });

        if let Some(ttl) = options.ttl {
            query["ttl"] = ttl.into();
        }

        if let Some(ref include) = options.include {
            query["include"] = json!(include.join(","));
        } else {
            query["include"] = "license.entitlements".into();
        }

        let client = self.get_client()?;
        let response = client
            .post(
                &format!("machines/{}/actions/check-out", self.id),
                None::<&()>,
                Some(&query),
            )
            .await?;

        let machine_file_response: CertificateFileResponse = serde_json::from_value(response.body)?;
        let machine_file = MachineFile::from(machine_file_response.data);
        Ok(machine_file)
    }

    pub async fn ping(&self) -> Result<Machine, Error> {
        let client = self.get_client()?;
        let response: Response<MachineResponse> = client
            .post(
                &format!("machines/{}/actions/ping", self.id),
                None::<&()>,
                None::<&()>,
            )
            .await?;
        let machine = Machine::from(response.body.data).with_config(
            self.config
                .as_ref()
                .ok_or(Error::MissingConfiguration)?
                .as_ref()
                .clone(),
        );
        Ok(machine)
    }

    #[cfg(not(target_arch = "wasm32"))]
    pub fn monitor(
        self: Arc<Self>,
        heartbeat_interval: Duration,
        tx: Option<mpsc::Sender<Result<Machine, Error>>>,
        mut cancel_rx: Option<mpsc::Receiver<()>>,
    ) -> BoxFuture<'static, ()> {
        async move {
            async fn send(
                tx: &Option<mpsc::Sender<Result<Machine, Error>>>,
                result: Result<Machine, Error>,
            ) {
                if let Some(tx) = tx {
                    let _ = tx.send(result).await;
                }
            }

            let mut interval = tokio::time::interval(heartbeat_interval);
            interval.tick().await;

            send(&tx, self.ping().await).await;

            loop {
                tokio::select! {
                    _ = interval.tick() => {
                        send(&tx, self.ping().await).await;
                    }
                    _ = async {
                        if let Some(ref mut rx) = cancel_rx {
                            rx.recv().await
                        } else {
                            std::future::pending::<Option<()>>().await
                        }
                    } => {
                        break;
                    }
                }
            }
        }
        .boxed()
    }

    /// Create a new machine
    #[cfg(feature = "token")]
    pub async fn create(request: MachineCreateRequest) -> Result<Machine, Error> {
        let config = get_config()?;
        let client = Client::new(ClientOptions::from(config))?;
        let mut attributes = serde_json::Map::new();
        attributes.insert("fingerprint".to_string(), json!(request.fingerprint));

        insert_optional(&mut attributes, "name", request.name)?;
        insert_optional(&mut attributes, "platform", request.platform)?;
        insert_optional(&mut attributes, "hostname", request.hostname)?;
        insert_optional(&mut attributes, "ip", request.ip)?;
        insert_optional(&mut attributes, "cores", request.cores)?;
        insert_optional(&mut attributes, "metadata", request.metadata)?;

        let body = json!({
            "data": {
                "type": "machines",
                "attributes": attributes,
                "relationships": {
                    "license": {
                        "data": {
                            "type": "licenses",
                            "id": request.license_id
                        }
                    }
                }
            }
        });

        let response = client.post("machines", Some(&body), None::<&()>).await?;
        let machine_response: MachineResponse = serde_json::from_value(response.body)?;
        Ok(Machine::from(machine_response.data))
    }

    /// List machines with optional filters
    #[cfg(feature = "token")]
    pub async fn list(filters: Option<MachineListFilters>) -> Result<Vec<Machine>, Error> {
        let config = get_config()?;
        let client = Client::new(ClientOptions::from(config))?;
        let mut query_params = Vec::new();
        if let Some(filters) = filters {
            if let Some(license) = filters.license {
                query_params.push(("license".to_string(), license));
            }
            if let Some(user) = filters.user {
                query_params.push(("user".to_string(), user));
            }
            if let Some(platform) = filters.platform {
                query_params.push(("platform".to_string(), platform));
            }
            if let Some(name) = filters.name {
                query_params.push(("name".to_string(), name));
            }
            if let Some(fingerprint) = filters.fingerprint {
                query_params.push(("fingerprint".to_string(), fingerprint));
            }
            if let Some(ip) = filters.ip {
                query_params.push(("ip".to_string(), ip));
            }
            if let Some(hostname) = filters.hostname {
                query_params.push(("hostname".to_string(), hostname));
            }
            if let Some(product) = filters.product {
                query_params.push(("product".to_string(), product));
            }
            if let Some(owner) = filters.owner {
                query_params.push(("owner".to_string(), owner));
            }
            if let Some(group) = filters.group {
                query_params.push(("group".to_string(), group));
            }
            if let Some(policy) = filters.policy {
                query_params.push(("policy".to_string(), policy));
            }
            if let Some(key) = filters.key {
                query_params.push(("key".to_string(), key));
            }
            if let Some(metadata) = filters.metadata {
                for (key, value) in metadata {
                    query_params.push((format!("metadata[{key}]"), value.to_string()));
                }
            }
            // Add pagination parameters
            if let Some(page_number) = filters.page_number {
                query_params.push(("page[number]".to_string(), page_number.to_string()));
            }
            if let Some(page_size) = filters.page_size {
                query_params.push(("page[size]".to_string(), page_size.to_string()));
            }
            if let Some(limit) = filters.limit {
                query_params.push(("limit".to_string(), limit.to_string()));
            }
        }

        let query = if query_params.is_empty() {
            None
        } else {
            Some(
                query_params
                    .into_iter()
                    .collect::<HashMap<String, String>>(),
            )
        };

        let response = client.get("machines", query.as_ref()).await?;
        let machines_response: MachinesResponse = serde_json::from_value(response.body)?;
        Ok(machines_response
            .data
            .into_iter()
            .map(Machine::from)
            .collect())
    }

    /// Get a machine by ID
    #[cfg(feature = "token")]
    pub async fn get(id: &str) -> Result<Machine, Error> {
        let config = get_config()?;
        let client = Client::new(ClientOptions::from(config))?;
        let endpoint = format!("machines/{id}");
        let response = client.get(&endpoint, None::<&()>).await?;
        let machine_response: MachineResponse = serde_json::from_value(response.body)?;
        Ok(Machine::from(machine_response.data))
    }

    /// Update a machine
    #[cfg(feature = "token")]
    pub async fn update(&self, request: MachineUpdateRequest) -> Result<Machine, Error> {
        let client = self.get_client()?;
        let endpoint = format!("machines/{}", self.id);

        let mut attributes = serde_json::Map::new();
        insert_optional(&mut attributes, "name", request.name)?;
        insert_optional(&mut attributes, "platform", request.platform)?;
        insert_optional(&mut attributes, "hostname", request.hostname)?;
        insert_optional(&mut attributes, "ip", request.ip)?;
        insert_optional(&mut attributes, "cores", request.cores)?;
        insert_optional(&mut attributes, "metadata", request.metadata)?;

        let body = json!({
            "data": {
                "type": "machines",
                "attributes": attributes
            }
        });

        let response = client.patch(&endpoint, Some(&body), None::<&()>).await?;
        let machine_response: MachineResponse = serde_json::from_value(response.body)?;
        Ok(Machine::from(machine_response.data))
    }

    /// Reset machine heartbeat
    #[cfg(feature = "token")]
    pub async fn reset(&self) -> Result<Machine, Error> {
        let client = self.get_client()?;
        let endpoint = format!("machines/{}/actions/reset", self.id);
        let response = client.post(&endpoint, None::<&()>, None::<&()>).await?;
        let machine_response: MachineResponse = serde_json::from_value(response.body)?;
        Ok(Machine::from(machine_response.data))
    }

    /// Change the machine owner.
    #[cfg(feature = "token")]
    pub async fn change_owner(&self, owner_id: &str) -> Result<Machine, Error> {
        let client = self.get_client()?;
        let endpoint = format!("machines/{}/owner", self.id);
        let body = json!({
            "data": {
                "type": "users",
                "id": owner_id
            }
        });
        let response = client.put(&endpoint, Some(&body), None::<&()>).await?;
        let machine_response: MachineResponse = serde_json::from_value(response.body)?;
        Ok(Machine::from(machine_response.data))
    }

    /// Change the machine group.
    #[cfg(feature = "token")]
    pub async fn change_group(&self, group_id: &str) -> Result<Machine, Error> {
        let client = self.get_client()?;
        let endpoint = format!("machines/{}/group", self.id);
        let body = json!({
            "data": {
                "type": "groups",
                "id": group_id
            }
        });
        let response = client.put(&endpoint, Some(&body), None::<&()>).await?;
        let machine_response: MachineResponse = serde_json::from_value(response.body)?;
        Ok(Machine::from(machine_response.data))
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::{
        KeygenRelationship, KeygenRelationshipData, KeygenRelationships, KeygenResponseData,
    };
    use chrono::Utc;

    #[test]
    fn test_machine_relationships() {
        // Test that all relationship IDs are properly extracted
        let machine_data = KeygenResponseData {
            id: "test-machine-id".to_string(),
            r#type: "machines".to_string(),
            attributes: MachineAttributes {
                fingerprint: "test-fingerprint".to_string(),
                name: Some("Test Machine".to_string()),
                platform: Some("linux".to_string()),
                hostname: Some("test-host".to_string()),
                ip: Some("192.168.1.1".to_string()),
                cores: Some(8),
                metadata: Some(HashMap::new()),
                require_heartbeat: true,
                heartbeat_status: "ALIVE".to_string(),
                heartbeat_duration: Some(3600),
                created: Utc::now(),
                updated: Utc::now(),
            },
            relationships: KeygenRelationships {
                policy: None,
                account: Some(KeygenRelationship {
                    data: Some(KeygenRelationshipData {
                        r#type: "accounts".to_string(),
                        id: "test-account-id".to_string(),
                    }),
                    links: None,
                }),
                product: Some(KeygenRelationship {
                    data: Some(KeygenRelationshipData {
                        r#type: "products".to_string(),
                        id: "test-product-id".to_string(),
                    }),
                    links: None,
                }),
                group: Some(KeygenRelationship {
                    data: Some(KeygenRelationshipData {
                        r#type: "groups".to_string(),
                        id: "test-group-id".to_string(),
                    }),
                    links: None,
                }),
                owner: Some(KeygenRelationship {
                    data: Some(KeygenRelationshipData {
                        r#type: "users".to_string(),
                        id: "test-owner-id".to_string(),
                    }),
                    links: None,
                }),
                users: None,
                machines: None,
                environment: Some(KeygenRelationship {
                    data: Some(KeygenRelationshipData {
                        r#type: "environments".to_string(),
                        id: "test-environment-id".to_string(),
                    }),
                    links: None,
                }),
                license: Some(KeygenRelationship {
                    data: Some(KeygenRelationshipData {
                        r#type: "licenses".to_string(),
                        id: "test-license-id".to_string(),
                    }),
                    links: None,
                }),
                release: None,
                other: HashMap::new(),
            },
        };

        let machine = Machine::from(machine_data);

        assert_eq!(machine.account_id, Some("test-account-id".to_string()));
        assert_eq!(
            machine.environment_id,
            Some("test-environment-id".to_string())
        );
        assert_eq!(machine.product_id, Some("test-product-id".to_string()));
        assert_eq!(machine.license_id, Some("test-license-id".to_string()));
        assert_eq!(machine.owner_id, Some("test-owner-id".to_string()));
        assert_eq!(machine.group_id, Some("test-group-id".to_string()));
        assert_eq!(machine.id, "test-machine-id");
        assert_eq!(machine.fingerprint, "test-fingerprint");
    }

    #[test]
    fn test_machine_without_relationships() {
        // Test that all relationship IDs are None when no relationships exist
        let machine_data = KeygenResponseData {
            id: "test-machine-id".to_string(),
            r#type: "machines".to_string(),
            attributes: MachineAttributes {
                fingerprint: "test-fingerprint".to_string(),
                name: Some("Test Machine".to_string()),
                platform: Some("linux".to_string()),
                hostname: Some("test-host".to_string()),
                ip: Some("192.168.1.1".to_string()),
                cores: Some(8),
                metadata: Some(HashMap::new()),
                require_heartbeat: true,
                heartbeat_status: "ALIVE".to_string(),
                heartbeat_duration: Some(3600),
                created: Utc::now(),
                updated: Utc::now(),
            },
            relationships: KeygenRelationships {
                policy: None,
                account: None,
                product: None,
                group: None,
                owner: None,
                users: None,
                machines: None,
                environment: None,
                license: None,
                release: None,
                other: HashMap::new(),
            },
        };

        let machine = Machine::from(machine_data);

        assert_eq!(machine.account_id, None);
        assert_eq!(machine.environment_id, None);
        assert_eq!(machine.product_id, None);
        assert_eq!(machine.license_id, None);
        assert_eq!(machine.owner_id, None);
        assert_eq!(machine.group_id, None);
    }

    #[test]
    fn test_heartbeat_status_parse() {
        assert_eq!(
            HeartbeatStatus::parse("ALIVE"),
            Some(HeartbeatStatus::Alive)
        );
        assert_eq!(HeartbeatStatus::parse("DEAD"), Some(HeartbeatStatus::Dead));
        assert_eq!(
            HeartbeatStatus::parse("NOT_STARTED"),
            Some(HeartbeatStatus::NotStarted)
        );
        assert_eq!(
            HeartbeatStatus::parse("RESURRECTED"),
            Some(HeartbeatStatus::Resurrected)
        );
        assert_eq!(HeartbeatStatus::parse("UNKNOWN"), None);
    }

    #[test]
    fn test_heartbeat_status_parse_case_insensitive() {
        assert_eq!(
            HeartbeatStatus::parse("alive"),
            Some(HeartbeatStatus::Alive)
        );
        assert_eq!(HeartbeatStatus::parse("dead"), Some(HeartbeatStatus::Dead));
        assert_eq!(
            HeartbeatStatus::parse("not_started"),
            Some(HeartbeatStatus::NotStarted)
        );
        assert_eq!(
            HeartbeatStatus::parse("resurrected"),
            Some(HeartbeatStatus::Resurrected)
        );
    }

    #[test]
    fn test_heartbeat_status_serialize() {
        assert_eq!(
            serde_json::to_string(&HeartbeatStatus::Alive).unwrap(),
            "\"ALIVE\""
        );
        assert_eq!(
            serde_json::to_string(&HeartbeatStatus::Dead).unwrap(),
            "\"DEAD\""
        );
        assert_eq!(
            serde_json::to_string(&HeartbeatStatus::NotStarted).unwrap(),
            "\"NOT_STARTED\""
        );
        assert_eq!(
            serde_json::to_string(&HeartbeatStatus::Resurrected).unwrap(),
            "\"RESURRECTED\""
        );
    }

    #[test]
    fn test_heartbeat_status_deserialize() {
        assert_eq!(
            serde_json::from_str::<HeartbeatStatus>("\"ALIVE\"").unwrap(),
            HeartbeatStatus::Alive
        );
        assert_eq!(
            serde_json::from_str::<HeartbeatStatus>("\"DEAD\"").unwrap(),
            HeartbeatStatus::Dead
        );
        assert_eq!(
            serde_json::from_str::<HeartbeatStatus>("\"NOT_STARTED\"").unwrap(),
            HeartbeatStatus::NotStarted
        );
        assert_eq!(
            serde_json::from_str::<HeartbeatStatus>("\"RESURRECTED\"").unwrap(),
            HeartbeatStatus::Resurrected
        );
    }
}