1use 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#[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 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 pub fn with_config(mut self, config: KeygenConfig) -> Self {
196 self.config = Some(Arc::new(config));
197 self
198 }
199
200 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 #[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 #[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 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 #[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 #[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 #[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 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 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}