1use std::sync::Arc;
14
15use async_trait::async_trait;
16use http::StatusCode;
17use serde_json::{json, Value};
18use tokio::sync::Mutex as AsyncMutex;
19
20use fakecloud_aws::arn::Arn;
21use fakecloud_core::service::{AwsRequest, AwsResponse, AwsService, AwsServiceError};
22use fakecloud_persistence::SnapshotStore;
23
24use crate::state::{SharedShieldState, ShieldState};
25
26pub const SHIELD_ACTIONS: &[&str] = &[
27 "AssociateDRTLogBucket",
28 "AssociateDRTRole",
29 "AssociateHealthCheck",
30 "AssociateProactiveEngagementDetails",
31 "CreateProtection",
32 "CreateProtectionGroup",
33 "CreateSubscription",
34 "DeleteProtection",
35 "DeleteProtectionGroup",
36 "DeleteSubscription",
37 "DescribeAttack",
38 "DescribeAttackStatistics",
39 "DescribeDRTAccess",
40 "DescribeEmergencyContactSettings",
41 "DescribeProtection",
42 "DescribeProtectionGroup",
43 "DescribeSubscription",
44 "DisableApplicationLayerAutomaticResponse",
45 "DisableProactiveEngagement",
46 "DisassociateDRTLogBucket",
47 "DisassociateDRTRole",
48 "DisassociateHealthCheck",
49 "EnableApplicationLayerAutomaticResponse",
50 "EnableProactiveEngagement",
51 "GetSubscriptionState",
52 "ListAttacks",
53 "ListProtectionGroups",
54 "ListProtections",
55 "ListResourcesInProtectionGroup",
56 "ListTagsForResource",
57 "TagResource",
58 "UntagResource",
59 "UpdateApplicationLayerAutomaticResponse",
60 "UpdateEmergencyContactSettings",
61 "UpdateProtectionGroup",
62 "UpdateSubscription",
63];
64
65fn is_mutating_action(action: &str) -> bool {
69 const READ_PREFIXES: &[&str] = &["Get", "List", "Describe"];
70 !READ_PREFIXES.iter().any(|p| action.starts_with(p))
71}
72
73pub struct ShieldService {
74 state: SharedShieldState,
75 snapshot_store: Option<Arc<dyn SnapshotStore>>,
76 snapshot_lock: Arc<AsyncMutex<()>>,
77}
78
79impl ShieldService {
80 pub fn new(state: SharedShieldState) -> Self {
81 Self {
82 state,
83 snapshot_store: None,
84 snapshot_lock: Arc::new(AsyncMutex::new(())),
85 }
86 }
87
88 pub fn with_snapshot_store(mut self, store: Arc<dyn SnapshotStore>) -> Self {
89 self.snapshot_store = Some(store);
90 self
91 }
92
93 async fn save(&self) {
94 crate::persistence::save_snapshot(
95 &self.state,
96 self.snapshot_store.clone(),
97 &self.snapshot_lock,
98 )
99 .await;
100 }
101
102 pub fn snapshot_hook(&self) -> Option<fakecloud_persistence::SnapshotHook> {
104 let store = self.snapshot_store.clone()?;
105 let state = self.state.clone();
106 let lock = self.snapshot_lock.clone();
107 Some(Arc::new(move || {
108 let state = state.clone();
109 let store = store.clone();
110 let lock = lock.clone();
111 Box::pin(async move {
112 crate::persistence::save_snapshot(&state, Some(store), &lock).await;
113 })
114 }))
115 }
116}
117
118#[async_trait]
119impl AwsService for ShieldService {
120 fn service_name(&self) -> &str {
121 "shield"
122 }
123
124 async fn handle(&self, req: AwsRequest) -> Result<AwsResponse, AwsServiceError> {
125 let action = req.action.clone();
126 if let Err(msg) = crate::validate::validate_input(&action, &req.json_body()) {
129 return Err(invalid_parameter(msg));
130 }
131 let result = self.dispatch(&action, &req);
132 if is_mutating_action(&action)
133 && matches!(result.as_ref(), Ok(resp) if resp.status.is_success())
134 {
135 self.save().await;
136 }
137 result
138 }
139
140 fn supported_actions(&self) -> &[&str] {
141 SHIELD_ACTIONS
142 }
143}
144
145impl ShieldService {
146 #[allow(clippy::too_many_lines)]
147 fn dispatch(&self, action: &str, req: &AwsRequest) -> Result<AwsResponse, AwsServiceError> {
148 match action {
149 "CreateProtection" => self.create_protection(req),
151 "DescribeProtection" => self.describe_protection(req),
152 "DeleteProtection" => self.delete_protection(req),
153 "ListProtections" => self.list_protections(req),
154 "CreateProtectionGroup" => self.create_protection_group(req),
156 "UpdateProtectionGroup" => self.update_protection_group(req),
157 "DeleteProtectionGroup" => self.delete_protection_group(req),
158 "DescribeProtectionGroup" => self.describe_protection_group(req),
159 "ListProtectionGroups" => self.list_protection_groups(req),
160 "ListResourcesInProtectionGroup" => self.list_resources_in_protection_group(req),
161 "CreateSubscription" => self.create_subscription(req),
163 "DescribeSubscription" => self.describe_subscription(req),
164 "UpdateSubscription" => self.update_subscription(req),
165 "DeleteSubscription" => self.delete_subscription(req),
166 "GetSubscriptionState" => self.get_subscription_state(req),
167 "ListAttacks" => self.list_attacks(req),
169 "DescribeAttack" => self.describe_attack(req),
170 "DescribeAttackStatistics" => self.describe_attack_statistics(req),
171 "AssociateDRTRole" => self.associate_drt_role(req),
173 "DisassociateDRTRole" => self.disassociate_drt_role(req),
174 "AssociateDRTLogBucket" => self.associate_drt_log_bucket(req),
175 "DisassociateDRTLogBucket" => self.disassociate_drt_log_bucket(req),
176 "DescribeDRTAccess" => self.describe_drt_access(req),
177 "AssociateHealthCheck" => self.associate_health_check(req),
179 "DisassociateHealthCheck" => self.disassociate_health_check(req),
180 "AssociateProactiveEngagementDetails" => {
182 self.associate_proactive_engagement_details(req)
183 }
184 "EnableProactiveEngagement" => self.enable_proactive_engagement(req),
185 "DisableProactiveEngagement" => self.disable_proactive_engagement(req),
186 "UpdateEmergencyContactSettings" => self.update_emergency_contact_settings(req),
187 "DescribeEmergencyContactSettings" => self.describe_emergency_contact_settings(req),
188 "EnableApplicationLayerAutomaticResponse" => {
190 self.enable_application_layer_automatic_response(req)
191 }
192 "UpdateApplicationLayerAutomaticResponse" => {
193 self.update_application_layer_automatic_response(req)
194 }
195 "DisableApplicationLayerAutomaticResponse" => {
196 self.disable_application_layer_automatic_response(req)
197 }
198 "TagResource" => self.tag_resource(req),
200 "UntagResource" => self.untag_resource(req),
201 "ListTagsForResource" => self.list_tags_for_resource(req),
202 _ => Err(AwsServiceError::action_not_implemented(
203 self.service_name(),
204 action,
205 )),
206 }
207 }
208}
209
210fn err(code: &str, msg: impl Into<String>) -> AwsServiceError {
213 AwsServiceError::aws_error(StatusCode::BAD_REQUEST, code, msg.into())
214}
215
216pub(crate) fn invalid_parameter(msg: impl Into<String>) -> AwsServiceError {
217 err("InvalidParameterException", msg)
218}
219
220pub(crate) fn resource_not_found(msg: impl Into<String>) -> AwsServiceError {
221 err("ResourceNotFoundException", msg)
222}
223
224pub(crate) fn resource_already_exists(msg: impl Into<String>) -> AwsServiceError {
225 err("ResourceAlreadyExistsException", msg)
226}
227
228pub(crate) fn invalid_resource(msg: impl Into<String>) -> AwsServiceError {
229 err("InvalidResourceException", msg)
230}
231
232pub(crate) fn limits_exceeded(msg: impl Into<String>) -> AwsServiceError {
233 err("LimitsExceededException", msg)
234}
235
236pub(crate) fn no_associated_role(msg: impl Into<String>) -> AwsServiceError {
237 err("NoAssociatedRoleException", msg)
238}
239
240pub(crate) fn locked_subscription(msg: impl Into<String>) -> AwsServiceError {
241 err("LockedSubscriptionException", msg)
242}
243
244pub(crate) fn ok(value: Value) -> Result<AwsResponse, AwsServiceError> {
245 Ok(AwsResponse::ok_json(value))
246}
247
248pub(crate) fn now_epoch() -> f64 {
250 chrono::Utc::now().timestamp_millis() as f64 / 1000.0
251}
252
253pub(crate) const ONE_YEAR_SECS: f64 = 365.0 * 24.0 * 60.0 * 60.0;
255
256pub fn protection_arn(region: &str, account: &str, id: &str) -> String {
258 Arn::global_in(region, "shield", account, &format!("protection/{id}")).to_string()
259}
260
261pub fn protection_group_arn(region: &str, account: &str, id: &str) -> String {
262 Arn::global_in(region, "shield", account, &format!("protection-group/{id}")).to_string()
263}
264
265pub fn subscription_arn(region: &str, account: &str) -> String {
266 Arn::global_in(region, "shield", account, "subscription").to_string()
267}
268
269impl ShieldService {
270 pub(crate) fn with_account_mut<R>(
272 &self,
273 req: &AwsRequest,
274 f: impl FnOnce(&mut ShieldState) -> R,
275 ) -> R {
276 let mut guard = self.state.write();
277 let acct = guard.get_or_create(&req.account_id);
278 f(acct)
279 }
280
281 pub(crate) fn with_account<R>(&self, req: &AwsRequest, f: impl FnOnce(&ShieldState) -> R) -> R {
284 let guard = self.state.read();
285 match guard.get(&req.account_id) {
286 Some(acct) => f(acct),
287 None => f(&ShieldState::default()),
288 }
289 }
290}
291
292include!("handlers.rs");