1use anyhow::Context;
2use chrono::{DateTime, Utc};
3use kumo_mac_address::get_mac_address;
4use reqwest::header::{HeaderMap, HeaderValue};
5use reqwest::Method;
6use serde::Deserialize;
7use serde_with::serde_as;
8use sysinfo::System;
9
10#[derive(Debug)]
11pub struct MachineInfo {
12 pub hostname: String,
13 pub mac_address: String,
14 pub machine_uid: Option<String>,
15 pub node_id: Option<String>,
16 pub num_cores: usize,
17 pub kernel_version: Option<String>,
18 pub platform: String,
19 pub distribution: String,
20 pub os_version: String,
21 pub total_memory_bytes: u64,
22 pub container_runtime: Option<String>,
23 pub cpu_brand: String,
24 pub cloud_provider: Option<CloudProvider>,
25}
26
27#[derive(Debug)]
28pub enum CloudProvider {
29 Aws(aws::IdentityDocument),
31 Azure(azure::InstanceMetadata),
33 Gcp(gcp::InstanceMetadata),
35}
36
37impl CloudProvider {
38 fn augment_fingerprint(&self, components: &mut Vec<String>) {
39 match self {
40 Self::Aws(id) => {
41 components.push(format!("aws_instance_id={}", id.instance_id));
42 }
43 Self::Azure(instance) => {
44 components.push(format!("azure_vm_id={}", instance.compute.vm_id));
45 }
46 Self::Gcp(instance) => {
47 components.push(format!("gcp_id={}", instance.instance_id));
48 }
49 }
50 }
51}
52
53impl MachineInfo {
54 pub fn fingerprint(&self) -> String {
55 let mut components = vec![];
56 if let Some(provider) = &self.cloud_provider {
57 provider.augment_fingerprint(&mut components);
58 }
59 if let Some(uid) = &self.machine_uid {
60 components.push(format!("machine_uid={uid}"));
61 }
62 if let Some(id) = &self.node_id {
63 components.push(format!("node_id={id}"));
64 }
65 if components.is_empty() {
66 components.push(format!("mac={}", self.mac_address));
67 }
68 components.join(",")
69 }
70
71 pub fn new() -> Self {
72 let hostname = gethostname::gethostname().to_string_lossy().to_string();
73 let mac = get_mac_address();
74 let mac_address = format!(
75 "{:02X}:{:02X}:{:02X}:{:02X}:{:02X}:{:02X}",
76 mac[0], mac[1], mac[2], mac[3], mac[4], mac[5]
77 );
78
79 let machine_uid = machine_uid::get().ok();
80
81 let arch = System::cpu_arch();
82
83 let mut system = System::new();
84 system.refresh_memory();
85 system.refresh_cpu_all();
86
87 let mut cpu_info = vec![];
88 for cpu in system.cpus() {
89 let info = cpu.brand().to_string();
90 if cpu_info.contains(&info) {
91 continue;
92 }
93 cpu_info.push(info);
94 }
95 let cpu_brand = cpu_info.join(", ");
96
97 Self {
98 hostname,
99 machine_uid,
100 mac_address,
101 node_id: None,
102 num_cores: num_cpus::get(),
103 platform: format!("{}/{arch}", std::env::consts::OS),
104 distribution: System::distribution_id(),
105 os_version: System::long_os_version()
106 .unwrap_or_else(|| std::env::consts::OS.to_string()),
107 total_memory_bytes: system.total_memory(),
108 container_runtime: in_container::get_container_runtime().map(|r| r.to_string()),
109 kernel_version: System::kernel_version(),
110 cpu_brand,
111 cloud_provider: None,
112 }
113 }
114
115 pub async fn query_cloud_provider(&mut self) {
117 let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel();
118 {
119 let tx = tx.clone();
120 tokio::task::spawn(async move {
121 if let Ok(id) = aws::IdentityDocument::query().await {
122 tx.send(CloudProvider::Aws(id)).ok();
123 }
124 });
125 }
126 {
127 let tx = tx.clone();
128 tokio::task::spawn(async move {
129 if let Ok(id) = azure::InstanceMetadata::query().await {
130 tx.send(CloudProvider::Azure(id)).ok();
131 }
132 });
133 }
134 {
135 let tx = tx.clone();
136 tokio::task::spawn(async move {
137 if let Ok(id) = gcp::InstanceMetadata::query().await {
138 tx.send(CloudProvider::Gcp(id)).ok();
139 }
140 });
141 }
142 drop(tx);
143
144 tokio::select! {
145 biased;
146
147 provider = rx.recv() => {
149 self.cloud_provider = provider;
150 }
151
152 _ = tokio::time::sleep(tokio::time::Duration::from_secs(3)) => {}
154 };
155 }
156}
157
158pub mod azure {
159 use super::*;
160 use serde_json::Value;
161 use std::collections::BTreeMap;
162
163 #[serde_as]
164 #[derive(Deserialize, Debug)]
165 #[serde(rename_all = "camelCase")]
166 pub struct InstanceMetadata {
167 pub compute: Compute,
168 pub network: Network,
169 }
170
171 #[serde_as]
172 #[derive(Deserialize, Debug)]
173 #[serde(rename_all = "camelCase")]
174 pub struct Compute {
175 pub az_environment: String,
176 pub vm_id: String,
177 pub vm_size: String,
178 pub location: String,
179 #[serde(flatten)]
180 pub unknown_: BTreeMap<String, Value>,
181 }
182
183 #[serde_as]
184 #[derive(Deserialize, Debug)]
185 #[serde(rename_all = "camelCase")]
186 pub struct Network {
187 pub interface: Vec<NetworkInterface>,
188 }
189
190 #[serde_as]
191 #[derive(Deserialize, Debug)]
192 #[serde(rename_all = "camelCase")]
193 pub struct NetworkInterface {
194 pub ipv4: Ipv4Info,
195 pub mac_address: String,
196 #[serde(flatten)]
197 pub unknown_: BTreeMap<String, Value>,
198 }
199
200 #[serde_as]
201 #[derive(Deserialize, Debug)]
202 #[serde(rename_all = "camelCase")]
203 pub struct Ipv4Info {
204 pub ip_address: Vec<IpAddressInfo>,
205 pub subnet: Vec<SubnetInfo>,
206 }
207
208 #[serde_as]
209 #[derive(Deserialize, Debug)]
210 #[serde(rename_all = "camelCase")]
211 pub struct IpAddressInfo {
212 pub private_ip_address: String,
213 #[serde(default)]
214 pub public_ip_address: String,
215 }
216 #[serde_as]
217 #[derive(Deserialize, Debug)]
218 #[serde(rename_all = "camelCase")]
219 pub struct SubnetInfo {
220 pub address: String,
221 pub prefix: String,
222 }
223
224 impl InstanceMetadata {
225 pub async fn query_via(base_url: &str) -> anyhow::Result<Self> {
226 let client = reqwest::Client::builder()
227 .no_proxy()
228 .timeout(std::time::Duration::from_secs(1))
229 .build()
230 .unwrap();
231
232 let mut headers = HeaderMap::new();
233 headers.insert("Metadata", HeaderValue::from_static("true"));
234
235 let request = client
236 .request(
237 Method::GET,
238 format!("{base_url}/metadata/instance?api-version=2021-02-01"),
239 )
240 .headers(headers)
241 .build()?;
242 let response = client.execute(request).await?;
243
244 let status = response.status();
245
246 let body_text = response
247 .text()
248 .await
249 .context("failed to read response body")?;
250 if status.is_client_error() || status.is_server_error() {
251 anyhow::bail!("failed to query identity: {status:?} {body_text}");
252 }
253
254 Ok(serde_json::from_str(&body_text)?)
255 }
256
257 pub async fn query() -> anyhow::Result<Self> {
258 Self::query_via("http://169.254.169.254").await
259 }
260 }
261
262 #[cfg(test)]
263 #[tokio::test]
264 async fn test_metadata() {
265 use mockito::Server;
266
267 let mut server = Server::new_async().await;
268 let _mock = server
269 .mock("GET", "/metadata/instance?api-version=2021-02-01")
270 .match_header("Metadata", "true")
271 .with_status(200)
272 .with_body(
273 r#"{
274 "compute": {
275 "azEnvironment": "AZUREPUBLICCLOUD",
276 "additionalCapabilities": {
277 "hibernationEnabled": "true"
278 },
279 "hostGroup": {
280 "id": "testHostGroupId"
281 },
282 "extendedLocation": {
283 "type": "edgeZone",
284 "name": "microsoftlosangeles"
285 },
286 "evictionPolicy": "",
287 "isHostCompatibilityLayerVm": "true",
288 "licenseType": "",
289 "location": "westus",
290 "name": "examplevmname",
291 "offer": "UbuntuServer",
292 "osProfile": {
293 "adminUsername": "admin",
294 "computerName": "examplevmname",
295 "disablePasswordAuthentication": "true"
296 },
297 "osType": "Linux",
298 "placementGroupId": "f67c14ab-e92c-408c-ae2d-da15866ec79a",
299 "plan": {
300 "name": "planName",
301 "product": "planProduct",
302 "publisher": "planPublisher"
303 },
304 "platformFaultDomain": "36",
305 "platformSubFaultDomain": "",
306 "platformUpdateDomain": "42",
307 "priority": "Regular",
308 "publicKeys": [{
309 "keyData": "ssh-rsa 0",
310 "path": "/home/user/.ssh/authorized_keys0"
311 },
312 {
313 "keyData": "ssh-rsa 1",
314 "path": "/home/user/.ssh/authorized_keys1"
315 }
316 ],
317 "publisher": "Canonical",
318 "resourceGroupName": "macikgo-test-may-23",
319 "resourceId": "/subscriptions/xxxxx-xxxx-xxxx-xxxx-xxxxxxxxxxx/resourceGroups/macikgo-test-may-23/providers/Microsoft.Compute/virtualMachines/examplevmname",
320 "securityProfile": {
321 "secureBootEnabled": "true",
322 "virtualTpmEnabled": "false",
323 "encryptionAtHost": "true",
324 "securityType": "TrustedLaunch"
325 },
326 "sku": "18.04-LTS",
327 "storageProfile": {
328 "dataDisks": [{
329 "bytesPerSecondThrottle": "979202048",
330 "caching": "None",
331 "createOption": "Empty",
332 "diskCapacityBytes": "274877906944",
333 "diskSizeGB": "1024",
334 "image": {
335 "uri": ""
336 },
337 "isSharedDisk": "false",
338 "isUltraDisk": "true",
339 "lun": "0",
340 "managedDisk": {
341 "id": "/subscriptions/xxxxx-xxxx-xxxx-xxxx-xxxxxxxxxxx/resourceGroups/macikgo-test-may-23/providers/Microsoft.Compute/disks/exampledatadiskname",
342 "storageAccountType": "StandardSSD_LRS"
343 },
344 "name": "exampledatadiskname",
345 "opsPerSecondThrottle": "65280",
346 "vhd": {
347 "uri": ""
348 },
349 "writeAcceleratorEnabled": "false"
350 }],
351 "imageReference": {
352 "id": "",
353 "offer": "UbuntuServer",
354 "publisher": "Canonical",
355 "sku": "16.04.0-LTS",
356 "version": "latest",
357 "communityGalleryImageId": "/CommunityGalleries/testgallery/Images/1804Gen2/Versions/latest",
358 "sharedGalleryImageId": "/SharedGalleries/1P/Images/gen2/Versions/latest",
359 "exactVersion": "1.1686127202.30113"
360 },
361 "osDisk": {
362 "caching": "ReadWrite",
363 "createOption": "FromImage",
364 "diskSizeGB": "30",
365 "diffDiskSettings": {
366 "option": "Local"
367 },
368 "encryptionSettings": {
369 "enabled": "false",
370 "diskEncryptionKey": {
371 "sourceVault": {
372 "id": "/subscriptions/test-source-guid/resourceGroups/testrg/providers/Microsoft.KeyVault/vaults/test-kv"
373 },
374 "secretUrl": "https://test-disk.vault.azure.net/secrets/xxxxx-xxxx-xxxx-xxxx-xxxxxxxxxxx/xxxxx-xxxx-xxxx-xxxx-xxxxxxxxxxx"
375 },
376 "keyEncryptionKey": {
377 "sourceVault": {
378 "id": "/subscriptions/test-key-guid/resourceGroups/testrg/providers/Microsoft.KeyVault/vaults/test-kv"
379 },
380 "keyUrl": "https://test-key.vault.azure.net/secrets/xxxxx-xxxx-xxxx-xxxx-xxxxxxxxxxx/xxxxx-xxxx-xxxx-xxxx-xxxxxxxxxxx"
381 }
382 },
383 "image": {
384 "uri": ""
385 },
386 "managedDisk": {
387 "id": "/subscriptions/xxxxx-xxxx-xxxx-xxxx-xxxxxxxxxxx/resourceGroups/macikgo-test-may-23/providers/Microsoft.Compute/disks/exampleosdiskname",
388 "storageAccountType": "StandardSSD_LRS"
389 },
390 "name": "exampleosdiskname",
391 "osType": "Linux",
392 "vhd": {
393 "uri": ""
394 },
395 "writeAcceleratorEnabled": "false"
396 },
397 "resourceDisk": {
398 "size": "4096"
399 }
400 },
401 "subscriptionId": "xxxxx-xxxx-xxxx-xxxx-xxxxxxxxxxx",
402 "tags": "baz:bash;foo:bar",
403 "version": "15.05.22",
404 "virtualMachineScaleSet": {
405 "id": "/subscriptions/xxxxxxxx-xxxxx-xxx-xxx-xxxx/resourceGroups/resource-group-name/providers/Microsoft.Compute/virtualMachineScaleSets/virtual-machine-scale-set-name"
406 },
407 "vmId": "02aab8a4-74ef-476e-8182-f6d2ba4166a6",
408 "vmScaleSetName": "crpteste9vflji9",
409 "vmSize": "Standard_A3",
410 "zone": ""
411 },
412 "network": {
413 "interface": [{
414 "ipv4": {
415 "ipAddress": [{
416 "privateIpAddress": "10.144.133.132",
417 "publicIpAddress": ""
418 }],
419 "subnet": [{
420 "address": "10.144.133.128",
421 "prefix": "26"
422 }]
423 },
424 "ipv6": {
425 "ipAddress": [
426 ]
427 },
428 "macAddress": "0011AAFFBB22"
429 }]
430 }
431}"#,
432 )
433 .create_async()
434 .await;
435
436 let id = InstanceMetadata::query_via(&server.url()).await.unwrap();
437 eprintln!("{id:#?}");
438 assert_eq!(id.compute.vm_id, "02aab8a4-74ef-476e-8182-f6d2ba4166a6");
439 }
440}
441
442pub mod aws {
443 use super::*;
444
445 #[serde_as]
446 #[derive(Deserialize, Debug)]
447 #[serde(rename_all = "camelCase")]
448 pub struct IdentityDocument {
449 #[serde(default)]
450 #[serde_as(as = "serde_with::DefaultOnNull<_>")]
451 pub devpay_product_codes: Vec<String>,
452 #[serde(default)]
453 #[serde_as(as = "serde_with::DefaultOnNull<_>")]
454 pub marketplace_product_codes: Vec<String>,
455 pub availability_zone: String,
456 pub private_ip: String,
457 pub version: String,
458 pub instance_id: String,
459 #[serde(default)]
460 #[serde_as(as = "serde_with::DefaultOnNull<_>")]
461 pub billing_products: Vec<String>,
462 pub instance_type: String,
463 pub account_id: String,
464 pub image_id: String,
465 pub pending_time: DateTime<Utc>,
466 pub architecture: String,
467 pub kernel_id: Option<String>,
468 pub ramdisk_id: Option<String>,
469 pub region: String,
470 }
471
472 impl IdentityDocument {
473 pub async fn query_via(base_url: &str) -> anyhow::Result<Self> {
474 let client = reqwest::Client::builder()
475 .no_proxy()
476 .timeout(std::time::Duration::from_secs(1))
477 .build()
478 .unwrap();
479
480 let mut token = None;
486
487 {
488 let mut headers = HeaderMap::new();
489 headers.insert(
490 "X-aws-ec2-metadata-token-ttl-seconds",
491 HeaderValue::from_static("60"),
492 );
493
494 let request = client
495 .request(Method::PUT, format!("{base_url}/latest/api/token"))
496 .headers(headers)
497 .build()?;
498
499 let response = client.execute(request).await?;
506
507 if response.status().is_success() {
508 if let Ok(content) = response.text().await {
509 token.replace(content.trim().to_string());
510 }
511 } else {
512 }
515 }
516
517 let mut headers = HeaderMap::new();
518 if let Some(token) = token.as_deref() {
519 headers.insert("X-aws-ec2-metadata-token", HeaderValue::from_str(token)?);
520 }
521
522 let request = client
523 .request(
524 Method::GET,
525 format!("{base_url}/latest/dynamic/instance-identity/document"),
526 )
527 .headers(headers)
528 .build()?;
529 let response = client.execute(request).await?;
530
531 let status = response.status();
532
533 let body_text = response
534 .text()
535 .await
536 .context("failed to read response body")?;
537 if status.is_client_error() || status.is_server_error() {
538 anyhow::bail!("failed to query identity: {status:?} {body_text}");
539 }
540
541 Ok(serde_json::from_str(&body_text)?)
542 }
543
544 pub async fn query() -> anyhow::Result<Self> {
545 Self::query_via("http://169.254.169.254").await
546 }
547 }
548
549 #[cfg(test)]
550 #[tokio::test]
551 async fn test_aws_identity_v1() {
552 use mockito::Server;
553
554 let mut server = Server::new_async().await;
555 let _mock = server
556 .mock("GET", "/latest/dynamic/instance-identity/document")
557 .with_status(200)
558 .with_body(
559 r#"{
560 "devpayProductCodes" : null,
561 "marketplaceProductCodes" : [ "1abc2defghijklm3nopqrs4tu" ],
562 "availabilityZone" : "us-west-2b",
563 "privateIp" : "10.158.112.84",
564 "version" : "2017-09-30",
565 "instanceId" : "i-1234567890abcdef0",
566 "billingProducts" : null,
567 "instanceType" : "t2.micro",
568 "accountId" : "123456789012",
569 "imageId" : "ami-5fb8c835",
570 "pendingTime" : "2016-11-19T16:32:11Z",
571 "architecture" : "x86_64",
572 "kernelId" : null,
573 "ramdiskId" : null,
574 "region" : "us-west-2"
575}"#,
576 )
577 .create_async()
578 .await;
579
580 let id = IdentityDocument::query_via(&server.url()).await.unwrap();
581 eprintln!("{id:#?}");
582 assert_eq!(id.instance_id, "i-1234567890abcdef0");
583 }
584
585 #[cfg(test)]
586 #[tokio::test]
587 async fn test_aws_identity_v2() {
588 use mockito::Server;
589
590 let token = "fake-token";
591 let mut server = Server::new_async().await;
592 let _mock = server
593 .mock("PUT", "/latest/api/token")
594 .match_header("X-aws-ec2-metadata-token-ttl-seconds", "60")
595 .with_status(200)
596 .with_body(token)
597 .create_async()
598 .await;
599 let _mock = server
600 .mock("GET", "/latest/dynamic/instance-identity/document")
601 .with_status(200)
602 .match_header("X-aws-ec2-metadata-token", token)
603 .with_body(
604 r#"{
605 "devpayProductCodes" : null,
606 "marketplaceProductCodes" : [ "1abc2defghijklm3nopqrs4tu" ],
607 "availabilityZone" : "us-west-2b",
608 "privateIp" : "10.158.112.84",
609 "version" : "2017-09-30",
610 "instanceId" : "i-1234567890abcdef0",
611 "billingProducts" : null,
612 "instanceType" : "t2.micro",
613 "accountId" : "123456789012",
614 "imageId" : "ami-5fb8c835",
615 "pendingTime" : "2016-11-19T16:32:11Z",
616 "architecture" : "x86_64",
617 "kernelId" : null,
618 "ramdiskId" : null,
619 "region" : "us-west-2"
620}"#,
621 )
622 .create_async()
623 .await;
624
625 let id = IdentityDocument::query_via(&server.url()).await.unwrap();
626 eprintln!("{id:#?}");
627 assert_eq!(id.instance_id, "i-1234567890abcdef0");
628 }
629}
630
631pub mod gcp {
632 use super::*;
633
634 #[derive(Debug)]
635 pub struct InstanceMetadata {
636 pub instance_id: String,
637 }
638
639 impl InstanceMetadata {
640 pub async fn query_via(base_url: &str) -> anyhow::Result<Self> {
641 let client = reqwest::Client::builder()
642 .no_proxy()
643 .timeout(std::time::Duration::from_secs(1))
644 .build()
645 .unwrap();
646
647 let mut headers = HeaderMap::new();
648 headers.insert("Metadata-Flavor", HeaderValue::from_static("Google"));
649
650 let request = client
651 .request(
652 Method::GET,
653 format!("{base_url}/computeMetadata/v1/instance/id"),
654 )
655 .headers(headers)
656 .build()?;
657 let response = client.execute(request).await?;
658 let status = response.status();
659
660 let instance_id = response
661 .text()
662 .await
663 .context("failed to read response body")?
664 .trim()
665 .to_string();
666 if status.is_client_error() || status.is_server_error() {
667 anyhow::bail!("failed to query identity: {status:?} {instance_id}");
668 }
669
670 Ok(Self { instance_id })
671 }
672
673 pub async fn query() -> anyhow::Result<Self> {
674 Self::query_via("http://metadata.google.internal").await
675 }
676 }
677
678 #[cfg(test)]
679 #[tokio::test]
680 async fn test_gcp() {
681 use mockito::Server;
682
683 let mut server = Server::new_async().await;
684 let _mock = server
685 .mock("GET", "/computeMetadata/v1/instance/id")
686 .with_status(200)
687 .with_body("some_id")
688 .create_async()
689 .await;
690
691 let id = InstanceMetadata::query_via(&server.url()).await.unwrap();
692 eprintln!("{id:#?}");
693 assert_eq!(id.instance_id, "some_id");
694 }
695}
696
697#[cfg(test)]
698mod test {
699 #[test]
700 fn test_machine_info() {
701 use super::*;
702 let info = MachineInfo::new();
703 eprintln!("{}", info.fingerprint());
704 eprintln!("{info:#?}");
705 }
712}