Skip to main content

rusticity_core/
kms.rs

1use crate::config::AwsConfig;
2use anyhow::Result;
3
4#[derive(Clone, Debug)]
5pub struct KmsKey {
6    pub key_id: String,
7    pub key_arn: String,
8    pub description: String,
9    pub key_state: String,
10    pub key_usage: String,
11    pub key_spec: String,
12    pub key_manager: String, // "AWS" | "CUSTOMER"
13    pub creation_date: String,
14    pub expiration_date: String, // for EXTERNAL origin keys
15    pub deletion_date: String,   // when pending deletion
16    pub custom_key_store_id: String,
17    pub origin: String,
18    pub enabled: bool,
19    pub alias: String,
20    pub multi_region: bool,
21}
22
23pub struct KmsClient {
24    config: AwsConfig,
25}
26
27impl KmsClient {
28    pub fn new(config: AwsConfig) -> Self {
29        Self { config }
30    }
31
32    pub async fn list_keys(&self) -> Result<Vec<KmsKey>> {
33        let client = self.config.kms_client();
34
35        // Step 1: list all key IDs with pagination
36        let mut key_ids: Vec<String> = Vec::new();
37        let mut marker: Option<String> = None;
38
39        loop {
40            let mut request = client.list_keys();
41            if let Some(m) = marker {
42                request = request.marker(m);
43            }
44
45            let response = request.send().await?;
46
47            if let Some(keys) = response.keys {
48                for entry in keys {
49                    if let Some(id) = entry.key_id {
50                        key_ids.push(id);
51                    }
52                }
53            }
54
55            if !response.truncated {
56                break;
57            }
58            marker = response.next_marker;
59            if marker.is_none() {
60                break;
61            }
62        }
63
64        // Step 2: fetch aliases once to build a key_id → alias map
65        let mut alias_map: std::collections::HashMap<String, String> =
66            std::collections::HashMap::new();
67        let mut alias_marker: Option<String> = None;
68
69        loop {
70            let mut request = client.list_aliases();
71            if let Some(m) = alias_marker {
72                request = request.marker(m);
73            }
74            let response = request.send().await?;
75
76            if let Some(aliases) = response.aliases {
77                for a in aliases {
78                    if let (Some(name), Some(target_id)) = (a.alias_name, a.target_key_id) {
79                        let display = name.strip_prefix("alias/").unwrap_or(&name).to_string();
80                        alias_map
81                            .entry(target_id)
82                            .and_modify(|v| {
83                                v.push_str(", ");
84                                v.push_str(&display);
85                            })
86                            .or_insert(display);
87                    }
88                }
89            }
90
91            if !response.truncated {
92                break;
93            }
94            alias_marker = response.next_marker;
95            if alias_marker.is_none() {
96                break;
97            }
98        }
99
100        // Step 3: describe each key in parallel (10 concurrent to avoid throttling)
101        use futures::stream::{self, StreamExt};
102
103        let mut result: Vec<KmsKey> = stream::iter(key_ids)
104            .map(|key_id| {
105                let client = self.config.kms_client();
106                let alias = alias_map.get(&key_id).cloned().unwrap_or_default();
107                async move {
108                    let resp = client.describe_key().key_id(&key_id).send().await.ok()?;
109                    let meta = resp.key_metadata?;
110
111                    let key_manager = meta
112                        .key_manager
113                        .as_ref()
114                        .map(|m| format!("{:?}", m))
115                        .unwrap_or_default();
116                    let key_state = meta
117                        .key_state
118                        .as_ref()
119                        .map(|s| format!("{:?}", s))
120                        .unwrap_or_default();
121                    let key_usage = meta
122                        .key_usage
123                        .as_ref()
124                        .map(|u| format!("{:?}", u))
125                        .unwrap_or_default();
126                    let key_spec = meta
127                        .key_spec
128                        .as_ref()
129                        .map(|s| format!("{:?}", s))
130                        .unwrap_or_default();
131                    let origin = meta
132                        .origin
133                        .as_ref()
134                        .map(|o| format!("{:?}", o))
135                        .unwrap_or_default();
136
137                    let fmt_dt = |dt: aws_smithy_types::DateTime| {
138                        dt.fmt(aws_smithy_types::date_time::Format::DateTime)
139                            .unwrap_or_default()
140                    };
141
142                    let creation_date = meta.creation_date.map(fmt_dt).unwrap_or_default();
143                    let expiration_date = meta.valid_to.map(fmt_dt).unwrap_or_default();
144                    let deletion_date = meta.deletion_date.map(fmt_dt).unwrap_or_default();
145                    let custom_key_store_id = meta.custom_key_store_id.unwrap_or_default();
146
147                    Some(KmsKey {
148                        key_id: meta.key_id,
149                        key_arn: meta.arn.unwrap_or_default(),
150                        description: meta.description.unwrap_or_default(),
151                        enabled: meta.enabled,
152                        alias,
153                        key_state,
154                        key_usage,
155                        key_spec,
156                        key_manager,
157                        creation_date,
158                        expiration_date,
159                        deletion_date,
160                        custom_key_store_id,
161                        origin,
162                        multi_region: meta.multi_region.unwrap_or(false),
163                    })
164                }
165            })
166            .buffer_unordered(10)
167            .filter_map(|k| async move { k })
168            .collect()
169            .await;
170
171        result.sort_by(|a, b| a.alias.cmp(&b.alias).then_with(|| a.key_id.cmp(&b.key_id)));
172
173        Ok(result)
174    }
175}