1use std::collections::HashMap;
4
5use chrono::{DateTime, Utc};
6use serde::{Deserialize, Serialize};
7use winterbaume_core::{StateChangeNotifier, StateViewError, StatefulService};
8
9use crate::handlers::ShieldService;
10use crate::state::ShieldState;
11use crate::types::{AutoRenew, Protection, Subscription, Tag};
12
13#[derive(Debug, Clone, Serialize, Deserialize, Default)]
14pub struct ShieldStateView {
15 pub subscription: Option<SubscriptionView>,
16 #[serde(default)]
17 pub protections: HashMap<String, ProtectionView>,
18 #[serde(default)]
19 pub tags: HashMap<String, Vec<TagView>>,
20}
21
22#[derive(Debug, Clone, Serialize, Deserialize)]
23pub struct SubscriptionView {
24 pub start_time: DateTime<Utc>,
25 pub end_time: Option<DateTime<Utc>>,
26 pub time_commitment_in_seconds: i64,
27 pub auto_renew: String,
28 pub subscription_arn: String,
29}
30
31#[derive(Debug, Clone, Serialize, Deserialize)]
32pub struct ProtectionView {
33 pub id: String,
34 pub name: String,
35 pub resource_arn: String,
36 pub protection_arn: String,
37 #[serde(default)]
38 pub health_check_ids: Vec<String>,
39 #[serde(default)]
40 pub tags: Vec<TagView>,
41}
42
43#[derive(Debug, Clone, Serialize, Deserialize)]
44pub struct TagView {
45 pub key: String,
46 pub value: String,
47}
48
49impl From<&ShieldState> for ShieldStateView {
50 fn from(state: &ShieldState) -> Self {
51 ShieldStateView {
52 subscription: state.subscription.as_ref().map(|s| SubscriptionView {
53 start_time: s.start_time,
54 end_time: s.end_time,
55 time_commitment_in_seconds: s.time_commitment_in_seconds,
56 auto_renew: s.auto_renew.as_str().to_string(),
57 subscription_arn: s.subscription_arn.clone(),
58 }),
59 protections: state
60 .protections
61 .iter()
62 .map(|(k, v)| {
63 (
64 k.clone(),
65 ProtectionView {
66 id: v.id.clone(),
67 name: v.name.clone(),
68 resource_arn: v.resource_arn.clone(),
69 protection_arn: v.protection_arn.clone(),
70 health_check_ids: v.health_check_ids.clone(),
71 tags: v
72 .tags
73 .iter()
74 .map(|t| TagView {
75 key: t.key.clone(),
76 value: t.value.clone(),
77 })
78 .collect(),
79 },
80 )
81 })
82 .collect(),
83 tags: state
84 .tags
85 .iter()
86 .map(|(k, v)| {
87 (
88 k.clone(),
89 v.iter()
90 .map(|t| TagView {
91 key: t.key.clone(),
92 value: t.value.clone(),
93 })
94 .collect(),
95 )
96 })
97 .collect(),
98 }
99 }
100}
101
102impl From<ShieldStateView> for ShieldState {
103 fn from(view: ShieldStateView) -> Self {
104 ShieldState {
105 subscription: view.subscription.map(|s| Subscription {
106 start_time: s.start_time,
107 end_time: s.end_time,
108 time_commitment_in_seconds: s.time_commitment_in_seconds,
109 auto_renew: if s.auto_renew == "ENABLED" {
110 AutoRenew::Enabled
111 } else {
112 AutoRenew::Disabled
113 },
114 subscription_arn: s.subscription_arn,
115 }),
116 protections: view
117 .protections
118 .into_iter()
119 .map(|(k, v)| {
120 (
121 k,
122 Protection {
123 id: v.id,
124 name: v.name,
125 resource_arn: v.resource_arn,
126 protection_arn: v.protection_arn,
127 health_check_ids: v.health_check_ids,
128 tags: v
129 .tags
130 .into_iter()
131 .map(|t| Tag {
132 key: t.key,
133 value: t.value,
134 })
135 .collect(),
136 },
137 )
138 })
139 .collect(),
140 tags: view
141 .tags
142 .into_iter()
143 .map(|(k, v)| {
144 (
145 k,
146 v.into_iter()
147 .map(|t| Tag {
148 key: t.key,
149 value: t.value,
150 })
151 .collect(),
152 )
153 })
154 .collect(),
155 }
156 }
157}
158
159impl StatefulService for ShieldService {
160 type StateView = ShieldStateView;
161
162 async fn snapshot(&self, account_id: &str, region: &str) -> Self::StateView {
163 let state = self.state.get(account_id, region);
164 let guard = state.read().await;
165 ShieldStateView::from(&*guard)
166 }
167
168 async fn restore(
169 &self,
170 account_id: &str,
171 region: &str,
172 view: Self::StateView,
173 ) -> Result<(), StateViewError> {
174 let state = self.state.get(account_id, region);
175 {
176 let mut guard = state.write().await;
177 *guard = ShieldState::from(view);
178 }
179 self.notify_state_changed(account_id, region).await;
180 Ok(())
181 }
182
183 async fn merge(
184 &self,
185 account_id: &str,
186 region: &str,
187 view: Self::StateView,
188 ) -> Result<(), StateViewError> {
189 let state = self.state.get(account_id, region);
190 {
191 let mut guard = state.write().await;
192 if let Some(sub) = view.subscription {
193 guard.subscription = Some(Subscription {
194 start_time: sub.start_time,
195 end_time: sub.end_time,
196 time_commitment_in_seconds: sub.time_commitment_in_seconds,
197 auto_renew: if sub.auto_renew == "ENABLED" {
198 AutoRenew::Enabled
199 } else {
200 AutoRenew::Disabled
201 },
202 subscription_arn: sub.subscription_arn,
203 });
204 }
205 for (k, v) in view.protections {
206 guard.protections.insert(
207 k,
208 Protection {
209 id: v.id,
210 name: v.name,
211 resource_arn: v.resource_arn,
212 protection_arn: v.protection_arn,
213 health_check_ids: v.health_check_ids,
214 tags: v
215 .tags
216 .into_iter()
217 .map(|t| Tag {
218 key: t.key,
219 value: t.value,
220 })
221 .collect(),
222 },
223 );
224 }
225 }
226 self.notify_state_changed(account_id, region).await;
227 Ok(())
228 }
229
230 fn notifier(&self) -> &StateChangeNotifier<Self::StateView> {
231 &self.notifier
232 }
233}