use anyhow::Result;
use async_trait::async_trait;
use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
use std::sync::Arc;
use std::time::Duration;
#[cfg(feature = "enhanced")]
pub mod enhanced;
pub mod memory;
#[cfg(feature = "enhanced")]
pub use enhanced::EnhancedStorage;
pub use memory::InMemoryStorage;
use crate::traits::{RateLimitKey, SecurityEvent};
#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub struct EventId(pub String);
#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub struct SnapshotId(pub String);
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct EventFilter {
pub client_id: Option<String>,
pub event_type: Option<String>,
pub from_time: Option<DateTime<Utc>>,
pub to_time: Option<DateTime<Utc>>,
pub limit: Option<usize>,
pub min_severity: Option<String>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RateLimitState {
pub tokens: f64,
pub last_refill: DateTime<Utc>,
pub request_count: u64,
pub penalty_multiplier: f64,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CorrelationState {
pub windows: Vec<CorrelationWindow>,
pub patterns: Vec<DetectedPattern>,
pub last_update: DateTime<Utc>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CorrelationWindow {
pub id: String,
pub start_time: DateTime<Utc>,
pub end_time: DateTime<Utc>,
pub events: Vec<EventId>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct DetectedPattern {
pub pattern_type: String,
pub confidence: f64,
pub events: Vec<EventId>,
pub detected_at: DateTime<Utc>,
}
#[async_trait]
pub trait StorageProvider: Send + Sync {
async fn store_event(&self, event: &SecurityEvent) -> Result<EventId>;
async fn get_event(&self, id: &EventId) -> Result<Option<SecurityEvent>>;
async fn query_events(&self, filter: EventFilter) -> Result<Vec<SecurityEvent>>;
async fn store_rate_limit_state(
&self,
key: &RateLimitKey,
state: &RateLimitState,
) -> Result<()>;
async fn get_rate_limit_state(&self, key: &RateLimitKey) -> Result<Option<RateLimitState>>;
async fn cleanup_rate_limit_states(&self, older_than: Duration) -> Result<u64>;
async fn store_correlation_state(
&self,
client_id: &str,
state: &CorrelationState,
) -> Result<()>;
async fn get_correlation_state(&self, client_id: &str) -> Result<Option<CorrelationState>>;
async fn create_snapshot(&self) -> Result<SnapshotId>;
async fn list_snapshots(&self) -> Result<Vec<(SnapshotId, DateTime<Utc>)>>;
async fn restore_snapshot(&self, id: &SnapshotId) -> Result<()>;
async fn delete_snapshot(&self, id: &SnapshotId) -> Result<()>;
async fn get_stats(&self) -> Result<StorageStats>;
async fn compact(&self) -> Result<()>;
}
#[async_trait]
pub trait ArchivalStorage: StorageProvider {
async fn archive_events(&self, older_than: Duration) -> Result<u64>;
async fn query_archived_events(&self, filter: EventFilter) -> Result<Vec<SecurityEvent>>;
async fn restore_from_archive(&self, filter: EventFilter) -> Result<u64>;
async fn get_archive_stats(&self) -> Result<ArchiveStats>;
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct StorageStats {
pub event_count: u64,
pub total_size: u64,
pub rate_limit_entries: u64,
pub correlation_states: u64,
pub storage_type: String,
pub metadata: serde_json::Value,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ArchiveStats {
pub archived_events: u64,
pub archive_size: u64,
pub oldest_event: Option<DateTime<Utc>>,
pub newest_event: Option<DateTime<Utc>>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct StorageConfig {
pub enabled: bool,
pub storage_type: StorageType,
pub data_dir: Option<String>,
pub connection_string: Option<String>,
pub retention_days: u32,
pub archive_after_days: Option<u32>,
pub max_storage_mb: Option<u64>,
pub compression: bool,
pub encryption_at_rest: bool,
}
impl Default for StorageConfig {
fn default() -> Self {
Self {
enabled: false,
storage_type: StorageType::Memory,
data_dir: None,
connection_string: None,
retention_days: 30,
archive_after_days: None,
max_storage_mb: Some(1024), compression: true,
encryption_at_rest: false,
}
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum StorageType {
Memory,
File,
RocksDb,
Redis,
Postgres,
S3,
#[cfg(feature = "enhanced")]
Enhanced,
}
pub trait StorageProviderFactory: Send + Sync {
fn create(&self, config: &StorageConfig) -> Result<Arc<dyn StorageProvider>>;
}
pub struct DefaultStorageFactory;
impl StorageProviderFactory for DefaultStorageFactory {
fn create(&self, config: &StorageConfig) -> Result<Arc<dyn StorageProvider>> {
match config.storage_type {
StorageType::Memory => Ok(Arc::new(InMemoryStorage::new())),
#[cfg(feature = "enhanced")]
StorageType::Enhanced => Ok(Arc::new(EnhancedStorage::new(config.clone())?)),
_ => Ok(Arc::new(InMemoryStorage::new())), }
}
}
pub fn create_storage_provider(_config: &crate::config::Config) -> Arc<dyn StorageProvider> {
let factory = DefaultStorageFactory;
let storage_config = StorageConfig::default();
factory
.create(&storage_config)
.unwrap_or_else(|_| Arc::new(InMemoryStorage::new()))
}