winterbaume_shield/
state.rs

1use std::collections::HashMap;
2
3use chrono::Utc;
4use thiserror::Error;
5
6use crate::types::*;
7
8/// In-memory state for the Shield service.
9#[derive(Debug, Default)]
10pub struct ShieldState {
11    /// The Shield Advanced subscription, if any.
12    pub subscription: Option<Subscription>,
13    /// Protections keyed by protection ID.
14    pub protections: HashMap<String, Protection>,
15    /// Tags keyed by resource ARN.
16    pub tags: HashMap<String, Vec<Tag>>,
17}
18
19/// Error type for Shield operations.
20#[derive(Debug, Error)]
21pub enum ShieldError {
22    #[error("Subscription already exists")]
23    SubscriptionAlreadyExists,
24
25    #[error("Subscription does not exist")]
26    SubscriptionNotFound,
27
28    #[error("A protection already exists for resource: {resource_arn}")]
29    ProtectionAlreadyExists { resource_arn: String },
30
31    #[error("Protection not found: '{id}'")]
32    ProtectionNotFound { id: String },
33
34    #[error("No protection found for resource: {arn}")]
35    ProtectionNotFoundForResource { arn: String },
36
37    #[error("Either ProtectionId or ResourceArn must be provided")]
38    MissingProtectionIdentifier,
39}
40
41impl ShieldState {
42    pub fn create_subscription(
43        &mut self,
44        account_id: &str,
45        region: &str,
46    ) -> Result<&Subscription, ShieldError> {
47        if self.subscription.is_some() {
48            return Err(ShieldError::SubscriptionAlreadyExists);
49        }
50
51        let subscription_arn = format!("arn:aws:shield:{region}:{account_id}:subscription");
52
53        let sub = Subscription {
54            start_time: Utc::now(),
55            end_time: None,
56            time_commitment_in_seconds: 31536000, // 1 year
57            auto_renew: AutoRenew::Enabled,
58            subscription_arn,
59        };
60
61        self.subscription = Some(sub);
62        Ok(self.subscription.as_ref().unwrap())
63    }
64
65    pub fn describe_subscription(&self) -> Result<&Subscription, ShieldError> {
66        self.subscription
67            .as_ref()
68            .ok_or(ShieldError::SubscriptionNotFound)
69    }
70
71    pub fn create_protection(
72        &mut self,
73        name: &str,
74        resource_arn: &str,
75        account_id: &str,
76        region: &str,
77    ) -> Result<&Protection, ShieldError> {
78        // Check if a protection already exists for this resource
79        if self
80            .protections
81            .values()
82            .any(|p| p.resource_arn == resource_arn)
83        {
84            return Err(ShieldError::ProtectionAlreadyExists {
85                resource_arn: resource_arn.to_string(),
86            });
87        }
88
89        let protection_id = uuid::Uuid::new_v4().to_string();
90        let protection_arn =
91            format!("arn:aws:shield:{region}:{account_id}:protection/{protection_id}");
92
93        let protection = Protection {
94            id: protection_id.clone(),
95            name: name.to_string(),
96            resource_arn: resource_arn.to_string(),
97            protection_arn,
98            health_check_ids: Vec::new(),
99            tags: Vec::new(),
100        };
101
102        self.protections.insert(protection_id.clone(), protection);
103        Ok(self.protections.get(&protection_id).unwrap())
104    }
105
106    pub fn describe_protection(
107        &self,
108        protection_id: Option<&str>,
109        resource_arn: Option<&str>,
110    ) -> Result<&Protection, ShieldError> {
111        if let Some(id) = protection_id {
112            return self
113                .protections
114                .get(id)
115                .ok_or_else(|| protection_not_found(id));
116        }
117
118        if let Some(arn) = resource_arn {
119            return self
120                .protections
121                .values()
122                .find(|p| p.resource_arn == arn)
123                .ok_or_else(|| ShieldError::ProtectionNotFoundForResource {
124                    arn: arn.to_string(),
125                });
126        }
127
128        Err(ShieldError::MissingProtectionIdentifier)
129    }
130
131    pub fn delete_protection(&mut self, protection_id: &str) -> Result<(), ShieldError> {
132        self.protections
133            .remove(protection_id)
134            .ok_or_else(|| protection_not_found(protection_id))?;
135        Ok(())
136    }
137
138    pub fn list_protections(&self) -> Vec<&Protection> {
139        self.protections.values().collect()
140    }
141}
142
143fn protection_not_found(id: &str) -> ShieldError {
144    ShieldError::ProtectionNotFound { id: id.to_string() }
145}
winterbaume_shield/state.rs

winterbaume_shield/
state.rs
