Skip to main content

burncloud_aws/
config.rs

1use std::env;
2
3#[derive(Debug, Clone)]
4pub struct Config {
5    pub aws_access_key_id: String,
6    pub aws_secret_access_key: String,
7    pub aws_region: String,
8    pub security_config: SecurityConfig,
9}
10
11#[derive(Debug, Clone)]
12pub struct SecurityConfig {
13    /// 是否启用自动AK/SK禁用
14    pub auto_credential_disable: bool,
15    /// 429错误检测窗口(秒)
16    pub throttle_detection_window: u64,
17    /// 触发AK/SK禁用的429错误次数
18    pub throttle_disable_threshold: u32,
19    /// 监控的模型列表
20    pub monitored_models: Vec<String>,
21    /// 是否启用试运行模式(不实际禁用)
22    pub dry_run_mode: bool,
23    /// SNS通知主题ARN
24    pub sns_topic_arn: Option<String>,
25    /// 是否启用详细日志记录
26    pub enable_detailed_logging: bool,
27    /// 凭据禁用记录保留天数
28    pub disable_record_retention_days: u32,
29}
30
31impl Default for SecurityConfig {
32    fn default() -> Self {
33        Self {
34            auto_credential_disable: true,
35            throttle_detection_window: 1,
36            throttle_disable_threshold: 1,
37            monitored_models: vec![
38                "anthropic.claude-3-5-sonnet-20240620-v1:0".to_string(),
39                "anthropic.claude-3-sonnet-20240229-v1:0".to_string(),
40                "anthropic.claude-3-haiku-20240307-v1:0".to_string(),
41                "anthropic.claude-sonnet-4-20250514-v1:0".to_string(),
42            ],
43            dry_run_mode: false,
44            sns_topic_arn: None,
45            enable_detailed_logging: true,
46            disable_record_retention_days: 30,
47        }
48    }
49}
50
51impl Config {
52    pub fn from_env() -> anyhow::Result<Self> {
53        dotenv::dotenv().ok();
54
55        let aws_access_key_id = env::var("AWS_ACCESS_KEY_ID")?;
56        let aws_secret_access_key = env::var("AWS_SECRET_ACCESS_KEY")?;
57        let aws_region = env::var("AWS_REGION")?;
58
59        // 安全配置
60        let security_config = SecurityConfig {
61            auto_credential_disable: env::var("BEDROCK_AUTO_CREDENTIAL_DISABLE")
62                .unwrap_or_else(|_| "true".to_string())
63                .parse()
64                .unwrap_or(true),
65
66            throttle_detection_window: env::var("BEDROCK_THROTTLE_DETECTION_WINDOW")
67                .unwrap_or_else(|_| "1".to_string())
68                .parse()
69                .unwrap_or(1),
70
71            throttle_disable_threshold: env::var("BEDROCK_THROTTLE_DISABLE_THRESHOLD")
72                .unwrap_or_else(|_| "1".to_string())
73                .parse()
74                .unwrap_or(1),
75
76            monitored_models: env::var("BEDROCK_MONITORED_MODELS")
77                .unwrap_or_else(|_| {
78                    "anthropic.claude-3-5-sonnet-20240620-v1:0,\
79                    anthropic.claude-3-sonnet-20240229-v1:0,\
80                    anthropic.claude-3-haiku-20240307-v1:0,\
81                    anthropic.claude-sonnet-4-20250514-v1:0".to_string()
82                })
83                .split(',')
84                .map(|s| s.trim().to_string())
85                .collect(),
86
87            dry_run_mode: env::var("BEDROCK_DRY_RUN_MODE")
88                .unwrap_or_else(|_| "false".to_string())
89                .parse()
90                .unwrap_or(false),
91
92            sns_topic_arn: env::var("BEDROCK_SNS_TOPIC_ARN").ok(),
93
94            enable_detailed_logging: env::var("BEDROCK_ENABLE_DETAILED_LOGGING")
95                .unwrap_or_else(|_| "true".to_string())
96                .parse()
97                .unwrap_or(true),
98
99            disable_record_retention_days: env::var("BEDROCK_DISABLE_RECORD_RETENTION_DAYS")
100                .unwrap_or_else(|_| "30".to_string())
101                .parse()
102                .unwrap_or(30),
103        };
104
105        Ok(Config {
106            aws_access_key_id,
107            aws_secret_access_key,
108            aws_region,
109            security_config,
110        })
111    }
112
113    /// 获取威胁检测配置
114    pub fn get_threat_detection_config(&self) -> ThreatDetectionConfig {
115        ThreatDetectionConfig {
116            enabled: self.security_config.auto_credential_disable,
117            detection_window: self.security_config.throttle_detection_window,
118            threshold: self.security_config.throttle_disable_threshold,
119            dry_run: self.security_config.dry_run_mode,
120        }
121    }
122}
123
124#[derive(Debug, Clone)]
125pub struct ThreatDetectionConfig {
126    pub enabled: bool,
127    pub detection_window: u64,
128    pub threshold: u32,
129    pub dry_run: bool,
130}
131
132impl ThreatDetectionConfig {
133    pub fn is_strict_mode(&self) -> bool {
134        self.enabled && self.detection_window <= 1 && self.threshold <= 1
135    }
136
137    pub fn should_trigger_immediately(&self, error_count: u32) -> bool {
138        self.enabled && error_count >= self.threshold
139    }
140}