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 pub auto_credential_disable: bool,
15 pub throttle_detection_window: u64,
17 pub throttle_disable_threshold: u32,
19 pub monitored_models: Vec<String>,
21 pub dry_run_mode: bool,
23 pub sns_topic_arn: Option<String>,
25 pub enable_detailed_logging: bool,
27 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 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 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}