use std::env;
#[derive(Debug, Clone)]
pub struct Config {
pub aws_access_key_id: String,
pub aws_secret_access_key: String,
pub aws_region: String,
pub security_config: SecurityConfig,
}
#[derive(Debug, Clone)]
pub struct SecurityConfig {
pub auto_credential_disable: bool,
pub throttle_detection_window: u64,
pub throttle_disable_threshold: u32,
pub monitored_models: Vec<String>,
pub dry_run_mode: bool,
pub sns_topic_arn: Option<String>,
pub enable_detailed_logging: bool,
pub disable_record_retention_days: u32,
}
impl Default for SecurityConfig {
fn default() -> Self {
Self {
auto_credential_disable: true,
throttle_detection_window: 1,
throttle_disable_threshold: 1,
monitored_models: vec![
"anthropic.claude-3-5-sonnet-20240620-v1:0".to_string(),
"anthropic.claude-3-sonnet-20240229-v1:0".to_string(),
"anthropic.claude-3-haiku-20240307-v1:0".to_string(),
"anthropic.claude-sonnet-4-20250514-v1:0".to_string(),
],
dry_run_mode: false,
sns_topic_arn: None,
enable_detailed_logging: true,
disable_record_retention_days: 30,
}
}
}
impl Config {
pub fn from_env() -> anyhow::Result<Self> {
dotenv::dotenv().ok();
let aws_access_key_id = env::var("AWS_ACCESS_KEY_ID")?;
let aws_secret_access_key = env::var("AWS_SECRET_ACCESS_KEY")?;
let aws_region = env::var("AWS_REGION")?;
let security_config = SecurityConfig {
auto_credential_disable: env::var("BEDROCK_AUTO_CREDENTIAL_DISABLE")
.unwrap_or_else(|_| "true".to_string())
.parse()
.unwrap_or(true),
throttle_detection_window: env::var("BEDROCK_THROTTLE_DETECTION_WINDOW")
.unwrap_or_else(|_| "1".to_string())
.parse()
.unwrap_or(1),
throttle_disable_threshold: env::var("BEDROCK_THROTTLE_DISABLE_THRESHOLD")
.unwrap_or_else(|_| "1".to_string())
.parse()
.unwrap_or(1),
monitored_models: env::var("BEDROCK_MONITORED_MODELS")
.unwrap_or_else(|_| {
"anthropic.claude-3-5-sonnet-20240620-v1:0,\
anthropic.claude-3-sonnet-20240229-v1:0,\
anthropic.claude-3-haiku-20240307-v1:0,\
anthropic.claude-sonnet-4-20250514-v1:0".to_string()
})
.split(',')
.map(|s| s.trim().to_string())
.collect(),
dry_run_mode: env::var("BEDROCK_DRY_RUN_MODE")
.unwrap_or_else(|_| "false".to_string())
.parse()
.unwrap_or(false),
sns_topic_arn: env::var("BEDROCK_SNS_TOPIC_ARN").ok(),
enable_detailed_logging: env::var("BEDROCK_ENABLE_DETAILED_LOGGING")
.unwrap_or_else(|_| "true".to_string())
.parse()
.unwrap_or(true),
disable_record_retention_days: env::var("BEDROCK_DISABLE_RECORD_RETENTION_DAYS")
.unwrap_or_else(|_| "30".to_string())
.parse()
.unwrap_or(30),
};
Ok(Config {
aws_access_key_id,
aws_secret_access_key,
aws_region,
security_config,
})
}
pub fn get_threat_detection_config(&self) -> ThreatDetectionConfig {
ThreatDetectionConfig {
enabled: self.security_config.auto_credential_disable,
detection_window: self.security_config.throttle_detection_window,
threshold: self.security_config.throttle_disable_threshold,
dry_run: self.security_config.dry_run_mode,
}
}
}
#[derive(Debug, Clone)]
pub struct ThreatDetectionConfig {
pub enabled: bool,
pub detection_window: u64,
pub threshold: u32,
pub dry_run: bool,
}
impl ThreatDetectionConfig {
pub fn is_strict_mode(&self) -> bool {
self.enabled && self.detection_window <= 1 && self.threshold <= 1
}
pub fn should_trigger_immediately(&self, error_count: u32) -> bool {
self.enabled && error_count >= self.threshold
}
}