keygen_rs/
machine.rs

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