use chrono::{DateTime, Duration, Utc};
use dashmap::DashMap;
use serde::{Deserialize, Serialize};
use std::sync::atomic::{AtomicU64, Ordering};
use std::time::{SystemTime, UNIX_EPOCH};
use thiserror::Error;
use tracing::warn;
pub const DEFAULT_MAX_CAMPAIGN_CAPACITY: usize = 10_000;
pub const DEFAULT_MAX_IP_MAPPING_CAPACITY: usize = 500_000;
#[derive(Debug, Error)]
pub enum CampaignError {
#[error("Campaign not found: {0}")]
NotFound(String),
#[error("Campaign already exists: {0}")]
AlreadyExists(String),
#[error("Actor not in campaign: {0}")]
ActorNotInCampaign(String),
#[error("Invalid campaign state: {0}")]
InvalidState(String),
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Hash)]
#[serde(rename_all = "lowercase")]
#[derive(Default)]
pub enum CampaignStatus {
#[default]
Detected,
Active,
Dormant,
Resolved,
}
impl std::fmt::Display for CampaignStatus {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::Detected => write!(f, "detected"),
Self::Active => write!(f, "active"),
Self::Dormant => write!(f, "dormant"),
Self::Resolved => write!(f, "resolved"),
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Hash)]
#[serde(rename_all = "snake_case")]
pub enum CorrelationType {
AttackSequence,
AuthToken,
HttpFingerprint,
TlsFingerprint,
BehavioralSimilarity,
TimingCorrelation,
NetworkProximity,
}
impl CorrelationType {
pub fn weight(&self) -> u8 {
match self {
Self::AttackSequence => 50,
Self::AuthToken => 45,
Self::HttpFingerprint => 40,
Self::TlsFingerprint => 35,
Self::BehavioralSimilarity => 30,
Self::TimingCorrelation => 25,
Self::NetworkProximity => 15,
}
}
pub fn all_by_weight() -> Vec<Self> {
vec![
Self::AttackSequence,
Self::AuthToken,
Self::HttpFingerprint,
Self::TlsFingerprint,
Self::BehavioralSimilarity,
Self::TimingCorrelation,
Self::NetworkProximity,
]
}
pub fn is_fingerprint(&self) -> bool {
matches!(self, Self::HttpFingerprint | Self::TlsFingerprint)
}
pub fn is_behavioral(&self) -> bool {
matches!(self, Self::BehavioralSimilarity | Self::TimingCorrelation)
}
pub fn display_name(&self) -> &'static str {
match self {
Self::AttackSequence => "Attack Sequence",
Self::AuthToken => "Auth Token",
Self::HttpFingerprint => "HTTP Fingerprint",
Self::TlsFingerprint => "TLS Fingerprint",
Self::BehavioralSimilarity => "Behavioral Similarity",
Self::TimingCorrelation => "Timing Correlation",
Self::NetworkProximity => "Network Proximity",
}
}
}
impl std::fmt::Display for CorrelationType {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::AttackSequence => write!(f, "attack_sequence"),
Self::AuthToken => write!(f, "auth_token"),
Self::HttpFingerprint => write!(f, "http_fingerprint"),
Self::TlsFingerprint => write!(f, "tls_fingerprint"),
Self::BehavioralSimilarity => write!(f, "behavioral_similarity"),
Self::TimingCorrelation => write!(f, "timing_correlation"),
Self::NetworkProximity => write!(f, "network_proximity"),
}
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CorrelationReason {
pub correlation_type: CorrelationType,
pub confidence: f64,
pub description: String,
pub evidence: Vec<String>,
}
impl CorrelationReason {
pub fn new(
correlation_type: CorrelationType,
confidence: f64,
description: impl Into<String>,
evidence: Vec<String>,
) -> Self {
Self {
correlation_type,
confidence: confidence.clamp(0.0, 1.0),
description: description.into(),
evidence,
}
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Campaign {
pub id: String,
pub status: CampaignStatus,
pub actors: Vec<String>,
pub actor_count: usize,
pub confidence: f64,
pub attack_types: Vec<String>,
pub correlation_reasons: Vec<CorrelationReason>,
pub first_seen: DateTime<Utc>,
pub last_activity: DateTime<Utc>,
pub total_requests: u64,
pub blocked_requests: u64,
pub rules_triggered: u64,
pub risk_score: u32,
#[serde(skip_serializing_if = "Option::is_none")]
pub resolved_at: Option<DateTime<Utc>>,
#[serde(skip_serializing_if = "Option::is_none")]
pub resolved_reason: Option<String>,
}
impl Campaign {
pub fn generate_id() -> String {
let timestamp = SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap_or_default()
.as_millis();
format!("camp-{:x}", timestamp)
}
pub fn new(id: String, actors: Vec<String>, confidence: f64) -> Self {
let now = Utc::now();
let actor_count = actors.len();
Self {
id,
status: CampaignStatus::Detected,
actors,
actor_count,
confidence: confidence.clamp(0.0, 1.0),
attack_types: Vec::new(),
correlation_reasons: Vec::new(),
first_seen: now,
last_activity: now,
total_requests: 0,
blocked_requests: 0,
rules_triggered: 0,
risk_score: 0,
resolved_at: None,
resolved_reason: None,
}
}
#[inline]
pub fn is_active(&self) -> bool {
matches!(
self.status,
CampaignStatus::Detected | CampaignStatus::Active
)
}
#[inline]
pub fn is_resolved(&self) -> bool {
self.status == CampaignStatus::Resolved
}
pub fn time_since_activity(&self) -> Duration {
Utc::now().signed_duration_since(self.last_activity)
}
pub fn block_rate(&self) -> f64 {
if self.total_requests == 0 {
0.0
} else {
self.blocked_requests as f64 / self.total_requests as f64
}
}
}
#[derive(Debug, Clone, Default)]
pub struct CampaignUpdate {
pub campaign_id: Option<String>,
pub status: Option<CampaignStatus>,
pub confidence: Option<f64>,
pub attack_types: Option<Vec<String>>,
pub add_member_ips: Option<Vec<String>>,
pub add_correlation_reason: Option<CorrelationReason>,
pub increment_requests: Option<u64>,
pub increment_blocked: Option<u64>,
pub increment_rules: Option<u64>,
pub risk_score: Option<u32>,
}
impl CampaignUpdate {
pub fn new() -> Self {
Self::default()
}
pub fn with_status(mut self, status: CampaignStatus) -> Self {
self.status = Some(status);
self
}
pub fn with_confidence(mut self, confidence: f64) -> Self {
self.confidence = Some(confidence);
self
}
pub fn with_request_increment(mut self, total: u64, blocked: u64) -> Self {
self.increment_requests = Some(total);
self.increment_blocked = Some(blocked);
self
}
}
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct CampaignStoreStats {
pub total_campaigns: usize,
pub active_campaigns: usize,
pub detected_campaigns: usize,
pub resolved_campaigns: usize,
pub total_actors: usize,
pub evicted_campaigns: u64,
}
#[derive(Debug, Clone)]
pub struct CampaignStoreConfig {
pub max_campaign_capacity: usize,
pub max_ip_mapping_capacity: usize,
}
impl Default for CampaignStoreConfig {
fn default() -> Self {
Self {
max_campaign_capacity: DEFAULT_MAX_CAMPAIGN_CAPACITY,
max_ip_mapping_capacity: DEFAULT_MAX_IP_MAPPING_CAPACITY,
}
}
}
pub struct CampaignStore {
campaigns: DashMap<String, Campaign>,
ip_to_campaign: DashMap<String, String>,
config: CampaignStoreConfig,
evicted_count: AtomicU64,
}
impl Default for CampaignStore {
fn default() -> Self {
Self::new()
}
}
impl CampaignStore {
pub fn new() -> Self {
Self::with_config(CampaignStoreConfig::default())
}
pub fn with_config(config: CampaignStoreConfig) -> Self {
Self {
campaigns: DashMap::new(),
ip_to_campaign: DashMap::new(),
config,
evicted_count: AtomicU64::new(0),
}
}
pub fn with_capacity(campaigns: usize, actors: usize) -> Self {
Self {
campaigns: DashMap::with_capacity(campaigns),
ip_to_campaign: DashMap::with_capacity(actors),
config: CampaignStoreConfig::default(),
evicted_count: AtomicU64::new(0),
}
}
fn check_and_evict_if_needed(&self) {
while self.campaigns.len() > self.config.max_campaign_capacity {
let evicted = self
.evict_one_by_status(CampaignStatus::Resolved)
.or_else(|| self.evict_one_by_status(CampaignStatus::Dormant))
.or_else(|| self.evict_one_by_status(CampaignStatus::Detected))
.or_else(|| self.evict_one_by_status(CampaignStatus::Active));
if evicted.is_none() {
break;
}
}
}
fn evict_one_by_status(&self, status: CampaignStatus) -> Option<Campaign> {
let mut oldest_id: Option<String> = None;
let mut oldest_time: Option<DateTime<Utc>> = None;
for entry in self.campaigns.iter() {
let campaign = entry.value();
if campaign.status == status && oldest_time.is_none_or(|t| campaign.last_activity < t) {
oldest_id = Some(entry.key().clone());
oldest_time = Some(campaign.last_activity);
}
}
if let Some(id) = oldest_id {
if let Some(campaign) = self.remove_campaign(&id) {
self.evicted_count.fetch_add(1, Ordering::Relaxed);
warn!(
campaign_id = %campaign.id,
status = %campaign.status,
actor_count = campaign.actor_count,
"Evicted campaign due to capacity limit"
);
return Some(campaign);
}
}
None
}
pub fn config(&self) -> &CampaignStoreConfig {
&self.config
}
pub fn create_campaign(&self, campaign: Campaign) -> Result<(), CampaignError> {
let id = campaign.id.clone();
if self.campaigns.contains_key(&id) {
return Err(CampaignError::AlreadyExists(id));
}
for actor in &campaign.actors {
self.ip_to_campaign.insert(actor.clone(), id.clone());
}
self.campaigns.insert(id, campaign);
self.check_and_evict_if_needed();
Ok(())
}
pub fn get_campaign(&self, id: &str) -> Option<Campaign> {
self.campaigns.get(id).map(|entry| entry.value().clone())
}
pub fn get_campaign_for_ip(&self, ip: &str) -> Option<String> {
self.ip_to_campaign
.get(ip)
.map(|entry| entry.value().clone())
}
pub fn update_campaign(&self, id: &str, update: CampaignUpdate) -> Result<(), CampaignError> {
let mut entry = self
.campaigns
.get_mut(id)
.ok_or_else(|| CampaignError::NotFound(id.to_string()))?;
let campaign = entry.value_mut();
if let Some(status) = update.status {
campaign.status = status;
}
if let Some(confidence) = update.confidence {
campaign.confidence = confidence.clamp(0.0, 1.0);
}
if let Some(attack_types) = update.attack_types {
campaign.attack_types = attack_types;
}
if let Some(reason) = update.add_correlation_reason {
campaign.correlation_reasons.push(reason);
}
if let Some(increment) = update.increment_requests {
campaign.total_requests = campaign.total_requests.saturating_add(increment);
}
if let Some(increment) = update.increment_blocked {
campaign.blocked_requests = campaign.blocked_requests.saturating_add(increment);
}
if let Some(increment) = update.increment_rules {
campaign.rules_triggered = campaign.rules_triggered.saturating_add(increment);
}
if let Some(risk_score) = update.risk_score {
campaign.risk_score = risk_score;
}
campaign.last_activity = Utc::now();
Ok(())
}
pub fn add_actor_to_campaign(&self, campaign_id: &str, ip: &str) -> Result<(), CampaignError> {
let mut entry = self
.campaigns
.get_mut(campaign_id)
.ok_or_else(|| CampaignError::NotFound(campaign_id.to_string()))?;
let campaign = entry.value_mut();
if !campaign.actors.contains(&ip.to_string()) {
campaign.actors.push(ip.to_string());
campaign.actor_count = campaign.actors.len();
campaign.last_activity = Utc::now();
self.ip_to_campaign
.insert(ip.to_string(), campaign_id.to_string());
}
Ok(())
}
pub fn remove_actor_from_campaign(
&self,
campaign_id: &str,
ip: &str,
) -> Result<(), CampaignError> {
let mut entry = self
.campaigns
.get_mut(campaign_id)
.ok_or_else(|| CampaignError::NotFound(campaign_id.to_string()))?;
let campaign = entry.value_mut();
let original_len = campaign.actors.len();
campaign.actors.retain(|actor| actor != ip);
if campaign.actors.len() == original_len {
return Err(CampaignError::ActorNotInCampaign(ip.to_string()));
}
campaign.actor_count = campaign.actors.len();
campaign.last_activity = Utc::now();
self.ip_to_campaign.remove(ip);
Ok(())
}
pub fn list_campaigns(&self, status_filter: Option<CampaignStatus>) -> Vec<Campaign> {
self.campaigns
.iter()
.filter(|entry| {
status_filter
.map(|status| entry.value().status == status)
.unwrap_or(true)
})
.map(|entry| entry.value().clone())
.collect()
}
pub fn list_active_campaigns(&self) -> Vec<Campaign> {
self.campaigns
.iter()
.filter(|entry| entry.value().is_active())
.map(|entry| entry.value().clone())
.collect()
}
pub fn resolve_campaign(&self, id: &str, reason: &str) -> Result<(), CampaignError> {
let mut entry = self
.campaigns
.get_mut(id)
.ok_or_else(|| CampaignError::NotFound(id.to_string()))?;
let campaign = entry.value_mut();
if campaign.status == CampaignStatus::Resolved {
return Err(CampaignError::InvalidState(format!(
"Campaign {} is already resolved",
id
)));
}
campaign.status = CampaignStatus::Resolved;
campaign.resolved_at = Some(Utc::now());
campaign.resolved_reason = Some(reason.to_string());
campaign.last_activity = Utc::now();
Ok(())
}
pub fn expire_dormant_campaigns(&self, max_age: Duration) -> Vec<String> {
let now = Utc::now();
let mut expired = Vec::new();
for mut entry in self.campaigns.iter_mut() {
let campaign = entry.value_mut();
if matches!(
campaign.status,
CampaignStatus::Resolved | CampaignStatus::Dormant
) {
continue;
}
let age = now.signed_duration_since(campaign.last_activity);
if age > max_age {
campaign.status = CampaignStatus::Dormant;
expired.push(campaign.id.clone());
}
}
expired
}
pub fn stats(&self) -> CampaignStoreStats {
let mut active = 0;
let mut detected = 0;
let mut resolved = 0;
let mut total_actors = 0;
for entry in self.campaigns.iter() {
let campaign = entry.value();
total_actors += campaign.actor_count;
match campaign.status {
CampaignStatus::Detected => {
detected += 1;
active += 1;
}
CampaignStatus::Active => active += 1,
CampaignStatus::Resolved => resolved += 1,
CampaignStatus::Dormant => {}
}
}
CampaignStoreStats {
total_campaigns: self.campaigns.len(),
active_campaigns: active,
detected_campaigns: detected,
resolved_campaigns: resolved,
total_actors,
evicted_campaigns: self.evicted_count.load(Ordering::Relaxed),
}
}
#[inline]
pub fn len(&self) -> usize {
self.campaigns.len()
}
#[inline]
pub fn is_empty(&self) -> bool {
self.campaigns.is_empty()
}
pub fn clear(&self) {
self.campaigns.clear();
self.ip_to_campaign.clear();
}
pub fn remove_campaign(&self, id: &str) -> Option<Campaign> {
if let Some((_, campaign)) = self.campaigns.remove(id) {
for actor in &campaign.actors {
self.ip_to_campaign.remove(actor);
}
Some(campaign)
} else {
None
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_campaign_generate_id() {
let id1 = Campaign::generate_id();
let id2 = Campaign::generate_id();
assert!(id1.starts_with("camp-"));
assert!(id2.starts_with("camp-"));
assert!(id1.len() > 5); }
#[test]
fn test_campaign_new() {
let actors = vec!["192.168.1.1".to_string(), "192.168.1.2".to_string()];
let campaign = Campaign::new("test-1".to_string(), actors.clone(), 0.85);
assert_eq!(campaign.id, "test-1");
assert_eq!(campaign.status, CampaignStatus::Detected);
assert_eq!(campaign.actors, actors);
assert_eq!(campaign.actor_count, 2);
assert!((campaign.confidence - 0.85).abs() < 0.001);
assert!(campaign.is_active());
assert!(!campaign.is_resolved());
}
#[test]
fn test_campaign_confidence_clamping() {
let campaign = Campaign::new("test-1".to_string(), vec![], 1.5);
assert!((campaign.confidence - 1.0).abs() < 0.001);
let campaign = Campaign::new("test-2".to_string(), vec![], -0.5);
assert!(campaign.confidence >= 0.0);
}
#[test]
fn test_campaign_block_rate() {
let mut campaign = Campaign::new("test-1".to_string(), vec![], 0.9);
assert!((campaign.block_rate() - 0.0).abs() < 0.001);
campaign.total_requests = 100;
campaign.blocked_requests = 50;
assert!((campaign.block_rate() - 0.5).abs() < 0.001);
campaign.blocked_requests = 100;
assert!((campaign.block_rate() - 1.0).abs() < 0.001);
}
#[test]
fn test_store_create_campaign() {
let store = CampaignStore::new();
let campaign = Campaign::new("camp-1".to_string(), vec!["192.168.1.1".to_string()], 0.9);
let result = store.create_campaign(campaign);
assert!(result.is_ok());
assert_eq!(store.len(), 1);
}
#[test]
fn test_store_create_duplicate_fails() {
let store = CampaignStore::new();
let campaign1 = Campaign::new("camp-1".to_string(), vec![], 0.9);
store.create_campaign(campaign1).unwrap();
let campaign2 = Campaign::new("camp-1".to_string(), vec![], 0.8);
let result = store.create_campaign(campaign2);
assert!(matches!(result, Err(CampaignError::AlreadyExists(_))));
}
#[test]
fn test_store_get_campaign() {
let store = CampaignStore::new();
let campaign = Campaign::new("camp-1".to_string(), vec!["10.0.0.1".to_string()], 0.9);
store.create_campaign(campaign).unwrap();
let retrieved = store.get_campaign("camp-1");
assert!(retrieved.is_some());
assert_eq!(retrieved.unwrap().id, "camp-1");
let not_found = store.get_campaign("nonexistent");
assert!(not_found.is_none());
}
#[test]
fn test_store_get_campaign_for_ip() {
let store = CampaignStore::new();
let campaign = Campaign::new(
"camp-1".to_string(),
vec!["192.168.1.1".to_string(), "192.168.1.2".to_string()],
0.9,
);
store.create_campaign(campaign).unwrap();
assert_eq!(
store.get_campaign_for_ip("192.168.1.1"),
Some("camp-1".to_string())
);
assert_eq!(
store.get_campaign_for_ip("192.168.1.2"),
Some("camp-1".to_string())
);
assert!(store.get_campaign_for_ip("10.0.0.1").is_none());
}
#[test]
fn test_store_update_campaign() {
let store = CampaignStore::new();
let campaign = Campaign::new("camp-1".to_string(), vec![], 0.5);
store.create_campaign(campaign).unwrap();
let update = CampaignUpdate::new()
.with_status(CampaignStatus::Active)
.with_confidence(0.95);
let result = store.update_campaign("camp-1", update);
assert!(result.is_ok());
let updated = store.get_campaign("camp-1").unwrap();
assert_eq!(updated.status, CampaignStatus::Active);
assert!((updated.confidence - 0.95).abs() < 0.001);
}
#[test]
fn test_store_update_nonexistent_fails() {
let store = CampaignStore::new();
let update = CampaignUpdate::new().with_confidence(0.9);
let result = store.update_campaign("nonexistent", update);
assert!(matches!(result, Err(CampaignError::NotFound(_))));
}
#[test]
fn test_store_update_increments() {
let store = CampaignStore::new();
let campaign = Campaign::new("camp-1".to_string(), vec![], 0.9);
store.create_campaign(campaign).unwrap();
let update = CampaignUpdate {
increment_requests: Some(100),
increment_blocked: Some(25),
increment_rules: Some(10),
..Default::default()
};
store.update_campaign("camp-1", update).unwrap();
let updated = store.get_campaign("camp-1").unwrap();
assert_eq!(updated.total_requests, 100);
assert_eq!(updated.blocked_requests, 25);
assert_eq!(updated.rules_triggered, 10);
}
#[test]
fn test_store_add_actor() {
let store = CampaignStore::new();
let campaign = Campaign::new("camp-1".to_string(), vec!["10.0.0.1".to_string()], 0.9);
store.create_campaign(campaign).unwrap();
let result = store.add_actor_to_campaign("camp-1", "10.0.0.2");
assert!(result.is_ok());
let updated = store.get_campaign("camp-1").unwrap();
assert_eq!(updated.actor_count, 2);
assert!(updated.actors.contains(&"10.0.0.2".to_string()));
assert_eq!(
store.get_campaign_for_ip("10.0.0.2"),
Some("camp-1".to_string())
);
}
#[test]
fn test_store_add_duplicate_actor_idempotent() {
let store = CampaignStore::new();
let campaign = Campaign::new("camp-1".to_string(), vec!["10.0.0.1".to_string()], 0.9);
store.create_campaign(campaign).unwrap();
store.add_actor_to_campaign("camp-1", "10.0.0.1").unwrap();
store.add_actor_to_campaign("camp-1", "10.0.0.1").unwrap();
let updated = store.get_campaign("camp-1").unwrap();
assert_eq!(updated.actor_count, 1); }
#[test]
fn test_store_remove_actor() {
let store = CampaignStore::new();
let campaign = Campaign::new(
"camp-1".to_string(),
vec!["10.0.0.1".to_string(), "10.0.0.2".to_string()],
0.9,
);
store.create_campaign(campaign).unwrap();
let result = store.remove_actor_from_campaign("camp-1", "10.0.0.1");
assert!(result.is_ok());
let updated = store.get_campaign("camp-1").unwrap();
assert_eq!(updated.actor_count, 1);
assert!(!updated.actors.contains(&"10.0.0.1".to_string()));
assert!(store.get_campaign_for_ip("10.0.0.1").is_none());
}
#[test]
fn test_store_remove_nonexistent_actor_fails() {
let store = CampaignStore::new();
let campaign = Campaign::new("camp-1".to_string(), vec!["10.0.0.1".to_string()], 0.9);
store.create_campaign(campaign).unwrap();
let result = store.remove_actor_from_campaign("camp-1", "10.0.0.2");
assert!(matches!(result, Err(CampaignError::ActorNotInCampaign(_))));
}
#[test]
fn test_store_list_campaigns() {
let store = CampaignStore::new();
let mut c1 = Campaign::new("camp-1".to_string(), vec![], 0.9);
c1.status = CampaignStatus::Detected;
let mut c2 = Campaign::new("camp-2".to_string(), vec![], 0.8);
c2.status = CampaignStatus::Active;
let mut c3 = Campaign::new("camp-3".to_string(), vec![], 0.7);
c3.status = CampaignStatus::Resolved;
store.create_campaign(c1).unwrap();
store.create_campaign(c2).unwrap();
store.create_campaign(c3).unwrap();
let all = store.list_campaigns(None);
assert_eq!(all.len(), 3);
let detected = store.list_campaigns(Some(CampaignStatus::Detected));
assert_eq!(detected.len(), 1);
assert_eq!(detected[0].id, "camp-1");
let resolved = store.list_campaigns(Some(CampaignStatus::Resolved));
assert_eq!(resolved.len(), 1);
assert_eq!(resolved[0].id, "camp-3");
}
#[test]
fn test_store_list_active_campaigns() {
let store = CampaignStore::new();
let mut c1 = Campaign::new("camp-1".to_string(), vec![], 0.9);
c1.status = CampaignStatus::Detected;
let mut c2 = Campaign::new("camp-2".to_string(), vec![], 0.8);
c2.status = CampaignStatus::Active;
let mut c3 = Campaign::new("camp-3".to_string(), vec![], 0.7);
c3.status = CampaignStatus::Resolved;
let mut c4 = Campaign::new("camp-4".to_string(), vec![], 0.6);
c4.status = CampaignStatus::Dormant;
store.create_campaign(c1).unwrap();
store.create_campaign(c2).unwrap();
store.create_campaign(c3).unwrap();
store.create_campaign(c4).unwrap();
let active = store.list_active_campaigns();
assert_eq!(active.len(), 2);
assert!(active.iter().all(|c| c.is_active()));
}
#[test]
fn test_store_resolve_campaign() {
let store = CampaignStore::new();
let campaign = Campaign::new("camp-1".to_string(), vec![], 0.9);
store.create_campaign(campaign).unwrap();
let result = store.resolve_campaign("camp-1", "Threat mitigated");
assert!(result.is_ok());
let resolved = store.get_campaign("camp-1").unwrap();
assert_eq!(resolved.status, CampaignStatus::Resolved);
assert!(resolved.resolved_at.is_some());
assert_eq!(
resolved.resolved_reason,
Some("Threat mitigated".to_string())
);
}
#[test]
fn test_store_resolve_already_resolved_fails() {
let store = CampaignStore::new();
let campaign = Campaign::new("camp-1".to_string(), vec![], 0.9);
store.create_campaign(campaign).unwrap();
store
.resolve_campaign("camp-1", "First resolution")
.unwrap();
let result = store.resolve_campaign("camp-1", "Second resolution");
assert!(matches!(result, Err(CampaignError::InvalidState(_))));
}
#[test]
fn test_store_expire_dormant_campaigns() {
let store = CampaignStore::new();
let mut old_campaign = Campaign::new("camp-old".to_string(), vec![], 0.9);
old_campaign.last_activity = Utc::now() - Duration::hours(2);
store.create_campaign(old_campaign).unwrap();
let recent_campaign = Campaign::new("camp-recent".to_string(), vec![], 0.9);
store.create_campaign(recent_campaign).unwrap();
let mut resolved_campaign = Campaign::new("camp-resolved".to_string(), vec![], 0.9);
resolved_campaign.status = CampaignStatus::Resolved;
resolved_campaign.last_activity = Utc::now() - Duration::hours(3);
store.create_campaign(resolved_campaign).unwrap();
let expired = store.expire_dormant_campaigns(Duration::hours(1));
assert_eq!(expired.len(), 1);
assert!(expired.contains(&"camp-old".to_string()));
let old = store.get_campaign("camp-old").unwrap();
assert_eq!(old.status, CampaignStatus::Dormant);
let recent = store.get_campaign("camp-recent").unwrap();
assert_eq!(recent.status, CampaignStatus::Detected);
let resolved = store.get_campaign("camp-resolved").unwrap();
assert_eq!(resolved.status, CampaignStatus::Resolved); }
#[test]
fn test_store_stats() {
let store = CampaignStore::new();
let mut c1 = Campaign::new(
"camp-1".to_string(),
vec!["10.0.0.1".to_string(), "10.0.0.2".to_string()],
0.9,
);
c1.status = CampaignStatus::Detected;
let mut c2 = Campaign::new("camp-2".to_string(), vec!["10.0.0.3".to_string()], 0.8);
c2.status = CampaignStatus::Active;
let mut c3 = Campaign::new(
"camp-3".to_string(),
vec!["10.0.0.4".to_string(), "10.0.0.5".to_string()],
0.7,
);
c3.status = CampaignStatus::Resolved;
store.create_campaign(c1).unwrap();
store.create_campaign(c2).unwrap();
store.create_campaign(c3).unwrap();
let stats = store.stats();
assert_eq!(stats.total_campaigns, 3);
assert_eq!(stats.active_campaigns, 2); assert_eq!(stats.detected_campaigns, 1);
assert_eq!(stats.resolved_campaigns, 1);
assert_eq!(stats.total_actors, 5);
}
#[test]
fn test_campaign_status_transitions() {
let store = CampaignStore::new();
let campaign = Campaign::new("camp-1".to_string(), vec![], 0.9);
store.create_campaign(campaign).unwrap();
let update = CampaignUpdate::new().with_status(CampaignStatus::Active);
store.update_campaign("camp-1", update).unwrap();
assert_eq!(
store.get_campaign("camp-1").unwrap().status,
CampaignStatus::Active
);
let update = CampaignUpdate::new().with_status(CampaignStatus::Dormant);
store.update_campaign("camp-1", update).unwrap();
assert_eq!(
store.get_campaign("camp-1").unwrap().status,
CampaignStatus::Dormant
);
let update = CampaignUpdate::new().with_status(CampaignStatus::Active);
store.update_campaign("camp-1", update).unwrap();
assert_eq!(
store.get_campaign("camp-1").unwrap().status,
CampaignStatus::Active
);
}
#[test]
fn test_campaign_status_display() {
assert_eq!(format!("{}", CampaignStatus::Detected), "detected");
assert_eq!(format!("{}", CampaignStatus::Active), "active");
assert_eq!(format!("{}", CampaignStatus::Dormant), "dormant");
assert_eq!(format!("{}", CampaignStatus::Resolved), "resolved");
}
#[test]
fn test_correlation_type_weights() {
assert_eq!(CorrelationType::AttackSequence.weight(), 50);
assert_eq!(CorrelationType::AuthToken.weight(), 45);
assert_eq!(CorrelationType::HttpFingerprint.weight(), 40);
assert_eq!(CorrelationType::TlsFingerprint.weight(), 35);
assert_eq!(CorrelationType::BehavioralSimilarity.weight(), 30);
assert_eq!(CorrelationType::TimingCorrelation.weight(), 25);
assert_eq!(CorrelationType::NetworkProximity.weight(), 15);
}
#[test]
fn test_correlation_type_all_by_weight() {
let all = CorrelationType::all_by_weight();
assert_eq!(all.len(), 7);
assert_eq!(all[0], CorrelationType::AttackSequence);
assert_eq!(all[6], CorrelationType::NetworkProximity);
}
#[test]
fn test_correlation_type_display() {
assert_eq!(
format!("{}", CorrelationType::AttackSequence),
"attack_sequence"
);
assert_eq!(format!("{}", CorrelationType::AuthToken), "auth_token");
assert_eq!(
format!("{}", CorrelationType::HttpFingerprint),
"http_fingerprint"
);
assert_eq!(
format!("{}", CorrelationType::TlsFingerprint),
"tls_fingerprint"
);
assert_eq!(
format!("{}", CorrelationType::BehavioralSimilarity),
"behavioral_similarity"
);
assert_eq!(
format!("{}", CorrelationType::TimingCorrelation),
"timing_correlation"
);
assert_eq!(
format!("{}", CorrelationType::NetworkProximity),
"network_proximity"
);
}
#[test]
fn test_correlation_type_is_fingerprint() {
assert!(!CorrelationType::AttackSequence.is_fingerprint());
assert!(!CorrelationType::AuthToken.is_fingerprint());
assert!(CorrelationType::HttpFingerprint.is_fingerprint());
assert!(CorrelationType::TlsFingerprint.is_fingerprint());
assert!(!CorrelationType::BehavioralSimilarity.is_fingerprint());
assert!(!CorrelationType::TimingCorrelation.is_fingerprint());
assert!(!CorrelationType::NetworkProximity.is_fingerprint());
}
#[test]
fn test_correlation_type_is_behavioral() {
assert!(!CorrelationType::AttackSequence.is_behavioral());
assert!(!CorrelationType::AuthToken.is_behavioral());
assert!(!CorrelationType::HttpFingerprint.is_behavioral());
assert!(!CorrelationType::TlsFingerprint.is_behavioral());
assert!(CorrelationType::BehavioralSimilarity.is_behavioral());
assert!(CorrelationType::TimingCorrelation.is_behavioral());
assert!(!CorrelationType::NetworkProximity.is_behavioral());
}
#[test]
fn test_correlation_type_display_name() {
assert_eq!(
CorrelationType::AttackSequence.display_name(),
"Attack Sequence"
);
assert_eq!(CorrelationType::AuthToken.display_name(), "Auth Token");
assert_eq!(
CorrelationType::HttpFingerprint.display_name(),
"HTTP Fingerprint"
);
assert_eq!(
CorrelationType::TlsFingerprint.display_name(),
"TLS Fingerprint"
);
assert_eq!(
CorrelationType::BehavioralSimilarity.display_name(),
"Behavioral Similarity"
);
assert_eq!(
CorrelationType::TimingCorrelation.display_name(),
"Timing Correlation"
);
assert_eq!(
CorrelationType::NetworkProximity.display_name(),
"Network Proximity"
);
}
#[test]
fn test_correlation_reason_new() {
let reason = CorrelationReason::new(
CorrelationType::HttpFingerprint,
0.95,
"Identical JA4H fingerprint detected",
vec!["192.168.1.1".to_string(), "192.168.1.2".to_string()],
);
assert_eq!(reason.correlation_type, CorrelationType::HttpFingerprint);
assert!((reason.confidence - 0.95).abs() < 0.001);
assert_eq!(reason.description, "Identical JA4H fingerprint detected");
assert_eq!(reason.evidence.len(), 2);
}
#[test]
fn test_correlation_reason_confidence_clamping() {
let reason = CorrelationReason::new(
CorrelationType::TimingCorrelation,
1.5, "Test",
vec![],
);
assert!((reason.confidence - 1.0).abs() < 0.001);
let reason = CorrelationReason::new(
CorrelationType::TimingCorrelation,
-0.5, "Test",
vec![],
);
assert!(reason.confidence >= 0.0);
}
#[test]
fn test_store_remove_campaign() {
let store = CampaignStore::new();
let campaign = Campaign::new(
"camp-1".to_string(),
vec!["10.0.0.1".to_string(), "10.0.0.2".to_string()],
0.9,
);
store.create_campaign(campaign).unwrap();
assert!(store.get_campaign_for_ip("10.0.0.1").is_some());
assert!(store.get_campaign_for_ip("10.0.0.2").is_some());
let removed = store.remove_campaign("camp-1");
assert!(removed.is_some());
assert_eq!(removed.unwrap().id, "camp-1");
assert!(store.get_campaign("camp-1").is_none());
assert!(store.get_campaign_for_ip("10.0.0.1").is_none());
assert!(store.get_campaign_for_ip("10.0.0.2").is_none());
}
#[test]
fn test_store_clear() {
let store = CampaignStore::new();
for i in 0..5 {
let campaign = Campaign::new(format!("camp-{}", i), vec![format!("10.0.0.{}", i)], 0.9);
store.create_campaign(campaign).unwrap();
}
assert_eq!(store.len(), 5);
store.clear();
assert!(store.is_empty());
assert!(store.get_campaign_for_ip("10.0.0.1").is_none());
}
#[test]
fn test_store_concurrent_access() {
use std::sync::Arc;
use std::thread;
let store = Arc::new(CampaignStore::new());
let mut handles = vec![];
for i in 0..10 {
let store = Arc::clone(&store);
handles.push(thread::spawn(move || {
let campaign =
Campaign::new(format!("camp-{}", i), vec![format!("10.0.{}.1", i)], 0.9);
let _ = store.create_campaign(campaign);
for j in 2..5 {
let _ = store.add_actor_to_campaign(
&format!("camp-{}", i),
&format!("10.0.{}.{}", i, j),
);
}
}));
}
for _ in 0..5 {
let store = Arc::clone(&store);
handles.push(thread::spawn(move || {
for _ in 0..100 {
let _ = store.list_campaigns(None);
let _ = store.stats();
}
}));
}
for handle in handles {
handle.join().unwrap();
}
assert!(store.len() > 0);
let stats = store.stats();
assert!(stats.total_actors >= store.len()); }
#[test]
fn test_campaign_time_since_activity() {
let mut campaign = Campaign::new("test".to_string(), vec![], 0.9);
campaign.last_activity = Utc::now() - Duration::hours(1);
let elapsed = campaign.time_since_activity();
assert!(elapsed.num_minutes() >= 59);
assert!(elapsed.num_minutes() <= 61);
}
#[test]
fn test_serialization_roundtrip() {
let campaign = Campaign {
id: "camp-test".to_string(),
status: CampaignStatus::Active,
actors: vec!["10.0.0.1".to_string()],
actor_count: 1,
confidence: 0.85,
attack_types: vec!["SQLi".to_string()],
correlation_reasons: vec![CorrelationReason {
correlation_type: CorrelationType::HttpFingerprint,
confidence: 0.9,
description: "Test correlation".to_string(),
evidence: vec!["10.0.0.1".to_string()],
}],
first_seen: Utc::now(),
last_activity: Utc::now(),
total_requests: 100,
blocked_requests: 25,
rules_triggered: 10,
risk_score: 75,
resolved_at: None,
resolved_reason: None,
};
let json = serde_json::to_string(&campaign).unwrap();
let deserialized: Campaign = serde_json::from_str(&json).unwrap();
assert_eq!(deserialized.id, campaign.id);
assert_eq!(deserialized.status, campaign.status);
assert_eq!(deserialized.actor_count, campaign.actor_count);
assert!((deserialized.confidence - campaign.confidence).abs() < 0.001);
}
#[test]
fn test_serialization_skip_none_fields() {
let campaign = Campaign::new("test".to_string(), vec![], 0.9);
let json = serde_json::to_string(&campaign).unwrap();
assert!(!json.contains("resolved_at"));
assert!(!json.contains("resolved_reason"));
}
#[test]
fn test_store_capacity_limit_evicts_resolved_first() {
let config = CampaignStoreConfig {
max_campaign_capacity: 3,
max_ip_mapping_capacity: 100,
};
let store = CampaignStore::with_config(config);
let mut resolved = Campaign::new(
"camp-resolved".to_string(),
vec!["10.0.0.1".to_string()],
0.9,
);
resolved.status = CampaignStatus::Resolved;
resolved.last_activity = Utc::now() - Duration::hours(2);
let mut dormant = Campaign::new(
"camp-dormant".to_string(),
vec!["10.0.0.2".to_string()],
0.8,
);
dormant.status = CampaignStatus::Dormant;
dormant.last_activity = Utc::now() - Duration::hours(1);
let active = Campaign::new("camp-active".to_string(), vec!["10.0.0.3".to_string()], 0.7);
store.create_campaign(resolved).unwrap();
store.create_campaign(dormant).unwrap();
store.create_campaign(active).unwrap();
assert_eq!(store.len(), 3);
let new_campaign = Campaign::new("camp-new".to_string(), vec!["10.0.0.4".to_string()], 0.6);
store.create_campaign(new_campaign).unwrap();
assert!(store.len() <= 3);
assert!(store.get_campaign("camp-resolved").is_none());
assert!(store.get_campaign("camp-active").is_some());
assert!(store.get_campaign("camp-dormant").is_some());
let stats = store.stats();
assert!(stats.evicted_campaigns >= 1);
}
#[test]
fn test_store_capacity_limit_evicts_oldest() {
let config = CampaignStoreConfig {
max_campaign_capacity: 2,
max_ip_mapping_capacity: 100,
};
let store = CampaignStore::with_config(config);
let mut old_resolved =
Campaign::new("camp-old".to_string(), vec!["10.0.0.1".to_string()], 0.9);
old_resolved.status = CampaignStatus::Resolved;
old_resolved.last_activity = Utc::now() - Duration::hours(3);
let mut new_resolved =
Campaign::new("camp-newer".to_string(), vec!["10.0.0.2".to_string()], 0.8);
new_resolved.status = CampaignStatus::Resolved;
new_resolved.last_activity = Utc::now() - Duration::hours(1);
store.create_campaign(old_resolved).unwrap();
store.create_campaign(new_resolved).unwrap();
let active = Campaign::new("camp-active".to_string(), vec!["10.0.0.3".to_string()], 0.7);
store.create_campaign(active).unwrap();
assert!(store.get_campaign("camp-old").is_none());
assert!(store.get_campaign("camp-newer").is_some());
}
#[test]
fn test_store_with_config() {
let config = CampaignStoreConfig {
max_campaign_capacity: 5_000,
max_ip_mapping_capacity: 250_000,
};
let store = CampaignStore::with_config(config);
assert_eq!(store.config().max_campaign_capacity, 5_000);
assert_eq!(store.config().max_ip_mapping_capacity, 250_000);
}
#[test]
fn test_store_default_config() {
let config = CampaignStoreConfig::default();
assert_eq!(config.max_campaign_capacity, DEFAULT_MAX_CAMPAIGN_CAPACITY);
assert_eq!(
config.max_ip_mapping_capacity,
DEFAULT_MAX_IP_MAPPING_CAPACITY
);
}
#[test]
fn test_store_stats_includes_eviction_counter() {
let config = CampaignStoreConfig {
max_campaign_capacity: 2,
max_ip_mapping_capacity: 100,
};
let store = CampaignStore::with_config(config);
for i in 0..3 {
let mut campaign =
Campaign::new(format!("camp-{}", i), vec![format!("10.0.0.{}", i)], 0.9);
campaign.status = CampaignStatus::Resolved;
store.create_campaign(campaign).unwrap();
}
let stats = store.stats();
assert!(stats.evicted_campaigns > 0);
}
#[test]
fn test_store_eviction_cleans_ip_mappings() {
let config = CampaignStoreConfig {
max_campaign_capacity: 1,
max_ip_mapping_capacity: 100,
};
let store = CampaignStore::with_config(config);
let campaign1 = Campaign::new(
"camp-1".to_string(),
vec!["10.0.0.1".to_string(), "10.0.0.2".to_string()],
0.9,
);
store.create_campaign(campaign1).unwrap();
assert!(store.get_campaign_for_ip("10.0.0.1").is_some());
assert!(store.get_campaign_for_ip("10.0.0.2").is_some());
let campaign2 = Campaign::new("camp-2".to_string(), vec!["10.0.0.3".to_string()], 0.8);
store.create_campaign(campaign2).unwrap();
assert!(store.get_campaign_for_ip("10.0.0.1").is_none());
assert!(store.get_campaign_for_ip("10.0.0.2").is_none());
assert!(store.get_campaign_for_ip("10.0.0.3").is_some());
}
}