Skip to main content

keygen_rs/
machine.rs

1//! Machine activation and management.
2//!
3//! This module provides functionality for activating machines against licenses,
4//! managing machine heartbeats, and checking out machine files for offline use.
5
6use crate::certificate::CertificateFileResponse;
7use crate::client::{Client, ClientOptions, Response};
8use crate::config::get_config;
9use crate::config::KeygenConfig;
10use crate::errors::Error;
11use crate::machine_file::MachineFile;
12use crate::KeygenResponseData;
13use chrono::{DateTime, Utc};
14use futures::future::{BoxFuture, FutureExt};
15use serde::{Deserialize, Serialize};
16use serde_json::{json, Value};
17use std::collections::HashMap;
18use std::sync::Arc;
19use std::time::Duration;
20use tokio::sync::mpsc;
21
22/// Heartbeat status as returned by the Keygen API
23#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
24#[serde(rename_all = "SCREAMING_SNAKE_CASE")]
25pub enum HeartbeatStatus {
26    Alive,
27    Dead,
28    NotStarted,
29    Resurrected,
30}
31
32impl HeartbeatStatus {
33    /// Parses a HeartbeatStatus from a string, returning None for unknown values
34    pub fn parse(s: &str) -> Option<Self> {
35        match s.to_uppercase().as_str() {
36            "ALIVE" => Some(Self::Alive),
37            "DEAD" => Some(Self::Dead),
38            "NOT_STARTED" => Some(Self::NotStarted),
39            "RESURRECTED" => Some(Self::Resurrected),
40            _ => None,
41        }
42    }
43}
44
45#[derive(Debug, Clone, Serialize, Deserialize)]
46pub(crate) struct MachineAttributes {
47    pub fingerprint: String,
48    pub name: Option<String>,
49    pub platform: Option<String>,
50    pub hostname: Option<String>,
51    pub ip: Option<String>,
52    pub cores: Option<i32>,
53    pub metadata: Option<HashMap<String, Value>>,
54    #[serde(rename = "requireHeartbeat")]
55    pub require_heartbeat: bool,
56    #[serde(rename = "heartbeatStatus")]
57    pub heartbeat_status: String,
58    #[serde(rename = "heartbeatDuration")]
59    pub heartbeat_duration: Option<i32>,
60    pub created: DateTime<Utc>,
61    pub updated: DateTime<Utc>,
62}
63
64#[derive(Debug, Clone, Serialize, Deserialize)]
65pub(crate) struct MachineResponse {
66    pub data: KeygenResponseData<MachineAttributes>,
67}
68
69#[derive(Debug, Clone, Serialize, Deserialize)]
70pub(crate) struct MachinesResponse {
71    pub data: Vec<KeygenResponseData<MachineAttributes>>,
72}
73
74#[derive(Debug, Clone, Serialize, Deserialize)]
75pub struct Machine {
76    pub id: String,
77    pub fingerprint: String,
78    pub name: Option<String>,
79    pub platform: Option<String>,
80    pub hostname: Option<String>,
81    pub ip: Option<String>,
82    pub cores: Option<i32>,
83    pub metadata: Option<HashMap<String, Value>>,
84    #[serde(rename = "requireHeartbeat")]
85    pub require_heartbeat: bool,
86    #[serde(rename = "heartbeatStatus")]
87    pub heartbeat_status: String,
88    #[serde(rename = "heartbeatDuration")]
89    pub heartbeat_duration: Option<i32>,
90    pub created: DateTime<Utc>,
91    pub updated: DateTime<Utc>,
92    pub account_id: Option<String>,
93    pub environment_id: Option<String>,
94    pub product_id: Option<String>,
95    pub license_id: Option<String>,
96    pub owner_id: Option<String>,
97    pub group_id: Option<String>,
98    #[serde(skip)]
99    pub config: Option<Arc<KeygenConfig>>,
100}
101
102#[derive(Debug, Clone, Default)]
103pub struct MachineCheckoutOpts {
104    pub ttl: Option<i64>,
105    pub include: Option<Vec<String>>,
106}
107
108impl MachineCheckoutOpts {
109    pub fn new() -> Self {
110        Self::default()
111    }
112
113    pub fn with_ttl(ttl: i64) -> Self {
114        Self {
115            ttl: Some(ttl),
116            ..Self::default()
117        }
118    }
119
120    pub fn with_include(include: Vec<String>) -> Self {
121        Self {
122            include: Some(include),
123            ..Self::default()
124        }
125    }
126}
127
128#[derive(Debug, Clone, Serialize, Deserialize, Default)]
129pub struct MachineListFilters {
130    pub license: Option<String>,
131    pub user: Option<String>,
132    pub platform: Option<String>,
133    pub name: Option<String>,
134    pub fingerprint: Option<String>,
135    pub ip: Option<String>,
136    pub hostname: Option<String>,
137    pub product: Option<String>,
138    pub owner: Option<String>,
139    pub group: Option<String>,
140    pub metadata: Option<HashMap<String, Value>>,
141    pub page_number: Option<i32>,
142    pub page_size: Option<i32>,
143    pub limit: Option<i32>,
144}
145
146#[derive(Debug, Clone, Serialize, Deserialize)]
147pub struct MachineCreateRequest {
148    pub fingerprint: String,
149    pub name: Option<String>,
150    pub platform: Option<String>,
151    pub hostname: Option<String>,
152    pub ip: Option<String>,
153    pub cores: Option<i32>,
154    pub metadata: Option<HashMap<String, Value>>,
155    pub license_id: String,
156}
157
158#[derive(Debug, Clone, Serialize, Deserialize)]
159pub struct MachineUpdateRequest {
160    pub name: Option<String>,
161    pub platform: Option<String>,
162    pub hostname: Option<String>,
163    pub ip: Option<String>,
164    pub cores: Option<i32>,
165    pub metadata: Option<HashMap<String, Value>>,
166}
167
168impl Machine {
169    pub(crate) fn from(data: KeygenResponseData<MachineAttributes>) -> Machine {
170        Machine {
171            id: data.id,
172            fingerprint: data.attributes.fingerprint,
173            name: data.attributes.name,
174            platform: data.attributes.platform,
175            hostname: data.attributes.hostname,
176            ip: data.attributes.ip,
177            cores: data.attributes.cores,
178            metadata: data.attributes.metadata,
179            require_heartbeat: data.attributes.require_heartbeat,
180            heartbeat_status: data.attributes.heartbeat_status,
181            heartbeat_duration: data.attributes.heartbeat_duration,
182            created: data.attributes.created,
183            updated: data.attributes.updated,
184            account_id: data.relationships.account_id(),
185            environment_id: data.relationships.environment_id(),
186            product_id: data.relationships.product_id(),
187            license_id: data.relationships.license_id(),
188            owner_id: data.relationships.owner_id(),
189            group_id: data.relationships.group_id(),
190            config: None,
191        }
192    }
193
194    /// Associates a configuration with this Machine
195    pub fn with_config(mut self, config: KeygenConfig) -> Self {
196        self.config = Some(Arc::new(config));
197        self
198    }
199
200    /// Gets a client for this machine, using the associated config or global config
201    fn get_client(&self) -> Result<Client, Error> {
202        let config = if let Some(ref cfg) = self.config {
203            cfg.as_ref().clone()
204        } else {
205            get_config()?
206        };
207        Client::new(ClientOptions::from(config))
208    }
209
210    pub async fn deactivate(&self) -> Result<(), Error> {
211        let client = self.get_client()?;
212        let _response = client
213            .delete::<(), serde_json::Value>(&format!("machines/{}", self.id), None::<&()>)
214            .await?;
215        Ok(())
216    }
217
218    pub async fn checkout(&self, options: &MachineCheckoutOpts) -> Result<MachineFile, Error> {
219        let mut query = json!({
220            "encrypt": 1
221        });
222
223        if let Some(ttl) = options.ttl {
224            query["ttl"] = ttl.into();
225        }
226
227        if let Some(ref include) = options.include {
228            query["include"] = json!(include.join(","));
229        } else {
230            query["include"] = "license.entitlements".into();
231        }
232
233        let client = self.get_client()?;
234        let response = client
235            .post(
236                &format!("machines/{}/actions/check-out", self.id),
237                None::<&()>,
238                Some(&query),
239            )
240            .await?;
241
242        let machine_file_response: CertificateFileResponse = serde_json::from_value(response.body)?;
243        let machine_file = MachineFile::from(machine_file_response.data);
244        Ok(machine_file)
245    }
246
247    pub async fn ping(&self) -> Result<Machine, Error> {
248        let client = self.get_client()?;
249        let response: Response<MachineResponse> = client
250            .post(
251                &format!("machines/{}/actions/ping", self.id),
252                None::<&()>,
253                None::<&()>,
254            )
255            .await?;
256        let machine = Machine::from(response.body.data).with_config(
257            self.config
258                .as_ref()
259                .ok_or(Error::MissingConfiguration)?
260                .as_ref()
261                .clone(),
262        );
263        Ok(machine)
264    }
265
266    pub fn monitor(
267        self: Arc<Self>,
268        heartbeat_interval: Duration,
269        tx: Option<mpsc::Sender<Result<Machine, Error>>>,
270        mut cancel_rx: Option<mpsc::Receiver<()>>,
271    ) -> BoxFuture<'static, ()> {
272        async move {
273            async fn send(
274                tx: &Option<mpsc::Sender<Result<Machine, Error>>>,
275                result: Result<Machine, Error>,
276            ) {
277                if let Some(tx) = tx {
278                    let _ = tx.send(result).await;
279                }
280            }
281
282            let mut interval = tokio::time::interval(heartbeat_interval);
283            interval.tick().await;
284
285            send(&tx, self.ping().await).await;
286
287            loop {
288                tokio::select! {
289                    _ = interval.tick() => {
290                        send(&tx, self.ping().await).await;
291                    }
292                    _ = async {
293                        if let Some(ref mut rx) = cancel_rx {
294                            rx.recv().await
295                        } else {
296                            std::future::pending::<Option<()>>().await
297                        }
298                    } => {
299                        break;
300                    }
301                }
302            }
303        }
304        .boxed()
305    }
306
307    /// Create a new machine
308    #[cfg(feature = "token")]
309    pub async fn create(request: MachineCreateRequest) -> Result<Machine, Error> {
310        let config = get_config()?;
311        let client = Client::new(ClientOptions::from(config))?;
312        let mut attributes = serde_json::Map::new();
313        attributes.insert("fingerprint".to_string(), json!(request.fingerprint));
314
315        if let Some(name) = request.name {
316            attributes.insert("name".to_string(), json!(name));
317        }
318        if let Some(platform) = request.platform {
319            attributes.insert("platform".to_string(), json!(platform));
320        }
321        if let Some(hostname) = request.hostname {
322            attributes.insert("hostname".to_string(), json!(hostname));
323        }
324        if let Some(ip) = request.ip {
325            attributes.insert("ip".to_string(), json!(ip));
326        }
327        if let Some(cores) = request.cores {
328            attributes.insert("cores".to_string(), json!(cores));
329        }
330        if let Some(metadata) = request.metadata {
331            attributes.insert("metadata".to_string(), json!(metadata));
332        }
333
334        let body = json!({
335            "data": {
336                "type": "machines",
337                "attributes": attributes,
338                "relationships": {
339                    "license": {
340                        "data": {
341                            "type": "licenses",
342                            "id": request.license_id
343                        }
344                    }
345                }
346            }
347        });
348
349        let response = client.post("machines", Some(&body), None::<&()>).await?;
350        let machine_response: MachineResponse = serde_json::from_value(response.body)?;
351        Ok(Machine::from(machine_response.data))
352    }
353
354    /// List machines with optional filters
355    #[cfg(feature = "token")]
356    pub async fn list(filters: Option<MachineListFilters>) -> Result<Vec<Machine>, Error> {
357        let config = get_config()?;
358        let client = Client::new(ClientOptions::from(config))?;
359        let mut query_params = Vec::new();
360        if let Some(filters) = filters {
361            if let Some(license) = filters.license {
362                query_params.push(("license".to_string(), license));
363            }
364            if let Some(user) = filters.user {
365                query_params.push(("user".to_string(), user));
366            }
367            if let Some(platform) = filters.platform {
368                query_params.push(("platform".to_string(), platform));
369            }
370            if let Some(name) = filters.name {
371                query_params.push(("name".to_string(), name));
372            }
373            if let Some(fingerprint) = filters.fingerprint {
374                query_params.push(("fingerprint".to_string(), fingerprint));
375            }
376            if let Some(ip) = filters.ip {
377                query_params.push(("ip".to_string(), ip));
378            }
379            if let Some(hostname) = filters.hostname {
380                query_params.push(("hostname".to_string(), hostname));
381            }
382            if let Some(product) = filters.product {
383                query_params.push(("product".to_string(), product));
384            }
385            if let Some(owner) = filters.owner {
386                query_params.push(("owner".to_string(), owner));
387            }
388            if let Some(group) = filters.group {
389                query_params.push(("group".to_string(), group));
390            }
391            if let Some(metadata) = filters.metadata {
392                for (key, value) in metadata {
393                    query_params.push((format!("metadata[{key}]"), value.to_string()));
394                }
395            }
396            // Add pagination parameters
397            if let Some(page_number) = filters.page_number {
398                query_params.push(("page[number]".to_string(), page_number.to_string()));
399            }
400            if let Some(page_size) = filters.page_size {
401                query_params.push(("page[size]".to_string(), page_size.to_string()));
402            }
403            if let Some(limit) = filters.limit {
404                query_params.push(("limit".to_string(), limit.to_string()));
405            }
406        }
407
408        let query = if query_params.is_empty() {
409            None
410        } else {
411            Some(
412                query_params
413                    .into_iter()
414                    .collect::<HashMap<String, String>>(),
415            )
416        };
417
418        let response = client.get("machines", query.as_ref()).await?;
419        let machines_response: MachinesResponse = serde_json::from_value(response.body)?;
420        Ok(machines_response
421            .data
422            .into_iter()
423            .map(Machine::from)
424            .collect())
425    }
426
427    /// Get a machine by ID
428    #[cfg(feature = "token")]
429    pub async fn get(id: &str) -> Result<Machine, Error> {
430        let config = get_config()?;
431        let client = Client::new(ClientOptions::from(config))?;
432        let endpoint = format!("machines/{id}");
433        let response = client.get(&endpoint, None::<&()>).await?;
434        let machine_response: MachineResponse = serde_json::from_value(response.body)?;
435        Ok(Machine::from(machine_response.data))
436    }
437
438    /// Update a machine
439    #[cfg(feature = "token")]
440    pub async fn update(&self, request: MachineUpdateRequest) -> Result<Machine, Error> {
441        let client = self.get_client()?;
442        let endpoint = format!("machines/{}", self.id);
443
444        let mut attributes = serde_json::Map::new();
445        if let Some(name) = request.name {
446            attributes.insert("name".to_string(), json!(name));
447        }
448        if let Some(platform) = request.platform {
449            attributes.insert("platform".to_string(), json!(platform));
450        }
451        if let Some(hostname) = request.hostname {
452            attributes.insert("hostname".to_string(), json!(hostname));
453        }
454        if let Some(ip) = request.ip {
455            attributes.insert("ip".to_string(), json!(ip));
456        }
457        if let Some(cores) = request.cores {
458            attributes.insert("cores".to_string(), json!(cores));
459        }
460        if let Some(metadata) = request.metadata {
461            attributes.insert("metadata".to_string(), json!(metadata));
462        }
463
464        let body = json!({
465            "data": {
466                "type": "machines",
467                "attributes": attributes
468            }
469        });
470
471        let response = client.patch(&endpoint, Some(&body), None::<&()>).await?;
472        let machine_response: MachineResponse = serde_json::from_value(response.body)?;
473        Ok(Machine::from(machine_response.data))
474    }
475
476    /// Reset machine heartbeat
477    #[cfg(feature = "token")]
478    pub async fn reset(&self) -> Result<Machine, Error> {
479        let client = self.get_client()?;
480        let endpoint = format!("machines/{}/actions/reset", self.id);
481        let response = client.post(&endpoint, None::<&()>, None::<&()>).await?;
482        let machine_response: MachineResponse = serde_json::from_value(response.body)?;
483        Ok(Machine::from(machine_response.data))
484    }
485}
486
487#[cfg(test)]
488mod tests {
489    use super::*;
490    use crate::{
491        KeygenRelationship, KeygenRelationshipData, KeygenRelationships, KeygenResponseData,
492    };
493    use chrono::Utc;
494
495    #[test]
496    fn test_machine_relationships() {
497        // Test that all relationship IDs are properly extracted
498        let machine_data = KeygenResponseData {
499            id: "test-machine-id".to_string(),
500            r#type: "machines".to_string(),
501            attributes: MachineAttributes {
502                fingerprint: "test-fingerprint".to_string(),
503                name: Some("Test Machine".to_string()),
504                platform: Some("linux".to_string()),
505                hostname: Some("test-host".to_string()),
506                ip: Some("192.168.1.1".to_string()),
507                cores: Some(8),
508                metadata: Some(HashMap::new()),
509                require_heartbeat: true,
510                heartbeat_status: "ALIVE".to_string(),
511                heartbeat_duration: Some(3600),
512                created: Utc::now(),
513                updated: Utc::now(),
514            },
515            relationships: KeygenRelationships {
516                policy: None,
517                account: Some(KeygenRelationship {
518                    data: Some(KeygenRelationshipData {
519                        r#type: "accounts".to_string(),
520                        id: "test-account-id".to_string(),
521                    }),
522                    links: None,
523                }),
524                product: Some(KeygenRelationship {
525                    data: Some(KeygenRelationshipData {
526                        r#type: "products".to_string(),
527                        id: "test-product-id".to_string(),
528                    }),
529                    links: None,
530                }),
531                group: Some(KeygenRelationship {
532                    data: Some(KeygenRelationshipData {
533                        r#type: "groups".to_string(),
534                        id: "test-group-id".to_string(),
535                    }),
536                    links: None,
537                }),
538                owner: Some(KeygenRelationship {
539                    data: Some(KeygenRelationshipData {
540                        r#type: "users".to_string(),
541                        id: "test-owner-id".to_string(),
542                    }),
543                    links: None,
544                }),
545                users: None,
546                machines: None,
547                environment: Some(KeygenRelationship {
548                    data: Some(KeygenRelationshipData {
549                        r#type: "environments".to_string(),
550                        id: "test-environment-id".to_string(),
551                    }),
552                    links: None,
553                }),
554                license: Some(KeygenRelationship {
555                    data: Some(KeygenRelationshipData {
556                        r#type: "licenses".to_string(),
557                        id: "test-license-id".to_string(),
558                    }),
559                    links: None,
560                }),
561                release: None,
562                other: HashMap::new(),
563            },
564        };
565
566        let machine = Machine::from(machine_data);
567
568        assert_eq!(machine.account_id, Some("test-account-id".to_string()));
569        assert_eq!(
570            machine.environment_id,
571            Some("test-environment-id".to_string())
572        );
573        assert_eq!(machine.product_id, Some("test-product-id".to_string()));
574        assert_eq!(machine.license_id, Some("test-license-id".to_string()));
575        assert_eq!(machine.owner_id, Some("test-owner-id".to_string()));
576        assert_eq!(machine.group_id, Some("test-group-id".to_string()));
577        assert_eq!(machine.id, "test-machine-id");
578        assert_eq!(machine.fingerprint, "test-fingerprint");
579    }
580
581    #[test]
582    fn test_machine_without_relationships() {
583        // Test that all relationship IDs are None when no relationships exist
584        let machine_data = KeygenResponseData {
585            id: "test-machine-id".to_string(),
586            r#type: "machines".to_string(),
587            attributes: MachineAttributes {
588                fingerprint: "test-fingerprint".to_string(),
589                name: Some("Test Machine".to_string()),
590                platform: Some("linux".to_string()),
591                hostname: Some("test-host".to_string()),
592                ip: Some("192.168.1.1".to_string()),
593                cores: Some(8),
594                metadata: Some(HashMap::new()),
595                require_heartbeat: true,
596                heartbeat_status: "ALIVE".to_string(),
597                heartbeat_duration: Some(3600),
598                created: Utc::now(),
599                updated: Utc::now(),
600            },
601            relationships: KeygenRelationships {
602                policy: None,
603                account: None,
604                product: None,
605                group: None,
606                owner: None,
607                users: None,
608                machines: None,
609                environment: None,
610                license: None,
611                release: None,
612                other: HashMap::new(),
613            },
614        };
615
616        let machine = Machine::from(machine_data);
617
618        assert_eq!(machine.account_id, None);
619        assert_eq!(machine.environment_id, None);
620        assert_eq!(machine.product_id, None);
621        assert_eq!(machine.license_id, None);
622        assert_eq!(machine.owner_id, None);
623        assert_eq!(machine.group_id, None);
624    }
625
626    #[test]
627    fn test_heartbeat_status_parse() {
628        assert_eq!(
629            HeartbeatStatus::parse("ALIVE"),
630            Some(HeartbeatStatus::Alive)
631        );
632        assert_eq!(HeartbeatStatus::parse("DEAD"), Some(HeartbeatStatus::Dead));
633        assert_eq!(
634            HeartbeatStatus::parse("NOT_STARTED"),
635            Some(HeartbeatStatus::NotStarted)
636        );
637        assert_eq!(
638            HeartbeatStatus::parse("RESURRECTED"),
639            Some(HeartbeatStatus::Resurrected)
640        );
641        assert_eq!(HeartbeatStatus::parse("UNKNOWN"), None);
642    }
643
644    #[test]
645    fn test_heartbeat_status_parse_case_insensitive() {
646        assert_eq!(
647            HeartbeatStatus::parse("alive"),
648            Some(HeartbeatStatus::Alive)
649        );
650        assert_eq!(HeartbeatStatus::parse("dead"), Some(HeartbeatStatus::Dead));
651        assert_eq!(
652            HeartbeatStatus::parse("not_started"),
653            Some(HeartbeatStatus::NotStarted)
654        );
655        assert_eq!(
656            HeartbeatStatus::parse("resurrected"),
657            Some(HeartbeatStatus::Resurrected)
658        );
659    }
660
661    #[test]
662    fn test_heartbeat_status_serialize() {
663        assert_eq!(
664            serde_json::to_string(&HeartbeatStatus::Alive).unwrap(),
665            "\"ALIVE\""
666        );
667        assert_eq!(
668            serde_json::to_string(&HeartbeatStatus::Dead).unwrap(),
669            "\"DEAD\""
670        );
671        assert_eq!(
672            serde_json::to_string(&HeartbeatStatus::NotStarted).unwrap(),
673            "\"NOT_STARTED\""
674        );
675        assert_eq!(
676            serde_json::to_string(&HeartbeatStatus::Resurrected).unwrap(),
677            "\"RESURRECTED\""
678        );
679    }
680
681    #[test]
682    fn test_heartbeat_status_deserialize() {
683        assert_eq!(
684            serde_json::from_str::<HeartbeatStatus>("\"ALIVE\"").unwrap(),
685            HeartbeatStatus::Alive
686        );
687        assert_eq!(
688            serde_json::from_str::<HeartbeatStatus>("\"DEAD\"").unwrap(),
689            HeartbeatStatus::Dead
690        );
691        assert_eq!(
692            serde_json::from_str::<HeartbeatStatus>("\"NOT_STARTED\"").unwrap(),
693            HeartbeatStatus::NotStarted
694        );
695        assert_eq!(
696            serde_json::from_str::<HeartbeatStatus>("\"RESURRECTED\"").unwrap(),
697            HeartbeatStatus::Resurrected
698        );
699    }
700}